#!/usr/bin/env python3

import getopt
import matplotlib.pyplot as plt
import numpy as np
import os
import re
import subprocess
import sys

opt = dict()

def measure_data(time):
	if not 'LD_LIBRARY_PATH' in os.environ:
		os.environ['LD_LIBRARY_PATH'] = '{}/var/projects/msp430/MSP430Flasher_1.3.7'.format(os.environ['HOME'])

	energytrace_cmd = '{}/var/source/energytrace-util/energytrace'.format(os.environ['HOME'])

	res = subprocess.run([energytrace_cmd, str(duration)], stdout = subprocess.PIPE, universal_newlines = True)

	return res.stdout

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

USAGE

msp430-etv [--load <file> | <measurement duration>] [--save <file>]
	[--skip <count>] [--plot] [--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.

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
  --plot
    Show power/time plot
  --stat
    Show mean voltage, current, and power as well as total energy consumption.
	''')

if __name__ == '__main__':
	try:
		optspec = ('help load= save= skip= plot stat')
		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 'skip' in opt:
			opt['skip'] = int(opt['skip'])
		else:
			opt['skip'] = 0

	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:
		with open(opt['load'], 'r') as f:
			log_data = f.read()
	else:
		log_data = measure_data(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))

	for i, line in enumerate(data_lines):
		if i >= opt['skip']:
			timestamp, current, voltage, total_energy = map(float, line.split(' '))
			data[i - opt['skip']] = [timestamp, current, voltage, total_energy]

	m_duration = data[-1, 0] - data[0, 0]
	m_energy = data[-1, 3] - data[0, 3]
	m_calc_energy = np.sum(data[1:, 1] * data[1:, 2] * (data[1:, 0] - data[:-1, 0]))
	m_energy_deviation = np.abs(m_energy - m_calc_energy) / np.max([m_energy, m_calc_energy])

	print('{:d} measurements in {:.2f} s = {:.0f} Hz sample rate'.format(
		data_count, m_duration, data_count / m_duration))

	print('Reported energy: E = {:f} J'.format(m_energy))
	print('Calculated energy: U*I*t = {:f} J'.format(m_calc_energy))
	print('Energy deviation: {:.1f}%'.format(m_energy_deviation * 100))

	if 'save' in opt:
		with open(opt['save'], 'w') as f:
			f.write(log_data)

	if 'stat' in opt:
		print('Mean voltage: {:f}'.format(np.mean(data[:, 2])))
		print('Mean current: {:f}'.format(np.mean(data[:, 1])))
		print('Mean power: {:f}'.format(np.mean(data[:, 1] * data[:, 2])))
		print('Total energy: {:f}'.format(m_energy))

	if 'plot' in opt:
		pwrhandle, = plt.plot(data[:, 0], data[:, 1] * data[:, 2], 'b-', label='U*I', markersize=1)
		#energyhandle, = plt.plot(data[1:, 0], (data[1:, 3] - data[:-1, 3]) / (data[1:, 0] - data[:-1, 0]), 'r-', label='E/Δt', markersize=1)
		plt.legend(handles=[pwrhandle])
		plt.xlabel('Time [s]')
		plt.ylabel('Power [W]')
		plt.grid(True)
		plt.show()