#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2021 Earth Sciences Department, BSC-CNS

""" Save time series """

import numpy as np
import pandas as pd
from sklearn.metrics import mean_squared_error
import json
import os
import sys
from datetime import datetime
from dateutil.relativedelta import relativedelta
from data_handler import ObsTimeSeriesHandler

CURRENT_PATH = os.path.abspath(os.path.dirname(__file__))
VARS = json.load(open(os.path.join(CURRENT_PATH, '../conf/vars.json')))
MODELS = json.load(open(os.path.join(CURRENT_PATH, '../conf/models.json')))
OBS = json.load(open(os.path.join(CURRENT_PATH, '../conf/obs.json')))
DTYPE = 'float32'


def dataframe_filter(df, column_name):
    """ Help to filter data: TODO read valid ranges from conf """
    df[column_name] = df[column_name].where(df[column_name].between(0, 10))
    return df.dropna(subset=[column_name]).set_index('time')

def dropna_and_setindex(df, grp=None, column_name=[]):
    """ This is a helper function for ret_scores to do as title says for the groupbys """
    if grp is not None:
        df = df.get_group(grp)
    return dataframe_filter(df, column_name)

def verify_df_size(grpname=None):
    """ Helper function for ret_scores to verify that dfs aren't empty """
    if grpname is None:
        return 'Total', {'BIAS':  '-', 'CORR': '-', 'RMSE': '-', 'FRGE': '-', 'TOTN': '0'}

    return None, {'BIAS':  '-', 'CORR': '-', 'RMSE': '-', 'FRGE': '-', 'TOTN': '0'}

def ret_scores(mod_df, obs_df, grp=None, grpname=None, obs='aeronet'):
    """ Function that calculates skill scores """
    decs = 2
    obs_stat = dropna_and_setindex(obs_df, grp, column_name=OBS[obs]['mod_var'])
    if obs_stat.size == 0:
        return verify_df_size(grpname)
    mod_stat = dropna_and_setindex(mod_df, grp, column_name=OBS[obs]['mod_var'])
    tot = obs_stat.merge(mod_stat, on=['time','station'], suffixes=('_obs', '_mod'))
    print("TOTAL DONE", tot)
    mvar, ovar = '{}_mod'.format(OBS[obs]['mod_var']), '{}_obs'.format(OBS[obs]['mod_var'])
    # print('******************** TOT MOD', tot[mvar])
    # print('******************** TOT OBS', tot[ovar])
    try:
        bias = (tot[mvar] - tot[ovar]).mean().round(decimals=decs)
        bias = not np.isnan(bias) and bias.astype(str) or '-'
    except:
        bias = '-'
    try:
        corr = tot[[ovar, mvar]].corr().values[0][1].round(decimals=decs)
        corr = not np.isnan(corr) and corr.astype(str) or '-'
    except:
        corr = '-'
    try:
        rmse = mean_squared_error(tot[ovar][~tot[ovar].isna()], tot[mvar][~tot[ovar].isna()],
                                  squared=False).round(decimals=decs)
        rmse = not np.isnan(rmse) and rmse.astype(str) or '-'
    except:
        rmse = '-'
    try:
        frge = (2*((tot[mvar] - tot[ovar])/(tot[mvar] + \
                                            tot[ovar]))).abs().mean().round(decimals=decs)
        frge = not np.isnan(frge) and frge.astype(str) or '-'
    except:
        frge = '-'
    try:
        totn = str(tot[ovar][tot[ovar].notnull()].shape[0])
    except:
        totn = '0'
    if grpname is None:
        ret_name = 'Total'
    else:
        ret_name = obs_stat[grpname][0]
        # print("-------------------------", ret_name, type(ret_name))
        if isinstance(ret_name, bytes):
            ret_name = ret_name.decode('utf-8')
        # print("-------------------------", ret_name, type(ret_name))
    return ret_name, {'BIAS':  bias, 'CORR': corr, 'RMSE': rmse, 'FRGE': frge, 'TOTN': totn}


def convert2timeseries(model, obs=None, months=[]):
    """ Convert data from daily netCDF to time series or interpolated
    observations timeseries """

    station_dict = {}

    if isinstance(model, str):
        model = [model]

    if not months:
        curr_year = str(datetime.now().year)
        months = np.arange(int("{}01".format(curr_year)),
                int("{}12".format(curr_year)))

    models = list(MODELS.keys())
    variable = OBS[obs]['mod_var']

    first_day = "{}01".format(months[0])
    last_month_fd = "{}01".format(months[-1])
    last_day = (datetime.strptime(last_month_fd, "%Y%m%d") + relativedelta(months=1) - relativedelta(days=1)).strftime("%Y%m%d")

    print("**********", months, first_day, last_day, "************")
    timeseries_obj = ObsTimeSeriesHandler(obs, first_day, last_day, variable, models=model)
    dataframes = timeseries_obj.dataframe
    print(dataframes)

    columns = ['station'] + models
    print(columns)
    bias_df = pd.DataFrame(columns=columns)
    rmse_df = pd.DataFrame(columns=columns)
    corr_df = pd.DataFrame(columns=columns)
    frge_df = pd.DataFrame(columns=columns)
    totn_df = pd.DataFrame(columns=columns)

    # read sites for a text file
    if 'sites' in OBS[obs]:
        sites = pd.read_csv(os.path.join(CURRENT_PATH, '../conf/',
            OBS[obs]['sites'])).sort_values(by='AREA')
    else:
        sites = None

    obs_df = None
    flt_df = None
    obs_grps_area = pd.DataFrame({}) 
    mod_grps_area = pd.DataFrame({})
    first_mod_df = pd.DataFrame({})
    try:
        if obs:
            obs_df = dataframes[0]
            # read filter variable
            if sites is not None:
                if 'flt_var' in OBS[obs]:
                    flt_df = dataframes[1]
                    print(dataframes[1])
                    # flt_df = flt_df.reset_index().dropna(subset=['station_name'])
                    flt_df = flt_df[['time', 'station', OBS[obs]['mod_var']]]  # .reset_index().dropna()

                print("Station selection")
                # select only stations in a given text file and remove others
                # obs_df = obs_df.reset_index().dropna(subset=['station_name'])
                obs_df = obs_df[['time', 'station', OBS[obs]['mod_var']]]

                print("Station adding areas")
                # adding area to each station
                for tmp_mod in dataframes[2:]:
                    if tmp_mod is not None:
                        first_mod_df = dataframes[2]
                        break 
                if 'longitude' in first_mod_df.columns:
                    lon_name = 'longitude'
                    lat_name = 'latitude'
                else:
                    lon_name = 'lon'
                    lat_name = 'lat'
                for sidx in sites.index:
                    site = sites.iloc[sidx]
                    print(site)
                    stat_num = first_mod_df.loc[(first_mod_df[lon_name] ==
                                                 site['LONGITUDE']) &
                                                (first_mod_df[lat_name] ==
                                                 site['LATITUDE']),
                                                'station'].values
                    if stat_num.size == 0:
                        continue
                    stat_num = stat_num[0]
                    obs_df.loc[obs_df['station'] == stat_num, 'station_name'] = site['SITE']
                    obs_df.loc[obs_df['station'] == stat_num, 'area'] = site['AREA']

                station_dict = dict(obs_df[['station', 'station_name']].drop_duplicates().values)
                print("STATION DICT", station_dict)
                obs_df = obs_df  # .reset_index().dropna(subset=['station_name'])
                # print(obs_ds.head())
                print("Lat/Lon retrieval")
                # retrieve lat and lon only with 'station' dimension to perform the interpolation
            else:
                obs_df = obs_df.reset_index()

            # apply filtering on the whole dataframe
            if flt_df is not None:
                print("OBS DF")
                print(obs_df.size, obs_df.head())
                print("FLT DF")
                print(flt_df.size, flt_df.head())
                # obs_df.loc[flt_df[OBS[obs]['flt_var']]>0.6, OBS[obs]['obs_var']] = np.nan
                cond = OBS[obs]['flt_expr']
                mod_var = OBS[obs]['mod_var']
                flt_df_var = flt_df[mod_var]
                obs_df = obs_df.loc[pd.eval("flt_df_var"+cond, target=obs_df)]  # = np.nan
                # group by area
                obs_grps_area = obs_df.sort_values(['area']).groupby('area')

    except Exception as err:
        print('Error', str(err))
        sys.exit(1)

    variable = OBS[obs]['mod_var']
    print('v', variable)
    # if observation perform interpolation
    print('mod var', OBS[obs]['mod_var'])
    print('obs var', OBS[obs]['obs_var'])

    #print(models_ds)
    print("OBS_GRPS", hasattr(obs_grps_area, 'groups') and obs_grps_area.groups.keys() or 'NO OBS GROUPS')
    for mod_idx, mod_df in enumerate(dataframes[2:]):
        stations = []
        bias = []
        rmse = []
        corr = []
        frge = []
        totn = []
        curr_mod = list(model)[mod_idx] 
        print('IDX', mod_idx, curr_mod)
        if mod_df is None:
            print('MOD is NONE!!! Continue')
            bias_df['station'] = ''
            rmse_df['station'] = ''
            corr_df['station'] = ''
            frge_df['station'] = ''
            totn_df['station'] = ''
            bias_df[curr_mod] = '-'
            rmse_df[curr_mod] = '-'
            corr_df[curr_mod] = '-'
            frge_df[curr_mod] = '-'
            totn_df[curr_mod] = 0
            continue

        var = mod_df.columns[-1]
        if variable != var.upper():
            continue
        if OBS[obs]['mod_var'].lower() == mod_df.columns[-1]:
            mod_df.rename(columns={OBS[obs]['mod_var'].lower(): OBS[obs]['mod_var']},
                          inplace=True)

        if obs_df is not None and 'sites' in OBS[obs]:
            # adding area to each station
            print("------------", mod_df['station'].unique(), obs_df['station'].unique(),
                  "------------")
            for st in mod_df['station'].unique():
                tmp_stations = obs_df.loc[obs_df['station'] == st, 'area'].values
                if tmp_stations.size == 0:
                    continue
                mod_df.loc[mod_df['station'] == st, 'area'] = tmp_stations[0]

            # print(mod_df)
            mod_grps_area = mod_df.sort_values(['area']).groupby('area')

        print("MOD_GRPS", hasattr(mod_grps_area, 'groups') and
              mod_grps_area.groups.keys() or 'NO MOD GROUPS')
        for obs_area in obs_grps_area:
            print("Running scores for AREA", obs_area[0])
            area, scores = ret_scores(mod_grps_area, obs_grps_area, obs_area[0], grpname='area')
            stations.append(area)
            bias.append(scores['BIAS'])
            rmse.append(scores['RMSE'])
            corr.append(scores['CORR'])
            frge.append(scores['FRGE'])
            totn.append(scores['TOTN'])
            mod_grps_stat = mod_grps_area.get_group(obs_area[0]).groupby('station')
            obs_grps_stat = obs_grps_area.get_group(obs_area[0]).groupby('station')
            for stat in obs_grps_stat:
                print("Running scores for station:", stat[0], station_dict[stat[0]])
                station, scores = ret_scores(mod_grps_stat, obs_grps_stat, stat[0],
                                             grpname='station_name')
                print("SCORES", scores)
                if station == None:
                    station = station_dict[stat[0]]  # .decode('utf-8')
                    print("NO DATA for:", station)
                stations.append(station)
                bias.append(scores['BIAS'])
                rmse.append(scores['RMSE'])
                corr.append(scores['CORR'])
                frge.append(scores['FRGE'])
                totn.append(scores['TOTN'])

        print("TOTAL")
        total, scores = ret_scores(mod_df, obs_df)
        print("DONE")
        stations.append(total)
        bias.append(scores['BIAS'])
        rmse.append(scores['RMSE'])
        corr.append(scores['CORR'])
        frge.append(scores['FRGE'])
        totn.append(scores['TOTN'])
        print(stations)
        print('CURR MODEL', curr_mod)
        print('BIAS', bias)
        print('RMSE', rmse)
        print('CORR', corr)
        print('FRGE', frge)
        print('TOTN', totn)

        bias_df['station'] = stations
        rmse_df['station'] = stations
        corr_df['station'] = stations
        frge_df['station'] = stations
        totn_df['station'] = stations
        bias_df[curr_mod] = bias
        rmse_df[curr_mod] = rmse
        corr_df[curr_mod] = corr
        frge_df[curr_mod] = frge
        totn_df[curr_mod] = totn


    if len(months) == 1:
        fname = "{}".format(months[0])
    else:
        fname = "{}-{}".format(months[0], months[-1])

    print("Writing h5 ...")
    bias_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_bias.h5")
    corr_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_corr.h5")
    rmse_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_rmse.h5")
    frge_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_frge.h5")
    totn_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_totn.h5")
    print("BIAS", bias_df_out.format(fname))
    bias_df.to_hdf(bias_df_out.format(fname), "bias_{}".format(fname), mode='w',
                   format='table', encoding='UTF-8')
    print("CORR", corr_df_out.format(fname))
    corr_df.to_hdf(corr_df_out.format(fname), "corr_{}".format(fname), mode='w',
                   format='table', encoding='UTF-8')
    print("RMSE", rmse_df_out.format(fname))
    rmse_df.to_hdf(rmse_df_out.format(fname), "rmse_{}".format(fname), mode='w',
                   format='table', encoding='UTF-8')
    print("FRGE", frge_df_out.format(fname))
    frge_df.to_hdf(frge_df_out.format(fname), "frge_{}".format(fname), mode='w',
                   format='table', encoding='UTF-8')
    print("TOTN", totn_df_out.format(fname))
    totn_df.to_hdf(totn_df_out.format(fname), "totn_{}".format(fname), mode='w',
                   format='table', encoding='UTF-8')


if __name__ == "__main__":
    sys.path.append(os.path.dirname(os.path.abspath('..')))
    if len(sys.argv) > 4:
        print("Error. Usage: {} [MODEL] [OBS] [MONTHS]".format(sys.argv[0]))
        sys.exit(1)
    if len(sys.argv) == 1:
        # for model in MODELS:
        convert2timeseries(MODELS.keys(), obs='aeronet')
    elif len(sys.argv) == 2:
        model = sys.argv[1].split(',')
        if model[0] == 'all':
            model = MODELS.keys()
        convert2timeseries(model=model, obs='aeronet')
    elif len(sys.argv) == 3:
        model = sys.argv[1].split(',')
        if model[0] == 'all':
            model = MODELS.keys()
        obs = sys.argv[2]
        convert2timeseries(model, obs=obs)
    else:
        model = sys.argv[1].split(',')
        if model[0] == 'all':
            model = MODELS.keys()
        obs = sys.argv[2]
        months = sys.argv[3].split(",")
        convert2timeseries(model, obs=obs, months=months)
