#!/usr/bin/env python

"""Plotting ODIAC emissions for Tokyo."""

from __future__ import print_function
import argparse
import pylab as pl
import matplotlib.colors
import matplotlib
import json
import scipy.ndimage.filters

parser = argparse.ArgumentParser("Calculating Tokyo Emissions")
parser.add_argument('datafiles', metavar='N', type=str, nargs='*',
                    default=['odiac_emissions.json'],
                    help='datafiles to parse')
args = parser.parse_args()


# moving average from https://matplotlib.org/examples/pylab_examples/finance_work2.html
def moving_average(x, n, type='simple'):
    """
    compute an n period moving average.

    type is 'simple' | 'exponential'

    """
    x = pl.asarray(x)
    if type == 'simple':
        weights = pl.ones(n)
    else:
        weights = pl.exp(pl.linspace(-1., 0., n))

    weights /= weights.sum()

    a = pl.convolve(x, weights, mode='full')[:len(x)]
    a[:n] = a[n]
    return a

for i in args.datafiles:
    with open(i) as f:
        data = json.load(f)
        fromtokyo = data['byfile_sum_fromtokyo_ton_C_per_month']
        bg = data['byfile_sum_bg_ton_C_per_month']
        byangle_sum_ton_C_per_month = data['byangle_sum_ton_C_per_month']
        # emission_angles_ton_C_per_month = data['emission_angles_ton_C_per_month']
    # pl.plot(*zip(*emission_angles_ton_C_per_month), marker="x")
    angles, emiss = zip(*byangle_sum_ton_C_per_month)
    pl.plot(angles, [i/1000. for i in emiss], color="blue", label="aggregated ODIAC", linewidth=0.5)
    # moving average
    sigma = 5
    pl.plot(angles, scipy.ndimage.filters.gaussian_filter1d([i/1000. for i in emiss], sigma, mode="wrap"), color="cyan", label="with gaussian filter, $\sigma = %d^{\circ}$" % sigma)
    pl.axvline(170, color="brown")
    pl.axvline(240, color="brown")
    pl.text((170+240)/2.0, 9, "from\nTokyo\narea", color="k", horizontalalignment='center', verticalalignment="top", fontsize=14, alpha=0.5)
    print(sum(j for i,j in byangle_sum_ton_C_per_month) * 12)
    from_tokyo_area = sum(j for i,j in byangle_sum_ton_C_per_month if i > 170 and i < 240) * 12
    background_nw = sum(j for i,j in byangle_sum_ton_C_per_month if i > 300) * 12
    print(from_tokyo_area, from_tokyo_area - background_nw*(70./60))
    pl.title("ODIAC emissions by wind direction")
    pl.ylabel("carbon emissions / kilotons per month")
    pl.xlabel("wind from direction / angle")
    # add measured emissions
    with open("tccon-data/tk20110804_20160330.public-tsukuba-residualsxco2.json") as f:
        meas = json.load(f)
    bgmedian = meas["bgmedian"]
    gms2ktmo = 0.646114909090909 # (define (mo x) (let ((s/month 2592000)(d_rect 914))(/ (* x (/ 12 44) d_rect s/month) 1e9))) (mo 1.)
    pl.plot(meas["fromtokyo_wdirs"],
            (pl.array(meas["fromtokyo_medians"]) - bgmedian)*gms2ktmo, linestyle="", marker="x", color="orange", alpha=0.6, label="")
    pl.plot(meas["wdirs_outer"],
            (pl.array(meas["medians_outer"]) - bgmedian)*gms2ktmo, linestyle="", marker="x", color="orange", alpha=0.6, label="measured - background")
    # pl.plot((pl.array(meas["fromtokyo_wdirs"])+15)%360, pl.array(meas["fromtokyo_medians"]) - bgmedian, linestyle="", marker="x", color="magenta", label="")
    # pl.plot((pl.array(meas["wdirs_outer"])+15)%360, pl.array(meas["medians_outer"]) - bgmedian, linestyle="", marker="x", color="magenta", label="m. meas - bg, shifted +15 degree")
    pl.legend(loc='upper left', fancybox=True, framealpha=0.5)
    pl.gcf().set_size_inches(4.4,3.3)
    pl.savefig("odiac-emissions-by-angle.pdf", bbox_inches='tight', transparent=False)
    pl.close()
    
    def path_to_yearmonth(path):
        "Extract the yearmonth from the path."
        idx = path.index("intl")
        return path[idx+5:idx+9]
    pl.plot(range(len(fromtokyo)), [i*12 for i in list(zip(*fromtokyo))[1]], label="from Tokyo area")
    pl.plot(range(len(bg)), [i*12 for i in list(zip(*bg))[1]], label="from background")
    pl.gca().set_xticklabels(map(path_to_yearmonth,
                                 pl.array(list(zip(*fromtokyo))[0])[
                                     [int(min(len(fromtokyo)-1, i))
                                      for i in matplotlib.pyplot.xticks()[0]]]))
    pl.legend()
    pl.gcf().set_size_inches(4.4,3.3)
    pl.savefig("odiac-emissions-by-month.pdf", bbox_inches='tight', transparent=False)
