#!/usr/bin/env python

import argparse
import datetime
import dateutil
import exceptions
import matplotlib
import matplotlib.dates
import numpy
import os
import requests
import sqlite3
import tzlocal

class Column(object):
    def __init__(self, rowname, title, table, units, limits = (None, None), conv = None, annofn = None):
        self.rowname = rowname
        self.title = title
        self.table = table
        self.units = units
        self.limits = limits
        self.conv = conv
        self.annofn = annofn

columns = [
    Column('main_voltage', 'Battery Voltage', 'eprolog', 'Vdc'),
    Column('aux_voltage', 'Aux Voltage', 'eprolog', 'Vdc'),
    Column('battery_curr', 'Battery Current', 'eprolog', 'A'),
    Column('amp_hours', 'Battery Amp Hours', 'eprolog', 'Ah'),
    Column('state_of_charge', 'State of Charge', 'eprolog', '%', (0, 100), annofn = lambda xdata, ydata: 'DoD: %.1f' % (100 - ydata.min())),
    Column('time_remaining', 'Time Remaining', 'eprolog', 'min'),
    Column('battery_temp', 'Battery Temperature', 'eprolog', 'C'),

    Column('ac_act_power', 'Active Power', 'giantlog', 'W'),
    Column('ac_app_power', 'Apparent Power', 'giantlog', 'W'),
    Column('ac_frequency', 'AC Frequency', 'giantlog', 'Hz'),
    Column('ac_volts', 'AC Voltage', 'giantlog', 'Vac'),
    Column('batt_chr_curr', 'Discharge Current', 'giantlog', 'A'),
    Column('batt_dis_curr', 'Charge Current', 'giantlog', 'A'),
    Column('battery_cap', 'Battery Capacity', 'giantlog', '%', (0, 100)),
    Column('battery_volts', 'Battery Voltage', 'giantlog', 'Vdc'),
    Column('grid_frequency', 'Grid Frequency', 'giantlog', 'Hz'),
    Column('grid_volts', 'Grid Voltage', 'giantlog', 'Vac'),
    Column('hs_temperature', 'HS Temperature', 'giantlog', 'C'),
    Column('load_pct', 'Load', 'giantlog', '%', (0, 100)),
]

def valid_date(s):
    try:
        return datetime.datetime.strptime(s, "%Y-%m-%d")
    except ValueError:
        raise argparse.ArgumentTypeError("Not a valid date: '{0}'.".format(s))

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-f', '--filename', help = 'Path to database', type = str, required = True)
    parser.add_argument('-g', '--graphfn', help = 'File to write graph to', type = str)
    parser.add_argument('-d', '--days', help = 'Days ago to graph', type = int)
    parser.add_argument('-s', '--start', help = 'Start date for graph (YYYY-MM-DD)', type = valid_date)
    parser.add_argument('-e', '--end', help = 'End date for graph (YYYY-MM-DD)', type = valid_date)
    parser.add_argument('-c', '--column', help = 'Column to plot (can be specified multiple times)', type = str, action = 'append')
    parser.add_argument('-z', '--timezone', help = 'Timezone to operate in in (Oslon format) default: localtime)', type = str)

    args = parser.parse_args()

    if args.days is not None and args.days < 0:
        parser.error('days must be non-negative')

    selector = [args.start is not None, args.end is not None, args.days is not None]
    if selector == [True, True, False]: # Start and end
        pass
    elif selector == [True, False, True] or selector == [True, False, False]: # Start and days or start
        if args.days == None:
            args.days = 1
        args.end = args.start + datetime.timedelta(days = args.days)
    elif selector == [False, True, True] or selector == [False, True, False]: # End and days or end
        if args.days == None:
            args.days = 1
        args.start = args.end - datetime.timedelta(days = args.days)
    elif selector == [False, False, True]: # Days
        args.end = datetime.date.today() + datetime.timedelta(days = 1)
        args.end = datetime.datetime(args.end.year, args.end.month, args.end.day)
        args.start = args.end - datetime.timedelta(days = args.days)
    elif selector == [False, False, False]: # Nothing
        args.start = datetime.date.today()
        args.start = datetime.datetime(args.start.year, args.start.month, args.start.day)
        args.end = args.start + datetime.timedelta(days = 1)
    else:
        parser.error('can\'t specify days, start and end simultaneously')

    if args.start >= args.end:
        parser.error('Start must be before end')

    cols = args.column
    if cols == None:
        cols = ['state_of_charge', 'load_pct', 'main_voltage', 'aux_voltage']

    dbh = sqlite3.connect(args.filename, detect_types = sqlite3.PARSE_DECLTYPES)
    cur = dbh.cursor()

    # Get local timezone name and convert start/end to it
    # Why is this so hard...
    if args.timezone is None:
        args.timezone = tzlocal.get_localzone().zone

    lt = dateutil.tz.gettz(args.timezone)
    if lt == None:
        parser.error('Unknown timezone')

    utc = dateutil.tz.gettz('UTC')
    matplotlib.rcParams['timezone'] = args.timezone

    if args.start.tzinfo == None:
        args.start = args.start.replace(tzinfo = lt)
    if args.end.tzinfo == None:
        args.end = args.end.replace(tzinfo = lt)
    graph(args.graphfn, cur, cols, args.start, args.end, lt, utc)

def graph(fname, cur, _cols, start, end, lt, utc):
    import numpy
    import matplotlib
    import matplotlib.dates

    startepoch = int(start.strftime('%s'))
    endepoch = int(end.strftime('%s'))

    colourlist = ['b','g','r','c','m','y','k']

    cols = []

    yaxisunits1 = None
    yaxisunits2 = None
    ax1lines = []
    ax2lines = []
    colouridx = 0
    for col in _cols:
        # Check the column exists
        for c in columns:
            if col == c.rowname:
                cols.append(c)
                break
        else:
            raise exceptions.Exception('Unknown column name ' + c)

        # Work out what axes we are using
        if yaxisunits1 == None:
            yaxisunits1 = c.units
        if yaxisunits2 == None:
            if c.units != yaxisunits1:
                yaxisunits2 = c.units
        else:
            if c.units != yaxisunits1 and c.units != yaxisunits2:
                raise exceptions.Exception('Asked to graph >2 different units')

    for c in cols:
        # Get the data
        cur.execute('SELECT tstamp, ' + c.rowname + ' FROM ' + c.table + ' WHERE tstamp > ? AND tstamp < ? ORDER BY tstamp',
                    (startepoch, endepoch))
        ary = numpy.array(cur.fetchall())
        if ary.shape[0] == 0:
            print('No data for ' + c.rowname)
            continue

        # Convert epoch stamp to datetime with UTC timezone
        c.xdata = map(lambda f: datetime.datetime.utcfromtimestamp(f).replace(tzinfo = utc), ary[:,0])
        c.ydata = ary[:,1]
        if c.conv != None:
            c.ydata = map(c.conv, c.ydata)

        scale_min, scale_max = c.limits

        # DoD?
        if c.annofn != None:
            c.annotation = c.annofn(c.xdata, c.ydata)
        else:
            c.annotation = None

        # Work out which axis to plot on
        if c.units == yaxisunits1:
            ax = ax1lines
        else:
            ax = ax2lines
        c.colour = colourlist[colouridx]
        colouridx += 1
        ax.append(c)

    if len(ax1lines) == 0 and len(ax2lines) == 0:
        print('Nothing to plot')
        return

    # Load the right backend for display or save
    if fname == None:
        import matplotlib.pylab
        fig = matplotlib.pylab.figure()
    else:
        import matplotlib.backends.backend_agg
        fig = matplotlib.figure.Figure(figsize = (12, 6), dpi = 75)

    # Do the plot
    ax1 = fig.add_subplot(111)
    ax1.set_ylabel(yaxisunits1)

    annotations = []
    for line in ax1lines:
        ax1.plot(line.xdata, line.ydata, label = line.title, color = line.colour)
        if line.limits[0] != None or line.limits[1] != None:
            ax1.set_ylim(line.limits[0], line.limits[1])
        if line.annotation != None:
            annotations.append(line.annotation)
    ax1.legend(loc = 'center left')

    if len(ax2lines) > 0:
        ax2 = ax1.twinx()
        ax2.set_ylabel(yaxisunits2)

        for line in ax2lines:
            ax2.plot(line.xdata, line.ydata, label = line.title, color = line.colour)
            if line.limits[0] != None or line.limits[1] != None:
                ax2.set_ylim(line.limits[0], line.limits[1])
            if line.annotation != None:
                annotations.append(line.annotation)

        ax2.legend(loc = 'center right')

    if len(annotations) > 0:
        ax1.text(0.02, 0.3, reduce(lambda a, b: a + '\n' + b, annotations),
                    transform = ax1.transAxes, bbox = dict(facecolor = 'blue', alpha = 0.5),
                    ha = 'left', va = 'top')
    ndays = int(max(1, round(((end - start).total_seconds()) / 86400)))
    once = False
    for ax in fig.get_axes():
        if not once:
            once = True
            if ndays > 1:
                ax.set_title('%s to %s' % (start.astimezone(lt).strftime('%Y-%m-%d'), end.astimezone(lt).strftime('%Y-%m-%d')))
            else:
                ax.set_title('%s' % (start.astimezone(lt).strftime('%Y-%m-%d')))
        ax.set_xlim([start.astimezone(utc), end.astimezone(utc)])
        ax.xaxis.grid(True)
        ax.xaxis.axis_date(lt)
        ax.xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%d %b\n%H:%M'))
        ax.xaxis.set_major_locator(matplotlib.dates.HourLocator(interval = 2 * ndays))
        ax.xaxis.set_minor_locator(matplotlib.dates.MinuteLocator(interval = 5 * ndays))
        for label in ax.get_xticklabels():
            label.set_ha('center')
            label.set_rotation(90)

    # Fudge margins to give more graph and less space
    fig.subplots_adjust(left = 0.10, right = 0.88, top = 0.95, bottom = 0.15)
    if fname == None:
        matplotlib.pyplot.show()
    else:
        canvas = matplotlib.backends.backend_agg.FigureCanvasAgg(fig) # Sets canvas in fig too
        fig.savefig(startdt.strftime(fname))

if __name__ == '__main__':
    main()