#!/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/b8e08afb-818f-4d2d-9d28-5afe8fc76a32'
# ?endDate=2017-08-23&granularity=Minute&metrics=read&startDate=2017-08-23&units=W'
logouturl = 'https://command.aglsolar.com.au/api/v2/Account/Logout'

class UTC(datetime.tzinfo):
    def utcoffset(self, dt):
        return datetime.timedelta(0)

    def dst(self, dt):
        return datetime.timedelta(0)

    def tzname(self, dt):
        return "UTC"


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', type = str)
    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 (args.start is None) ^ (args.end is None):
        parser.error('Must specify start and end or neither')

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

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

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

    dbh = sqlite3.connect(dbfn, detect_types = sqlite3.PARSE_DECLTYPES)
    cur = dbh.cursor()
    if args.update:
        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, date, date)
            if data == None:
                #print('Getting new token')
                token = gettoken(username, password)
                data = getdata(token, 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', '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'
}

names = {
    'battery_charge' : 'Battery Charge',
    'battery_power' : 'Battery Power',
    'power_consumed' : 'Power Consumed',
    'power_expected' : 'Power Expected',
    'power_exported' : 'Power Expected',
    '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)

    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 naive UTC to proper UTC then adjust to local time
    xdata = map(lambda f: f.replace(tzinfo = utc).astimezone(lt), ary[:,0])
    for idx in range(len(cols)):
        if units[cols[idx]] == yaxisunits1:
            ax1lines.append([xdata, ary[:,idx + 1], names[cols[idx]], colourlist[colouridx]])
        else:
            ax2lines.append([xdata, ary[:,idx + 1], names[cols[idx]], colourlist[colouridx]])
        colouridx += 1

    if fname == None:
        import matplotlib.pylab
        fig = matplotlib.pylab.figure()
    else:
        import matplotlib.backends.backend_agg
        fig = matplotlib.figure.Figure(figsize = (8, 6), dpi = 75)

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

    for line in ax1lines:
        ax1.plot(line[0], line[1], label = line[2])

    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])
        ax2.legend(loc = 'upper right')

    ndays = int(max(1, round(((end - start).total_seconds()) / 86400)))
    for ax in fig.get_axes():
        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(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, startdate, enddate):
    #print('getting ' + startdate.strftime('%Y-%m-%d'))
    reply = requests.request('GET', dataurl, 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()