#!/usr/bin/env python

import argparse
import ConfigParser
import datetime
import dateutil
import exceptions
import json
import os
import requests
import sqlite3
import sys
import tzlocal

loginurl = 'https://command.aglsolar.com.au/api/v2/Account/LoginUser'
dataurl = 'https://command.aglsolar.com.au/api/v2/graph/'
# ?endDate=2017-08-23&granularity=Minute&metrics=read&startDate=2017-08-23&units=W'
logouturl = 'https://command.aglsolar.com.au/api/v2/Account/Logout'

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('-u', '--update', help = 'Update data', action="store_true")
    parser.add_argument('-g', '--graph', help = 'Produce graph', action="store_true")
    parser.add_argument('-f', '--filename', help = 'Filename to save graph as (uses strftime on start)', 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)

    args = parser.parse_args()

    conf = ConfigParser.ConfigParser()
    confname = os.environ['HOME'] + '/.agl.ini'
    conf.read(confname)
    username = conf.get('DEFAULT', 'username')
    password = conf.get('DEFAULT', 'password')
    dbfn = conf.get('DEFAULT', 'db')

    if not args.update and not args.graph:
        parser.error('Nothing to do')

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

    start = args.start
    if start is None:
        start = datetime.date.today()
        start = datetime.datetime(start.year, start.month, start.day)

    if args.days is not None:
        start -= datetime.timedelta(days = args.days)

    end = args.end
    if end is None:
        end = start + datetime.timedelta(days = 1)
        end = datetime.datetime(end.year, end.month, end.day)

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

    dbh = sqlite3.connect(dbfn, detect_types = sqlite3.PARSE_DECLTYPES)
    cur = dbh.cursor()
    if args.update:
        installid = conf.get('DEFAULT', 'installid')
        date = start
        while date < end:
            if conf.has_option('DEFAULT', 'token'):
                token = conf.get('DEFAULT', 'token')
            else:
                token = gettoken(username, password)
                conf.set('DEFAULT', 'token', token)
                conf.write(file(confname, 'w'))

            data = getdata(token, installid, date, date)
            if data == None:
                #print('Getting new token')
                token = gettoken(username, password)
                data = getdata(token, installid, date, date)
            if data == None:
                print('Unable to fetch data')
            updatedb(cur, data)
            dbh.commit()
            date += datetime.timedelta(days = 1)

    if args.graph:
        graph(args.filename, cur,
                  ['battery_charge', 'battery_power', 'power_imported', 'power_exported', 'power_consumed', 'power_generated'],
                  start, end)

def mkdb(cur):
    cur.execute('''
CREATE TABLE IF NOT EXISTS agl (
    t_stamp	TIMESTAMP PRIMARY KEY,
    battery_charge NUMBER,
    battery_power NUMBER,
    power_consumed NUMBER,
    power_expected NUMBER,
    power_exported NUMBER,
    power_generated NUMBER,
    power_imported NUMBER,
    estimated_savings NUMBER,
    pv_forecast NUMBER,
    pv_gen_battery NUMBER,
    pv_gen_grid NUMBER,
    pv_gen_site NUMBER,
    site_cons_battery NUMBER,
    site_cons_grid NUMBER,
    site_cons_pv NUMBER
)''')

units = {
    'battery_charge' : '%',
    'battery_power' : 'Watt',
    'power_consumed' : 'Watt',
    'power_expected' : 'Watt',
    'power_exported' : 'Watt',
    'power_generated' : 'Watt',
    'power_imported' : 'Watt',
    'estimated_savings' : '$',
    'pv_forecast' : 'Watt',
    'pv_gen_battery' : 'Watt',
    'pv_gen_grid' : 'Watt',
    'pv_gen_site' : 'Watt',
    'site_cons_battery' : 'Watt',
    'site_cons_grid' : 'Watt',
    'site_cons_pv' : 'Watt'
}

convs = {
    'battery_power' : lambda a: a / 1000.0,
}

scale_limits = {
    'battery_charge' : (0, 100),
}

tarrifs = {
    'power_exported' : 0.163,
    'power_imported' : 0.380,
    'power_generated' : 0.163,
    'power_consumed' : 0.380,
}

names = {
    'battery_charge' : 'Battery Charge',
    'battery_power' : 'Battery Power',
    'power_consumed' : 'Power Consumed',
    'power_expected' : 'Power Expected',
    'power_exported' : 'Power Exported',
    'power_generated' : 'Power Generated',
    'power_imported' : 'Power Imported',
    'estimated_savings' : 'Estimated Savings',
    'pv_forecast' : 'PV Forecast',
    'pv_gen_battery' : 'PV Generation Battery',
    'pv_gen_grid' : 'PV Generation Grid',
    'pv_gen_site' : 'PV Generation Site',
    'site_cons_battery' : 'Site Consumption Batter',
    'site_cons_grid' : 'Site Consumption Grid',
    'site_cons_pv' : 'Site Consumption PV'
}

def graph(fname, cur, cols, start, end):
    import numpy
    import matplotlib
    import matplotlib.dates

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

    # Work out what axes we are using
    yaxisunits1 = None
    yaxisunits2 = None
    ax1lines = []
    ax2lines = []
    colouridx = 0
    for col in cols:
        unit = units[col]
        if yaxisunits1 == None:
            yaxisunits1 = unit
        if yaxisunits2 == None:
            if unit != yaxisunits1:
                yaxisunits2 = unit
        else:
            if unit != yaxisunits1 and unit != yaxisunits2:
                raise exceptions.Exception('Asked to graph >2 different units')

    ltname = tzlocal.get_localzone().zone # Why is this so hard..
    lt = dateutil.tz.gettz(ltname)
    utc = dateutil.tz.gettz('UTC')
    matplotlib.rcParams['timezone'] = ltname

    if start.tzinfo == None:
        start = start.replace(tzinfo = lt)
    if end.tzinfo == None:
        end = end.replace(tzinfo = lt)
    startlt = start
    endlt = end
    start = start.astimezone(utc)
    end = end.astimezone(utc)

    # Actually get the data
    colstr = reduce(lambda a, b: a + ', ' + b, cols)
    # Data is stored as naive datetime's which are in UTC so convert the requested time here
    cur.execute('SELECT t_stamp, ' + colstr + ' FROM agl WHERE t_stamp > ? AND t_stamp < ? ORDER BY t_stamp',
                    (start, end))
    ary = numpy.array(cur.fetchall())
    if ary.shape[0] == 0:
        print('No data')
        return
    # Convert TZ naive UTC to TZ aware UTC then adjust to local time
    xdata = map(lambda f: f.replace(tzinfo = utc).astimezone(lt), ary[:,0])
    xhours = matplotlib.dates.date2num(xdata) * 24
    xdeltas = xhours[1:] - xhours[0:-1]
    calcd = {}
    for idx in range(len(cols)):
        col = cols[idx]
        ydata = ary[:,idx + 1]

	ydata = [0 if v is None else v for v in ydata]
        if col in convs:
	    ydata = map(convs[col], ydata)
        if col in scale_limits:
            scale_min = scale_limits[col][0]
            scale_max = scale_limits[col][1]
        else:
            scale_min = None
            scale_max = None

        if col in tarrifs:
            calc = (ydata[1:] * xdeltas).sum() / 1000.0 * tarrifs[col]
            annotation = '%s: $%.2f' % (names[col], calc)
            calcd[col] = calc
        else:
            annotation = None

        if units[col] == yaxisunits1:
            ax = ax1lines
        else:
            ax = ax2lines
        ax.append([xdata, ydata, names[col], colourlist[colouridx], scale_min, scale_max, annotation])
        colouridx += 1

    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)

    ax1 = fig.add_subplot(111)
    ax1.set_ylabel(yaxisunits1)

    annotations = []
    for line in ax1lines:
        ax1.plot(line[0], line[1], label = line[2], color = line[3])
        if line[4] != None and line[5] != None:
            ax1.set_ylim((line[4], line[5]))
        if line[6] != None:
            annotations.append(line[6])
    ax1.legend(loc = 'upper left')

    if yaxisunits2 != None:
        ax2 = ax1.twinx()
        ax2.set_ylabel(yaxisunits2)

        for line in ax2lines:
            ax2.plot(line[0], line[1], label = line[2], color = line[3])
            if line[4] != None and line[5] != None:
                ax2.set_ylim(bottom = line[4], top = line[5])
            if line[6] != None:
                annotations.append(line[6])

        ax2.legend(loc = 'upper right')
    if all(map(lambda x: x in calcd, ('power_imported', 'power_exported', 'power_consumed', 'power_generated'))):
        annotations.append('Saved: $%0.2f' % ((calcd['power_consumed'] - calcd['power_generated']) -
                                                  (calcd['power_imported'] - calcd['power_exported'])))
    ax1.text(0.02, 0.9, reduce(lambda a, b: a + '\n' + b, annotations),
                 transform = ax1.transAxes, bbox = dict(facecolor = 'red', alpha = 0.5),
                 ha = 'left', va = 'top')
    ndays = int(max(1, round(((end - start).total_seconds()) / 86400)))
    for ax in fig.get_axes():
        if (endlt - startlt).total_seconds() > 86400:
            ax.set_title('%s to %s' % (startlt.strftime('%Y-%m-%d'), endlt.strftime('%Y-%m-%d')))
        else:
            ax.set_title('%s' % (startlt.strftime('%Y-%m-%d')))
        ax.set_xlim([start, end])
        ax.format_xdata = lambda d: matplotlib.dates.num2date(d).strftime('%d %b %H:%M')
        ax.xaxis.grid(True)
        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(startlt.strftime(fname))

def updatedb(cur, data):
    mkdb(cur)
    for d in data['reads']['data']:
        ts = datetime.datetime.strptime(d['t_stamp'], '%Y-%m-%dT%H:%M:%SZ')
        # Note we rename *energy* to *power* here to match what it actually means
        vals = [ts, d['battery_charge'], d['battery_energy'], d['energy_consumed'], d['energy_expected'], d['energy_exported'], d['energy_generated'],
                    d['energy_imported'], d['estimated_savings'], d['pv_forecast'], d['pv_generation']['battery_energy'],
                    d['pv_generation']['grid_energy'], d['pv_generation']['site_energy'], d['site_consumption']['battery_energy'],
                    d['site_consumption']['grid_energy'], d['site_consumption']['pv_energy']]
        skip = True
        for v in vals[1:]:
            if v != None:
                skip = False
                break
        if skip:
            print('Skipping empty record at ' + str(ts))
            continue
        cur.execute('INSERT OR IGNORE INTO agl(t_stamp, battery_charge, battery_power, power_consumed, power_expected, power_exported, power_generated, power_imported, estimated_savings, pv_forecast, pv_gen_battery, pv_gen_grid, pv_gen_site, site_cons_battery, site_cons_grid, site_cons_pv) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)', vals)

def gettoken(username, password):
    authblob = json.encoder.JSONEncoder().encode({'email' : username, 'password' : password})
    reply = requests.request('POST', loginurl, data = authblob, headers = {'Content-Type' : 'application/json'})
    if reply.status_code != 200:
        return None
    return json.decoder.JSONDecoder().decode(reply.content)['access_token']

def getdata(token, installid, startdate, enddate):
    #print('getting ' + startdate.strftime('%Y-%m-%d'))
    reply = requests.request('GET', dataurl + installid, params = {
        'startDate' :	startdate.strftime('%Y-%m-%d'),
        'endDate' :		enddate.strftime('%Y-%m-%d'),
        'granularity' :	'Minute',
        'metrics' :		'read',
        'units' :		'W',
        }, headers = { 'Authorization' : 'Bearer ' + token})

    if reply.status_code != 200:
        return None

    return json.decoder.JSONDecoder().decode(reply.content)

def logout(token):
    reply = requests.request('GET', logouturl, headers = { 'Authorization' : 'Bearer ' + token})
    return reply.status_code == 200

if __name__ == '__main__':
    main()