#!/usr/bin/env python

"""Calculating Tokyo emissions."""

from __future__ import print_function
import argparse
import os.path
import sys
import math
import collections
from osgeo import gdal
import pylab as pl
import matplotlib.colors
import json

datafiles = [
    os.path.join(
        "odiac", "odiac2016_1km_excl_intl_{:02d}{:02d}_rev1.tokyo.tif".format(year, month))
    for year in range(11,16) for month in range(1,13)
]

parser = argparse.ArgumentParser("Calculating Tokyo Emissions")
parser.add_argument("-p", "--print-required-datafiles", action="store_true")
parser.add_argument("-o", "--outfile", default="odiac_emissions.json")
args = parser.parse_args()

# show the required datafiles let them be prepared from a Makefile
if args.print_required_datafiles:
    print(" ".join(datafiles))
    sys.exit(0)

emission_sums_fromtokyo = collections.OrderedDict()
emission_sums_bg = collections.OrderedDict()
emission_angles_sum = collections.Counter()
emission_angles_count = collections.Counter()
emission_angles = []


for d in datafiles:
    D = gdal.Open(d)
    data = D.ReadAsArray()
    ullon, dx2lon, dy2lon, ullat, dx2lat, dy2lat = D.GetGeoTransform()
    Dx = D.RasterXSize
    Dy = D.RasterYSize
    print(ullon, dx2lon, dy2lon, ullat, dx2lon, dy2lat, Dx, Dy, ullon + dx2lon*Dx, ullat + dy2lat*Dy)
    lons = pl.zeros((Dx,))
    lats = pl.zeros((Dy,))
    for x in range(Dx):
        lons[x] = ullon + x*dx2lon
    for y in range(Dy):
        lats[y] = ullat + y*dy2lat
    # find Tsukuba
    lola = tsukuba_lonlat = 140.1215, 36.0513
    tsukuba_xy = (pl.searchsorted(lons, lola[0]),
                  -pl.searchsorted(lats, lola[1], sorter=range(len(lats)-1, -1, -1)))
    print (lola, tsukuba_xy, lons[tsukuba_xy[0]], lats[tsukuba_xy[1]])
    # mask all values not from inside Tokyo as defined in plot_tccon_by_direction.py
    masknotokyo = pl.zeros((Dy, Dx))
    masknobg = pl.zeros((Dy, Dx))
    tsukuba_tokyo_start_stop = 170, 240
    bg_start_stop = 270, 30
    for x in range(Dx):
        for y in range(Dy):
            dlon, dlat = lons[x] - tsukuba_lonlat[0], lats[y] - tsukuba_lonlat[1]
            # Angle from top, moving clockwards
            # North-East of Tsukuba
            if dlat > 0 and dlon > 0:
                angle = 90 - 360/2/math.pi * math.asin(dlat / math.sqrt(dlon**2 + dlat**2))
            # South-East of Tsukuba
            elif dlat <= 0 and dlon > 0:
                angle = 90 + 360/2/math.pi * math.asin(-dlat / math.sqrt(dlon**2 + dlat**2))
            # South-West of Tsukuba
            elif dlat <= 0 and dlon <= 0:
                angle = 270 - 360/2/math.pi * math.asin(-dlat / math.sqrt(dlon**2 + dlat**2))
            # North-West of Tsukuba
            elif dlat > 0 and dlon <= 0:
                angle = 270 + 360/2/math.pi * math.asin(dlat / math.sqrt(dlon**2 + dlat**2))
            else:
                raise ValueError("Missed Possibility in angle-if.")
            anglebinsize = 1
            anglebin = anglebinsize*int(angle//anglebinsize)
            emission_angles_sum[anglebin] += data[y,x]
            emission_angles_count[anglebin] += 1
            # emission_angles.append((angle, float(data[y,x])))
            if angle < tsukuba_tokyo_start_stop[0] or angle > tsukuba_tokyo_start_stop[1]:
                masknotokyo[y,x] = 1
            if not (angle > bg_start_stop[0] or angle < bg_start_stop[1]):
                masknobg[y,x] = 1
    data_fromtokyo = pl.ma.masked_array(data, mask=masknotokyo)
    data_frombg = pl.ma.masked_array(data, mask=masknobg)
    emission_sums_fromtokyo[d] = float(data_fromtokyo.sum())
    emission_sums_bg[d] = float(data_frombg.sum())
    if d == "odiac/odiac2016_1km_excl_intl_1501_rev1.tokyo.tif":
        pl.imshow(data, norm=matplotlib.colors.SymLogNorm(linthresh=0.1, linscale=1))
        pl.title("ODIAC emissions around Tokyo 2015-01")
        pl.plot(tsukuba_xy[0], Dy+tsukuba_xy[1], "xr", label="Tsukuba")
        pl.text(tsukuba_xy[0], Dy+tsukuba_xy[1] - 5, "Tsukuba", color="r", horizontalalignment='left', fontsize=14, alpha=0.9)
        pl.xlabel("geotiff pixels")
        pl.ylabel("geotiff pixels")
        pl.colorbar().set_label("carbon emissions / log$_{10}$ tonnes C month$^{-1}$")
        pl.savefig("plot-tokyo-2015-01.pdf", bbox_inches='tight', transparent=False)

for k in emission_angles_sum:
    # need to cut out the numpy values
    emission_angles_sum[k] = float(emission_angles_sum[k] / len(datafiles))
# angles, values = zip(*sorted(emission_angles_sum.items()))
# pl.plot(angles, values)
# pl.show()


with open(args.outfile, "w") as f:
    json.dump({"byfile_sum_fromtokyo_ton_C_per_month": list(emission_sums_fromtokyo.items()),
               "byfile_sum_bg_ton_C_per_month": list(emission_sums_bg.items()),
               "byangle_sum_ton_C_per_month": sorted(emission_angles_sum.items())},
              # "emission_angles_ton_C_per_month": emission_angles},
              f)
