# -*- coding: utf-8 -*-
# Copyright 2016 Earth Sciences Department, BSC-CNS

"""nc2geojson module.

This module provide conversion from netCDF to geojson format.

"""

from netCDF4 import Dataset as nc_file
import geopandas as gpd
import pandas as pd
from shapely import geometry
import numpy as np
import json
import orjson
import os
import os.path
import sys
from datetime import datetime
from dateutil.relativedelta import relativedelta

np.set_printoptions(precision=2)

DIR_PATH = os.path.abspath(os.path.dirname(__file__))
MODELS = json.load(open(os.path.join(DIR_PATH, '../conf/models.json')))
WAS = json.load(open(os.path.join(DIR_PATH, '../conf/was.json')))
DATES = json.load(open(os.path.join(DIR_PATH, '../conf/dates.json')))
NETCDF_TEMPLATE = "{}/archive/{}{}.nc"

DEBUG = True
DELAY = DATES['delay']['delayed']
DELAY_DATE = DATES['delay']['start_date']


class WasTables:
    """ Class to manage the figure creation """

    def __init__(self, was=None, model='median', variable='SCONC_DUST', selected_date=None):
        """ Initialize WasTables with shapefile and netCDF data """
        self.model = model
        self.was = was
        self.days = range(1, 3)

        if self.was:
            self.was_name = WAS[self.was]['name_field']
            self.colormap = WAS[self.was]['colors']
            # read shapefile
            print(WAS[self.was]['shp'])
            self.was_name = WAS[self.was]['name_field']
            self.was_df = gpd.read_file(WAS[self.was]['shp'])

        if self.model and selected_date:
            # 
            current_time_before = (DELAY_DATE and datetime.strptime(selected_date, "%Y%m%d") < datetime.strptime(DELAY_DATE, "%Y%m%d"))
            if (not DELAY and not current_time_before) or (DELAY and current_time_before):
                self.days = range(0, 3)

            # read nc file
            if DEBUG: print("MODEL", model)
            filepath = NETCDF_TEMPLATE.format(
                MODELS[self.model]['path'],
                selected_date,
                MODELS[self.model]['template']
            )
            if not os.path.exists(filepath):
                print("Filepath %s doesn't exist. Exit." % filepath)
                sys.exit(0)

            self.input_file = nc_file(filepath)
            if 'lon' in self.input_file.variables:
                lon = self.input_file.variables['lon'][:]
                lat = self.input_file.variables['lat'][:]
            else:
                lon = self.input_file.variables['longitude'][:]
                lat = self.input_file.variables['latitude'][:]
            time_obj = self.input_file.variables['time']
            self.tim = time_obj[:]
            time_units = time_obj.units.split()
            if len(time_units) == 3:
                self.what, _, rdate = time_units
                rtime = "00:00"
            # elif len(time_units) >= 4:
            else:
                self.what, _, rdate, rtime = time_units[:4]
            if len(rtime) > 5:
                rtime = rtime[:5]

            if self.what == 'seconds':
                self.rdatetime = datetime.strptime("{} {}" .format(rdate,
                    rtime), "%Y-%m-%d %H:%M") + relativedelta(seconds=int(self.tim[0]))
            elif self.what == 'hours':
                self.rdatetime = datetime.strptime("{} {}" .format(rdate,
                    rtime), "%Y-%m-%d %H:%M") + relativedelta(hours=int(self.tim[0]))

            self.xlon, self.ylat = np.meshgrid(lon, lat)
            self.vardata = self.input_file.variables[variable][:]*1e9

        if selected_date:
            self.selected_date_plain = selected_date

            self.selected_date = datetime.strptime(
                selected_date, "%Y%m%d").strftime("%Y-%m-%d")

        opath = WAS[was]['path'].format(was=was, format='h5', date=self.selected_date_plain)
        os.makedirs(opath, exist_ok=True)
        otpl = WAS[was]['template'].format(date=self.selected_date_plain, var=variable, format='h5')
        self.h5_output = os.path.join(opath, otpl)

        gpath = WAS[was]['path'].format(was=was, format='geojson', date=self.selected_date_plain)
        os.makedirs(gpath, exist_ok=True)
        gtpl = WAS[was]['template'].format(date=self.selected_date_plain, var=variable, format='geojson')
        self.geojson_output = os.path.join(gpath, gtpl)


    def set_data(self, day=1):
        """ Set time dependent data """

        d_idx = []
        d_date = self.rdatetime + relativedelta(days=day)

        for n, tstep in enumerate(self.tim[:]):
            inc = (n == 0) and 0 or n*(self.tim[n] - self.tim[n-1])
            print('TIME', n, tstep)
            if self.what == 'seconds':
                ctime = (self.rdatetime + relativedelta(seconds=float(inc)))
            # elif self.what == 'hours':
            else:
                ctime = (self.rdatetime + relativedelta(hours=float(inc)))

            print('CTIME', ctime, 'DTIME', d_date)
            if ctime.strftime("%Y%m%d") == d_date.strftime("%Y%m%d"):
                d_idx.append(n)

        print('D_IDX', d_idx, 'WHAT', self.what)
        # print(self.vardata[d_idx, :, :])
        return self.vardata[d_idx,:,:].max(axis=0)

    def get_regions_data(self):
        """ """
        names = []
        colors = []
        definitions = []

        if not hasattr(self, 'xlon'):
            return names, colors, definitions

        flon = self.xlon.flatten()
        print('SHAPE LON', flon.shape)
        flat = self.ylat.flatten()
        print('SHAPE LAT', flat.shape)
        out_df = pd.DataFrame(
                columns=['day', 'names', 'colors', 'definitions'])

        days = []

        corresp = {}
        exclude = []
        if 'correspondence' in WAS[self.was]:
            corresp = WAS[self.was]['correspondence']
        if 'exclude' in WAS[self.was]:
            exclude = WAS[self.was]['exclude']

        for day in self.days:
            data = self.set_data(day=day).flatten()
            tmp_was_df = self.was_df
            if DEBUG: print('SHAPE DATA', data.shape)

            for idx in self.was_df.index.values:
                name = self.was_df.iloc[idx][self.was_name]
                if DEBUG: print("NAME", name)
                if name not in WAS[self.was]['values']:
                    if DEBUG: print("ERROR NAME", name, "MISSING! Exit.")
                    continue
                geom = self.was_df.iloc[idx]['geometry']
                mask = np.array([geom.contains(geometry.Point(x, y)) for x, y in np.array([flon, flat]).T])
                mdata = np.ma.masked_array(data, mask=~mask)
                # if DEBUG: print(";;;;", mdata[~mdata.mask])
                if not mdata[~mdata.mask].all():
                    reg_var = 0
                else:
                    reg_var = mdata[~mdata.mask].max()
                perc = WAS[self.was]['values'][name]
                # if DEBUG: print("::::", name, reg_var, perc)
                color = ((reg_var<perc[0]) and 'green') or \
                    (((reg_var>=perc[0]) and (reg_var<perc[1])) and 'gold') or \
                    (((reg_var>=perc[1]) and (reg_var<perc[2])) and 'darkorange') or 'red'
                names.append(name)
                colors.append(color)
                definitions.append(WAS[self.was]['colors'][color])
                days.append('Day{}'.format(day))

            if corresp:
                curr = dict(zip(names, colors))
                for newname in corresp:
                    newcolor = curr[corresp[newname]]
                    names.append(newname)
                    colors.append(newcolor)
                    definitions.append(WAS[self.was]['colors'][newcolor])
                    days.append('Day{}'.format(day))

            if exclude:
                for oldname in exclude:
                    e_idx = names.index(oldname)
                    names.pop(e_idx)
                    colors.pop(e_idx)
                    definitions.pop(e_idx)
                    days.pop(e_idx)
                    ex_idx = tmp_was_df.loc[self.was_df[self.was_name] == oldname].index[0]
                    tmp_was_df = tmp_was_df.drop(ex_idx)

            geojson_data = orjson.loads(tmp_was_df.to_json())

            acolors = np.array(colors)
            for color_idx, (color, definition) in enumerate(self.colormap.items()):
                if DEBUG: print('******', names, color, definition, colors)
                items_idxs = np.where(acolors == color)[0]
                if DEBUG: print('******', items_idxs)
                items_names = np.array(names)[items_idxs]
                for feat_idx in range(len(geojson_data['features'])):
                    props = geojson_data['features'][feat_idx]['properties']
                    if props[self.was_name] in items_names:
                        props['value'] = int(color_idx)
                        props['tooltip'] = "Region: <b>{}</b><br>Level: <b>{}</b>".format(
                                props[self.was_name], definition)

            tmp_out = self.geojson_output.replace('.geojson', '_{}.geojson'.format(day))
            with open(tmp_out, 'w') as out:
                out.write(orjson.dumps(geojson_data).decode('utf-8'))


        print(names, colors, definitions, days)
        out_df['names'] = names
        out_df['colors'] = colors
        out_df['definitions'] = definitions
        out_df['day'] = days
        out_df.set_index('day')
        print(self.h5_output)
        out_df.to_hdf(self.h5_output, 'was_{}'.format(self.selected_date_plain), format='table')


if __name__ == "__main__":
    curdate = sys.argv[1]
    if len(sys.argv) == 3:
        country = sys.argv[2]
        WasTables(was=country, selected_date=curdate).get_regions_data()
    else:
        for country in WAS:
            # country = 'burkinafaso'
            WasTables(was=country, selected_date=curdate).get_regions_data()
