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

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

opt = dict()


def show_help():
    print(
        """msp430-etv - MSP430 EnergyTrace Visualizer

USAGE

msp430-etv [--load <file> | <measurement duration>] [--save <file>]
    [--skip <count>] [--threshold <power>] [--plot=U|I|P] [--stat]

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.

OPTIONS

  --load <file>
    Load data from <file>
  --save <file>
    Save measurement data in <file>
  --skip <count>
    Skip <count> data samples. This is useful to avoid startup code
    influencing the results of a long-running measurement
  --threshold <watts>|mean
    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 the threshold is set to "mean", the mean power of all measurements
    will be used
  --threshold-peakcount <num>
    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]
  --plot=U|I|P
    Plot voltage / current / power over time
  --stat
    Print mean voltage, current, and power as well as total energy consumption.
  --histogram=<n>
    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.

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).
    """
    )


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]))
    print(boundary_array)
    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


if __name__ == "__main__":
    try:
        optspec = "help load= save= skip= threshold= threshold-peakcount= plot= stat histogram="
        raw_opts, args = getopt.getopt(sys.argv[1:], "", optspec.split(" "))

        for option, parameter in raw_opts:
            optname = re.sub(r"^--", "", option)
            opt[optname] = parameter

        if "help" in opt:
            show_help()
            sys.exit(0)

        if not "load" in opt:
            duration = int(args[0])

        if not "save" in opt:
            opt["save"] = None

        if "skip" in opt:
            opt["skip"] = int(opt["skip"])
        else:
            opt["skip"] = 0

        if "threshold" in opt and opt["threshold"] != "mean":
            opt["threshold"] = float(opt["threshold"])

        if "threshold-peakcount" in opt:
            opt["threshold-peakcount"] = int(opt["threshold-peakcount"])

    except getopt.GetoptError as err:
        print(err)
        sys.exit(2)
    except IndexError:
        print("Usage: msp430-etv <duration>")
        sys.exit(2)
    except ValueError:
        print("Error: duration or skip is not a number")
        sys.exit(2)

    if "load" in opt:
        if ".xz" in opt["load"]:
            import lzma

            with lzma.open(opt["load"], "rt") as f:
                log_data = f.read()
        else:
            with open(opt["load"], "r") as f:
                log_data = f.read()
    else:
        log_data = measure_data(opt["save"], 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 - opt["skip"], 4))

    energy_overflow_count = 0
    prev_total_energy = 0
    for i, line in enumerate(data_lines):
        if i >= opt["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 - opt["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 "threshold-peakcount" in opt:
        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 == opt["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
                )
            )
            opt["threshold"] = threshold
        else:
            print("Found no working threshold")

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

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

        peaks = []
        peak_start = -1
        for i, dp in enumerate(power):
            if dp >= opt["threshold"] and peak_start == -1:
                peak_start = i
            elif dp < opt["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 "stat" in opt:
        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 "plot" in opt:
        if opt["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 opt["plot"] == "I":
            # 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 "load" in opt:
            plt.title(opt["load"])
        plt.show()

    if "histogram" in opt:
        bin_count = int(opt["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()