#!/usr/bin/env python3
# vim:tabstop=4:softtabstop=4:shiftwidth=4:textwidth=160:smarttab:expandtab

"""msp430-etv - MSP430 EnergyTrace Visualizer

DESCRIPTION

msp430-etv takes energy measurements from an MSP430 Launchpad or similar device
using MSP430 EnergyTrace technology. Measurements can be taken directly (by
specifying <measurement duration> in seconds) or loaded from a logfile using
--load <file>. Data can be plotted or aggregated on stdout.

This program is not affiliated with Texas Instruments. Use at your own risk.

DEPENDENCIES

For data measurements (i.e., any invocation not using --load),
energytrace-util <https://github.com/carrotIndustries/energytrace-util>
must be available in $PATH and libmsp430.so must be located in the
LD library search path (e.g. LD_LIBRARY_PATH=../MSP430Flasher).

OPTIONS
"""

import argparse
import itertools
import matplotlib.pyplot as plt
import numpy as np
import os
from shutil import which
import subprocess
import sys
import tempfile
import time

opt = dict()


def running_mean(x: np.ndarray, N: int) -> np.ndarray:
    """
    Compute `N` elements wide running average over `x`.

    :param x: 1-Dimensional NumPy array
    :param N: how many items to average. Should be even for optimal results.
    """

    # to ensure that output.shape == input.shape, we need to insert data
    # at the boundaries
    boundary_array = np.insert(x, 0, np.full((N // 2), x[0]))
    boundary_array = np.append(boundary_array, np.full((N // 2 + N % 2), x[-1]))
    cumsum = np.cumsum(boundary_array)
    return (cumsum[N:] - cumsum[:-N]) / N


def measure_data(filename, duration):
    # libmsp430.so must be available
    if not "LD_LIBRARY_PATH" in os.environ:
        os.environ[
            "LD_LIBRARY_PATH"
        ] = "{}/var/projects/msp430/MSP430Flasher_1.3.15".format(os.environ["HOME"])

    # https://github.com/carrotIndustries/energytrace-util must be available
    energytrace_cmd = "energytrace"
    if which(energytrace_cmd) is None:
        energytrace_cmd = "{}/var/source/energytrace-util/energytrace64".format(
            os.environ["HOME"]
        )

    if filename is not None:
        output_handle = open(filename, "w+")
    else:
        output_handle = tempfile.TemporaryFile("w+")

    energytrace = subprocess.Popen(
        [energytrace_cmd, str(duration)], stdout=output_handle, universal_newlines=True
    )

    try:
        if duration:
            time.sleep(duration)
        else:
            print("Press Ctrl+C to stop measurement")
            while True:
                time.sleep(3600)
    except KeyboardInterrupt:
        energytrace.send_signal(subprocess.signal.SIGTERM)

    energytrace.communicate(timeout=5)

    output_handle.seek(0)
    output = output_handle.read()
    output_handle.close()

    return output


def peak_search(data, lower, upper, direction_function):
    while upper - lower > 1e-6:
        bs_test = np.mean([lower, upper])
        peakcount = itertools.groupby(data, lambda x: x >= bs_test)
        peakcount = filter(lambda x: x[0] == True, peakcount)
        peakcount = sum(1 for i in peakcount)
        direction = direction_function(peakcount, bs_test)
        if direction == 0:
            return bs_test
        elif direction == 1:
            lower = bs_test
        else:
            upper = bs_test
    return None


def peak_search2(data, lower, upper, check_function):
    for power in np.arange(lower, upper, 1e-6):
        peakcount = itertools.groupby(data, lambda x: x >= power)
        peakcount = filter(lambda x: x[0] == True, peakcount)
        peakcount = sum(1 for i in peakcount)
        if check_function(peakcount, power) == 0:
            return power
    return None


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__
    )
    parser.add_argument("--load", metavar="FILE", type=str, help="Load data from FILE")
    parser.add_argument(
        "--save", metavar="FILE", type=str, help="Save measurement data in FILE"
    )
    parser.add_argument(
        "--skip",
        metavar="COUNT",
        type=int,
        default=0,
        help="Skip COUNT data samples. This is useful to avoid startup code influencing the results of a long-running measurement",
    )
    parser.add_argument(
        "--threshold",
        metavar="WATTS",
        type=str,
        help="Partition data into points with mean power >= WATTS and points with mean power < WATTS, and print some statistics. higher power is handled as peaks, whereas low-power measurements constitute the baseline. If WATTS is 'mean', the mean power of all measurements will be used",
    )
    parser.add_argument(
        "--threshold-peakcount",
        metavar="NUM",
        type=int,
        help="Automatically determine a threshold so that there are exactly NUM peaks. A peaks is a group of consecutive measurements with mean power >= threshold. WARNING: In general, there is more than one threshold value leading to exactly NUM peaks. If the difference between baseline and peak power is sufficiently high, this option should do what you mean[tm]",
    )
    parser.add_argument(
        "--plot",
        metavar="UNIT",
        choices=["U", "I", "P"],
        help="Plot voltage / current / power over time",
    )
    parser.add_argument(
        "--stat",
        action="store_true",
        help="Print mean voltage, current, and power as well as total energy consumption",
    )
    parser.add_argument(
        "--histogram",
        metavar="N",
        type=int,
        help="Draw histograms of reported energy values per measurement interval (i.e., the differences between each pair of consecutive total energy readings), measurement interval duration, and mean power values per measurement interval (calculated from energy difference and duration). Each histogram uses N buckets",
    )
    parser.add_argument(
        "duration", type=int, nargs="?", help="Measurement duration in seconds"
    )

    args = parser.parse_args()

    if args.load is None and args.duration is None:
        print("Either --load or duration must be specified", file=sys.stderr)
        sys.exit(1)

    if args.threshold is not None and args.threshold != "mean":
        args.threshold = float(args.threshold)

    if args.load:
        if args.load.endswith(".xz"):
            import lzma

            with lzma.open(args.load, "rt") as f:
                log_data = f.read()
        else:
            with open(args.load, "r") as f:
                log_data = f.read()
    else:
        log_data = measure_data(args.save, args.duration)

    lines = log_data.split("\n")
    data_count = sum(map(lambda x: len(x) > 0 and x[0] != "#", lines))
    data_lines = filter(lambda x: len(x) > 0 and x[0] != "#", lines)

    data = np.empty((data_count - args.skip, 4))

    energy_overflow_count = 0
    prev_total_energy = 0
    for i, line in enumerate(data_lines):
        if i >= args.skip:
            fields = line.split(" ")
            if len(fields) == 4:
                timestamp, current, voltage, total_energy = map(int, fields)
            elif len(fields) == 5:
                cpustate = fields[0]
                timestamp, current, voltage, total_energy = map(int, fields[1:])
            else:
                raise RuntimeError('cannot parse line "{}"'.format(line))
            if total_energy < 0 and prev_total_energy > 0:
                energy_overflow_count += 1
            prev_total_energy = total_energy
            total_energy += energy_overflow_count * (2 ** 32)
            data[i - args.skip] = [timestamp, current, voltage, total_energy]

    m_duration_us = data[-1, 0] - data[0, 0]
    m_energy_nj = data[-1, 3] - data[0, 3]

    print(
        "{:d} measurements in {:.2f} s = {:.0f} Hz sample rate".format(
            data_count, m_duration_us * 1e-6, data_count / (m_duration_us * 1e-6)
        )
    )

    print("Reported energy: E = {:f} J".format(m_energy_nj * 1e-9))

    # nJ / us = mW -> (nJ * 1e-9) / (us * 1e-6) = W
    # Do not use power = data[:, 1] * data[:, 2] * 1e-12 here: nA values provided by the EnergyTrace library in data[:, 1] are heavily filtered and mostly
    # useless for visualization and calculation. They often do not agree with the nJ values in data[:, 3].
    power = ((data[1:, 3] - data[:-1, 3]) * 1e-9) / (
        (data[1:, 0] - data[:-1, 0]) * 1e-6
    )

    if args.threshold_peakcount:
        bs_mean = np.mean(power)

        # Finding the correct threshold is tricky. If #peaks < peakcont, our
        # current threshold may be too low (extreme case: a single peaks
        # containing all measurements), but it may also be too high (extreme
        # case: a single peak containing just one data point). Similarly,
        # #peaks > peakcount may be due to baseline noise causing lots of
        # small peaks, or due to peak noise (if the threshold is already rather
        # high).
        # For now, we first try a simple binary search:
        # The threshold is probably somewhere around the mean, so if
        # #peaks != peakcount and threshold < mean, we go up, and if
        # #peaks != peakcount and threshold >= mean, we go down.
        # If that doesn't work, we fall back to a linear search in 1 µW steps
        def direction_function(peakcount, power):
            if peakcount == args.threshold - peakcount:
                return 0
            if power < bs_mean:
                return 1
            return -1

        threshold = peak_search(power, np.min(power), np.max(power), direction_function)
        if threshold == None:
            threshold = peak_search2(
                power, np.min(power), np.max(power), direction_function
            )

        if threshold != None:
            print(
                "Threshold set to {:.0f} µW         : {:.9f}".format(
                    threshold * 1e6, threshold
                )
            )
            args.threshold = threshold
        else:
            print("Found no working threshold")

    if args.threshold:
        if args.threshold == "mean":
            args.threshold = np.mean(power)
            print(
                "Threshold set to {:.0f} µW         : {:.9f}".format(
                    args.threshold * 1e6, args.threshold
                )
            )

        baseline_mean = 0
        if np.any(power < args.threshold):
            baseline_mean = np.mean(power[power < args.threshold])
            print(
                "Baseline mean: {:.0f} µW           : {:.9f}".format(
                    baseline_mean * 1e6, baseline_mean
                )
            )
        if np.any(power >= args.threshold):
            print(
                "Peak mean: {:.0f} µW               : {:.9f}".format(
                    np.mean(power[power >= args.threshold]) * 1e6,
                    np.mean(power[power >= args.threshold]),
                )
            )

        peaks = []
        peak_start = -1
        for i, dp in enumerate(power):
            if dp >= args.threshold and peak_start == -1:
                peak_start = i
            elif dp < args.threshold and peak_start != -1:
                peaks.append((peak_start, i))
                peak_start = -1

        total_energy = 0
        delta_energy = 0
        for peak in peaks:
            duration = data[peak[1] - 1, 0] - data[peak[0], 0]
            total_energy += np.mean(power[peak[0] : peak[1]]) * duration
            delta_energy += (
                np.mean(power[peak[0] : peak[1]]) - baseline_mean
            ) * duration
            print(
                "{:.2f}ms peak ({:f} -> {:f})".format(
                    duration * 1000, data[peak[0], 0], data[peak[1] - 1, 0]
                )
            )
            print(
                "    {:f} µJ / mean {:f} µW".format(
                    np.mean(power[peak[0] : peak[1]]) * duration * 1e6,
                    np.mean(power[peak[0] : peak[1]]) * 1e6,
                )
            )
        print(
            "Peak energy mean: {:.0f} µJ         : {:.9f}".format(
                total_energy * 1e6 / len(peaks), total_energy / len(peaks)
            )
        )
        print(
            "Average per-peak energy (delta over baseline): {:.0f} µJ         : {:.9f}".format(
                delta_energy * 1e6 / len(peaks), delta_energy / len(peaks)
            )
        )

    power_from_energy = ((data[1:, 3] - data[:-1, 3]) * 1e-9) / (
        (data[1:, 0] - data[:-1, 0]) * 1e-6
    )
    smooth_power = running_mean(power_from_energy, 10)

    if args.stat:
        mean_voltage = np.mean(data[:, 2] * 1e-3)
        mean_current = np.mean(data[:, 1] * 1e-9)
        mean_power = np.mean(data[:, 1] * data[:, 2] * 1e-12)
        print(
            "Mean voltage: {:.2f} V       : {:.9f}".format(mean_voltage, mean_voltage)
        )
        print(
            "Mean current: {:.0f} µA       : {:.9f}".format(
                mean_current * 1e6, mean_current
            )
        )
        print(
            "Mean power: {:.0f} µW       : {:.9f}".format(mean_power * 1e6, mean_power)
        )
        print(
            "Total energy: {:f} J       : {:.9f}".format(
                m_energy_nj * 1e-9, m_energy_nj * 1e-9
            )
        )

    if args.plot:
        if args.plot == "U":
            # mV
            (energyhandle,) = plt.plot(
                data[1:, 0] * 1e-6, data[1:, 2] * 1e-3, "b-", label="U", markersize=1
            )
            (meanhandle,) = plt.plot(
                data[1:, 0] * 1e-6,
                running_mean(data[1:, 2], 10) * 1e-3,
                "r-",
                label="mean(U, 10)",
                markersize=1,
            )
            plt.legend(handles=[energyhandle, meanhandle])
            plt.ylabel("Voltage [V]")
        elif args.plot == "I":
            print(
                "Warning: The current reported by energytrace is aggressively smoothed and often inaccurate."
            )
            # nA
            (energyhandle,) = plt.plot(
                data[1:, 0] * 1e-6, data[1:, 1] * 1e-9, "b-", label="I", markersize=1
            )
            (meanhandle,) = plt.plot(
                data[1:, 0] * 1e-6,
                running_mean(data[1:, 1], 10) * 1e-9,
                "r-",
                label="mean(I, 10)",
                markersize=1,
            )
            plt.legend(handles=[energyhandle, meanhandle])
            plt.ylabel("Current [A]")
        else:
            (energyhandle,) = plt.plot(
                data[1:, 0] * 1e-6,
                power_from_energy,
                "b-",
                label="P=ΔE/Δt",
                markersize=1,
            )
            (meanhandle,) = plt.plot(
                data[1:, 0] * 1e-6,
                smooth_power,
                "r-",
                label="mean(P, 10)",
                markersize=1,
            )
            plt.legend(handles=[energyhandle, meanhandle])
            plt.ylabel("Power [W]")
        plt.xlabel("Time [s]")
        plt.grid(True)
        if args.load:
            plt.title(args.load)
        plt.show()

    if args.histogram:
        bin_count = args.histogram

        plt.title("EnergyTrace Data Analysis")
        plt.xlabel("Reported Energy per Measurement Interval [J]")
        plt.ylabel("Count")
        plt.hist((data[1:, 3] - data[:-1, 3]) * 1e-9, bins=bin_count)
        plt.show()

        plt.title("EnergyTrace Data Analysis")
        plt.xlabel("Measurement Interval Duration [s]")
        plt.ylabel("Count")
        plt.hist((data[1:, 0] - data[:-1, 0]) * 1e-6, bins=bin_count)
        plt.show()

        plt.title("EnergyTrace Data Analysis")
        plt.xlabel("Mean Power per Measurement Interval [W]")
        plt.ylabel("Count")
        plt.hist(power_from_energy, bins=bin_count)
        plt.show()

        plt.title("Postprocessing via Running average (window size=10)")
        plt.xlabel("Mean Power per Measurement Interval [W]")
        plt.ylabel("Count")
        plt.hist(smooth_power, bins=bin_count)
        plt.show()


if __name__ == "__main__":
    main()