# -*- coding: utf-8 -*-
""" Timeseries Handler """

import os
from datetime import datetime as dt
from datetime import timedelta
from dateutil.relativedelta import relativedelta
import plotly.graph_objs as go
import numpy as np
import pandas as pd
import json
import logging

from utils import concat_dataframes
from utils import retrieve_timeseries
from utils import retrieve_single_point

from map_handler import DASH_LOG_LEVEL
from map_handler import MODELS
from map_handler import MODEBAR_CONFIG
from map_handler import MODEBAR_CONFIG_TS
from map_handler import MODEBAR_LAYOUT
from map_handler import MODEBAR_LAYOUT_TS 
from map_handler import DELAY
from map_handler import DELAY_DATE
from map_handler import FREQ
from map_handler import VARS
from map_handler import OBS

DIR_PATH = os.path.dirname(os.path.realpath(__file__))


class ForecastModelsTimeSeriesHandler:
    """ Class to handle forecast time series """

    def __init__(self, model, date, var):
        """ Initialize ForecastModelsTimeSeriesHandler class

        Parameters
        ----------
        model : str
            Model name
        date : str
            Selected date
        var : str
            Variable name
        """

        if isinstance(model, str):
            model = [model]
        self.model = model
        self.var = var
        self.fpaths = []
        try:
            self.month = dt.strptime(date, "%Y%m%d").strftime("%Y%m")
            self.currdate = dt.strptime(date, "%Y%m%d").strftime("%Y%m%d")
        except:
            self.month = dt.strptime(date, "%Y-%m-%d").strftime("%Y%m")
            self.currdate = dt.strptime(date, "%Y-%m-%d").strftime("%Y%m%d")

        return None

    def retrieve_single_point(self, tstep, lat, lon, model=None):
        """ Retrive data on single point

        Parameters
        ----------
        tstep : int
            Timestep
        lat : float
            Latitude
        lon : float
            Longitude
        model : str, optional
            Model name, by default None

        Returns
        -------
        float
            Variable value at closest available point from (lat, lon)
        """

        if not model:
            model = self.model[0]

        logging.debug("---------- %s", model)

        method = 'netcdf'
        if (MODELS[model]['start'] == 12 and not DELAY and DELAY_DATE and (dt.strptime(self.currdate, "%Y%m%d") >= dt.strptime(DELAY_DATE, "%Y%m%d"))) or \
                (MODELS[model]['start'] == 12 and not DELAY and not DELAY_DATE):
            mod_date = (dt.strptime(self.currdate, "%Y%m%d") -
                    timedelta(days=1)).strftime("%Y%m%d")
        else:
            mod_date = self.currdate
        path_template = '{}{}.nc'.format(mod_date,
                                         MODELS[model]['template'],
                                         self.var)

        fpath = os.path.join(MODELS[model]['path'], method, path_template)

        return retrieve_single_point(fpath, tstep, lat, lon, self.var)

    def retrieve_timeseries(self, lat, lon, model=None, method='netcdf', forecast=False):
        """ Retrieve timeseries plot for point in modal window

        Parameters
        ----------
        lat : float
            Latitude
        lon : float
            Longitude
        model : str, optional
            Model name, by default None
        method : str, optional
            Method name, by default 'netcdf'
        forecast : bool, optional
            Indicates if we are using forecast data, by default False

        Returns
        -------
        go.Figure
            Modal window with timeseries
        """

        if not model:
            model = self.model

        # logging.debug("---------- %s", model)

        obs_eval = model[0] not in MODELS and model[0] in OBS
        if obs_eval:
            all_models = [model[0]] + list(MODELS.keys())
        else:
            all_models = list(MODELS.keys())

        for mod in all_models:
            if obs_eval:
                filedir = OBS[model[0]]['path']
                path_tpl = '{}-{}-{}_interp.ft'  # 202010-median-OD550_DUST_interp.ft
            else:   # if mod in MODELS:
                filedir = MODELS[mod]['path']
                path_tpl = '{}-{}-{}.ft'  # 202010-median-OD550_DUST_interp.ft

            if method == 'feather':
                path_template = path_tpl.format(self.month, mod, self.var)
            elif method == 'netcdf':
                path_template = '{}*{}.nc'.format(self.month, MODELS[mod]['template'], self.var)

            if forecast:
                method = 'netcdf'
                if (MODELS[mod]['start'] == 12 and not DELAY and DELAY_DATE and (dt.strptime(self.currdate, "%Y%m%d") >= dt.strptime(DELAY_DATE, "%Y%m%d"))) or \
                        (MODELS[mod]['start'] == 12 and not DELAY and not DELAY_DATE):
                    mod_date = (dt.strptime(self.currdate, "%Y%m%d") -
                            timedelta(days=1)).strftime("%Y%m%d")
                else:
                    mod_date = self.currdate
                path_template = '{}{}.nc'.format(mod_date, MODELS[mod]['template'], self.var)

            fpath = os.path.join(filedir,
                                 method,
                                 path_template)
            self.fpaths.append(fpath)

        title = "{} @ lat = {} and lon = {}".format(
            VARS[self.var]['name'], round(lat, 2), round(lon, 2)
        )

        mul = VARS[self.var]['mul']

        fig = go.Figure()

        for mod, fpath in zip(all_models, self.fpaths):
            # logging.debug(' %s %s', mod, fpath)
            if mod not in MODELS and mod in OBS:
                var = OBS[mod]['obs_var']
            else:
                var = self.var

            if not os.path.exists(fpath):
                logging.debug("NOT retrieving %s File doesn't exist.", fpath)
                continue

            logging.debug('Retrieving *** FPATH *** %s', fpath)
            try:
                ts_lat, ts_lon, ts_index, ts_values = retrieve_timeseries(
                        fpath, lat, lon, var, method=method, forecast=forecast)
            except Exception as e:
                logging.debug("NOT retrieving %s ERROR: %s", fpath, str(e))
                continue

            if forecast is True or mod in ('cams', 'ema-regcm4'):
                ts = pd.Series(index=ts_index, data=ts_values)
            else:
                date_index = pd.date_range('{}01'.format(self.month),
                        '{}01'.format((dt.strptime(self.month, '%Y%m') +
                            relativedelta(days=31)).strftime('%Y%m')), freq=f'{FREQ}H')
                ts = pd.Series(index=ts_index, data=ts_values).reindex(date_index)

            if isinstance(ts_lat, np.ndarray):
                ts_lat = float(ts_lat)
                ts_lon = float(ts_lon)
            if isinstance(ts.values, np.ndarray):
                ts_values = (ts.values*mul).round(2)
            else:
                ts_values = round((ts.values*mul), 2)

            if obs_eval and mod == model[0]:
                sc_mode = 'markers'
                marker = {'size': 12, 'symbol': "triangle-up-dot", 'color': '#f0b450'}
                line = {}
                visible = True
                name = "{}".format(mod.upper())
            elif obs_eval:
                sc_mode = 'lines'
                marker = {}
                line = { 'color': MODELS[mod]['color'] }
                if mod in model:
                    visible = True
                else:
                    visible = 'legendonly'
                name = "{}".format(MODELS[mod]['name'])
            else:
                sc_mode = 'lines'
                marker = {}
                line = { 'color': MODELS[mod]['color'] }
                if mod in model:
                    visible = True
                else:
                    visible = 'legendonly'
                name = "{} ({}, {})".format(
                        MODELS[mod]['name'], round(ts_lat, 2), round(ts_lon, 2))

            if mod == 'median':
                line['dash'] = 'dash'
            else:
                line['dash'] = 'solid'

            fig.add_trace(dict(
                    type='scatter',
                    name=name,
                    x=ts.index,
                    y=ts_values,
                    mode=sc_mode,
                    marker=marker,
                    line=line,
                    visible=visible,
                )
            )

        fig.update_layout(
            title=dict(text=title, x=0.45, y=.99),
            yaxis=dict(exponentformat="none"),
            # uirevision=True,
            autosize=True,
            showlegend=True,
            plot_bgcolor='#F9F9F9',
            font_size=12,
            # hovermode="closest",        # highlight closest point on hover
            hovermode="x",        # highlight closest point on hover
            margin={"r": 10, "t": 35, "l": 10, "b": 10},
        )

        return fig


class EvaluationGroundTimeSeriesHandler:

    def __init__(self, obs, start_date, end_date, var, models=None):
        """ Initialize EvaluationGroundTimeSeriesHandler class

        Parameters
        ----------
        obs : str
            Observations name
        start_date : str
            Start date
        end_date : str
            End date
        var : str
            Variable name
        models : str, optional
            Model name, by default None
        """

        self.obs = obs
        if models is None:
            models = list(MODELS.keys())
        self.model = models
        self.var = var
        self.dataframe = []
        logging.debug("ObsTimeSeries %s %s", start_date, end_date)
        self.date_range = pd.date_range(start_date, end_date, freq='D')

        fname_tpl = os.path.join(OBS[obs]['path'],
                                'feather',
                                '{{}}-{dat}-{{}}_interp.ft')

        months = np.unique([d.strftime("%Y%m") for d in self.date_range.to_pydatetime()])

        self.date_index = pd.date_range('{}01'.format(months[0]),
                '{}01'.format((dt.strptime(months[-1], '%Y%m') +
                    relativedelta(days=31)).strftime('%Y%m')), freq=f'{FREQ}H')

        logging.debug('MONTHS %s', months)
        logging.debug('DATE_INDEX %s', self.date_index)

        fname_obs = fname_tpl.format(dat=obs)
        notnans, obs_df = concat_dataframes(fname_obs, months, self.var,
            rename_from=OBS[obs]['obs_var'])
        self.dataframe.append(obs_df)

        if 'flt_var' in OBS[obs]:
            fname_flt = fname_tpl.format(dat='{}_{}'.format(obs, OBS[obs]['flt_var']))
            _, flt_df = concat_dataframes(fname_flt, months, self.var,
                rename_from=OBS[obs]['flt_var'], notnans=notnans)
            self.dataframe.append(flt_df)

        for mod in self.model:
            fname_mod = fname_tpl.format(dat=mod)
            _, mod_df = concat_dataframes(fname_mod, months, self.var,
                    rename_from=None, notnans=notnans)
            self.dataframe.append(mod_df)

        return None

    def retrieve_timeseries(self, idx, station_name, model):
        """ Retrieve timeseries plot for station in modal window

        Parameters
        ----------
        idx : int
            Station index
        station_name : str
            Station name
        model : str, optional
            Model name, by default None

        Returns
        -------
        go.Figure
            Modal window with timeseries
        """

        #FIXME pop filtering variable for now
        self.dataframe.pop(1)
        old_indexes = self.dataframe[0]['station'].unique()
        new_indexes = np.arange(old_indexes.size)
        dict_idx = dict(zip(new_indexes, old_indexes))
        logging.debug("RETRIEVE TS %s %s %s %s", idx, dict_idx[idx], station_name, model)
        fig = go.Figure()
        for mod, df in zip([self.obs]+self.model, self.dataframe):
            if df is None:
                continue
            logging.debug("MOD %s COLS %s", mod, df.columns)
            if df.columns[-1].upper() == self.var:
                df = df.rename(columns = { df.columns[-1]: self.var })

            logging.debug("BUILDING TIME-SERIES")
            try:
                tmp_df = df[df['station']==dict_idx[idx]].set_index('time')
                tmp_df.drop_duplicates(keep='first')
                timeseries = \
                   tmp_df.reindex(self.date_index)
                if timeseries[self.var].isnull().all():
                    continue
            except Exception as e:
                logging.debug("ERROR timeseries %s", str(e))
                continue

            logging.debug("SELECTING COORDS")
            if 'lat' in df.columns:
                lat_col = 'lat'
                lon_col = 'lon'
            else:
                lat_col = 'latitude'
                lon_col = 'longitude'

            if mod == self.obs:
                logging.debug("OBSERVATION")
                sc_mode = 'markers'
                marker = {'size': 10, 'symbol': "triangle-up-dot", 'color': '#f0b450'}
                line = {}
                visible = True
                name = mod.upper()
            else:
                logging.debug("MODEL %s", mod)
                try:
                    sc_mode = 'lines'
                    marker = {}
                    line = { 'color': MODELS[mod]['color'] }
                    visible = (mod == model) and True or 'legendonly'
                    name = "{}".format(MODELS[mod]['name'])
                    cur_lat = round(timeseries[lat_col].dropna()[0], 2)
                    cur_lon = round(timeseries[lon_col].dropna()[0], 2)
                except Exception as e:
                    logging.debug("ERROR MODEL %s %s %s", mod, str(e), timeseries[lat_col].dropna())
                    continue

            if mod == 'median':
                line['dash'] = 'dash'

            logging.debug("ADD TRACE")
            fig.add_trace(dict(
                type='scatter',
                name=name,
                x=timeseries.index,
                y=timeseries[self.var].round(2),
                mode=sc_mode,
                marker=marker,
                line=line,
                visible=visible,
                connectgaps=False
                )
            )

        title = "{} @ {} (lat = {:.2f}, lon = {:.2f})".format(
            VARS[self.var]['name'], station_name, cur_lat, cur_lon,
        )

        fig.update_layout(
            title=dict(text=title, x=0.45, y=0.99),
            # uirevision=True,
            autosize=True,
            showlegend=True,
            plot_bgcolor='#F9F9F9',
            font_size=12,
            hovermode="x",        # highlight closest point on hover
            margin={"r": 10, "t": 35, "l": 10, "b": 10},
        )
        fig.update_xaxes(
            range=[self.date_range[0], self.date_range[-1]],
            rangeslider_visible=True,
            rangeselector=dict(
                buttons=list([
                    dict(step="all", label="all"),
                    dict(count=14, label="2w",
                         step="day", stepmode="backward"),
                    dict(count=7, # label="1w",
                         step="day", stepmode="backward"),
                ])
            )
        )

        logging.debug('FIG TYPE %s', type(fig))

        return fig
