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

""" Save time series """

import xarray as xr
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
from data_handler import STATS


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'

# STATS = OrderedDict({ 'bias': 'BIAS', 'corr': 'CORR', 'rmse': 'RMSE', 'frge': 'FRGE', 'totn': 'CASES' })

def ret_scores(mod_df, obs_df, axis=None):
    decs = 2
    total_df = obs_df.merge(mod_df, on=('time', 'lat', 'lon'), suffixes=('_obs', '_mod'))
    mod_var = OBS[obs]['mod_var']
    if not '{}_mod'.format(mod_var) in total_df.columns:
        mod_var = mod_var.lower()
    mod_da = total_df[f'{mod_var}_mod'].values
    obs_da = total_df[f'{mod_var}_obs'].values

    total_grps = total_df.groupby(['lat', 'lon'])
    # res = total_grps[[f'{mod_vat}_obs', f'{mod_var}_mod']]
    # .corr()[f'{mod_var}_obs'].unstack()['f{mod_var}_mod'].reset_index()

    try:
        if axis is None:
            bias = np.nanmean(mod_da - obs_da, axis=axis).round(decimals=decs).flatten().astype(str)
        else:
            bias = total_grps.apply(lambda row: (row['OD550_DUST_mod'] - row['OD550_DUST_obs']).mean()).reset_index().values.T.round(decimals=decs).astype(str)
            bias = pd.DataFrame({
                    'lat': bias[0],
                    'lon': bias[1],
                    'bias': bias[2]
                    })
        if type(bias) in (list, np.ndarray) and len(bias) == 1:
            bias = bias[0]
    except Exception as err:
        print('BIAS ERROR', str(err))
        bias = '-'
    try:
        if axis is None:
            corr = np.ma.corrcoef(mod_da.flatten(), obs_da.flatten()).data[0,1].round(decimals=decs).astype(str)
        else:
            corr = total_grps[[f'{mod_var}_mod', f'{mod_var}_obs']].corr()[f'{mod_var}_obs'].unstack()[f'{mod_var}_mod'].reset_index().rename(columns={f'{mod_var}_mod':'corr'}).values.T.round(decimals=decs).astype(str)
            corr = pd.DataFrame({
                    'lat': corr[0],
                    'lon': corr[1],
                    'corr': corr[2]
                    })

        if type(corr) in (list, np.ndarray) and len(corr) == 1:
            corr = corr[0]
    except Exception as err:
        print('CORR ERROR', str(err))
        corr = '-'
    try:
        if axis is None:
            rmse = mean_squared_error(obs_da, mod_da, squared=False).round(decimals=decs).astype(str)
        else:
            # rmse = np.sqrt(np.square((obs_da - mod_da).mean(axis=axis))).values.flatten().astype(str)
            rmse = total_grps.apply(lambda row: np.sqrt(np.square((row[f'{mod_var}_obs'] - row[f'{mod_var}_mod']).mean()))).reset_index().rename(columns={0:'rmse'}).values.T.round(decimals=decs).astype(str)
            rmse = pd.DataFrame({
                    'lat': rmse[0],
                    'lon': rmse[1],
                    'rmse': rmse[2]
                    })
        if type(rmse) in (list, np.ndarray) and len(rmse) == 1:
            rmse = rmse[0]
    except Exception as err:
        print('RMSE ERROR', str(err))
        rmse = '-'
    try:
        if axis is None:
            frge = np.abs((2*((mod_da - obs_da)/(mod_da + \
                                            obs_da)))).mean().round(decimals=decs)
            frge = not np.isnan(frge) and frge.astype(str) or '-'
            print("FRGE ->", frge)
        # frge_arr = np.abs(2*((mod_da[~(mod_da.mask+obs_da.mask)] - obs_da[~(mod_da.mask+obs_da.mask)])))
        else:
            frge = total_grps.apply(lambda row: np.nanmean(np.abs(2*(row[f'{mod_var}_obs'] - row[f'{mod_var}_mod'])))).reset_index().rename(columns={0:'frge'}).values.T.round(decimals=2).astype(str)
            frge = pd.DataFrame({
                    'lat': frge[0],
                    'lon': frge[1],
                    'frge': frge[2]
                    })
#        if frge['frge'].size == 0:
#            frge = '-'
#        if type(frge) in (list, np.ndarray) and len(frge) == 1:
#            frge = frge[0]
    except Exception as err:
        print('FRGE ERROR', str(err))
        frge = '-'
    try:
        totn = str(total_df[f'{mod_var}_obs'][total_df[f'{mod_var}_obs'].notnull()].shape[0])
        if type(totn) in (list, np.ndarray) and len(totn) == 1:
            totn = totn[0]
    except:
        totn = '0'
    #if axis is None:
    return {'bias':  bias, 'corr': corr, 'rmse': rmse, 'frge': frge, 'totn': totn}


def convert2timeseries(model, obs=None, months=None):
    """ Convert data from daily netCDF to time series or interpolated observations timeseries """
    decs = 2
    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)))

    variable = OBS[obs]['mod_var']
    obs_var = OBS[obs]['obs_var']
    # if observation perform interpolation
    print('mod var', OBS[obs]['mod_var'])
    print('obs var', OBS[obs]['obs_var'])
    filedir = OBS[obs]['path']
    path_tpl = '{}-{}-{}_interp.ft'  # 202010-median-OD550_DUST_interp.ft

    obs_paths = [os.path.join(filedir, 'feather', path_tpl.format(month, obs, obs_var))
            if obs else ''
            for month in months]
    # paths of all models
    mod_paths = [[os.path.join(filedir, 'feather', path_tpl.format(month, mod, variable)) for month in months] for mod in model]

    print("MOD", mod_paths)
    print("OBS", obs_paths)

    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)
    total_columns = ['model'] + list(STATS.keys())
    print(total_columns)
    total_df = pd.DataFrame(columns=total_columns)

    stat_columns = ['lon', 'lat'] + list(MODELS.keys())
    print(stat_columns)
    bias_df = pd.DataFrame(columns=stat_columns)
    rmse_df = pd.DataFrame(columns=stat_columns)
    corr_df = pd.DataFrame(columns=stat_columns)
    frge_df = pd.DataFrame(columns=stat_columns)
    # totn_df = pd.DataFrame(columns=stat_columns)

    models_df = dataframes[1:]
    obs_df = dataframes[0]
    print('OBS', obs_df.columns)
    # retrieve lat and lon to perform the interpolation
    obs_lat = obs_df['lat'].unique()
    obs_lon = obs_df['lon'].unique()
    xlon, xlat = np.meshgrid(obs_lon, obs_lat)
    coords_df = pd.DataFrame({
        'lon': xlon.flatten().round(decimals=decs).astype(str),
        'lat': xlat.flatten().round(decimals=decs).astype(str)
        })

    print(model)
    for mod_idx, (mod_df, mod) in enumerate(zip(models_df, model)):
        print('MOD', mod, 'IDX', mod_idx)
        if mod_df is None:
            continue
        print('VARIABLE', variable, mod_df.columns)
        if variable.lower() in mod_df.columns:
            mod_df = mod_df.rename(columns={ variable.lower(): variable.upper() })

        # TOTAL SCORES
        scores = ret_scores(mod_df, obs_df)
        scores.update({ 'model': mod })
        total_df.loc[mod_idx] = scores
        print(mod)
        print(total_df)

        # SCORES per point
        print('SCORES per point!')
        scores = ret_scores(mod_df, obs_df, axis=0)
        #print('lon', obs_lon.shape, 'lat', obs_lat.shape, 'BIAS', scores['bias'].shape)
        bias_df['lon'] = coords_df['lon'] 
        bias_df['lat'] = coords_df['lat']
        corr_df['lon'] = coords_df['lon']
        corr_df['lat'] = coords_df['lat']
        rmse_df['lon'] = coords_df['lon']
        rmse_df['lat'] = coords_df['lat']
        frge_df['lon'] = coords_df['lon']
        frge_df['lat'] = coords_df['lat']
        bias_df[stat_columns[mod_idx+2]] = coords_df.merge(scores['bias'], on=('lat', 'lon'), how='left')['bias']
        corr_df[stat_columns[mod_idx+2]] = coords_df.merge(scores['corr'], on=('lat', 'lon'), how='left')['corr']
        rmse_df[stat_columns[mod_idx+2]] = coords_df.merge(scores['rmse'], on=('lat', 'lon'), how='left')['rmse']
        frge_df[stat_columns[mod_idx+2]] = coords_df.merge(scores['frge'], on=('lat', 'lon'), how='left')['frge']

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

    print("Writing h5 tables for maps ...")
    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")
    print("Writing h5 tables for scores ...")
    total_df_out = os.path.join(OBS[obs]['path'], "h5", "{}_scores.h5")

    print("BIAS", bias_df_out.format(fname), bias_df.size)
    bias_df = bias_df[((bias_df[bias_df.columns[2:]]!='nan') & (bias_df[bias_df.columns[2:]].notnull())).any(axis=1)]
    print("AFTER", bias_df.size)
    bias_df.to_hdf(bias_df_out.format(fname), "bias_{}".format(fname), format='table')
    
    print("CORR", corr_df_out.format(fname), corr_df.size)
    corr_df = corr_df[((corr_df[corr_df.columns[2:]]!='nan') & (corr_df[corr_df.columns[2:]].notnull())).any(axis=1)]
    print("AFTER", corr_df.size)
    corr_df.to_hdf(corr_df_out.format(fname), "corr_{}".format(fname), format='table')
    
    print("RMSE", rmse_df_out.format(fname), rmse_df.size)
    rmse_df = rmse_df[((rmse_df[rmse_df.columns[2:]]!='nan') & (rmse_df[rmse_df.columns[2:]].notnull())).any(axis=1)]
    print("AFTER", rmse_df.size)
    rmse_df.to_hdf(rmse_df_out.format(fname), "rmse_{}".format(fname), format='table')
    
    print("FRGE", frge_df_out.format(fname), frge_df.size)
    frge_df = frge_df[((frge_df[frge_df.columns[2:]]!='nan') & (frge_df[frge_df.columns[2:]].notnull())).any(axis=1)]
    print("AFTER", frge_df.size)
    frge_df.to_hdf(frge_df_out.format(fname), "frge_{}".format(fname), format='table')
    
    print("TOTAL", total_df_out.format(fname))
    print(mod, total_df)
    total_df.to_hdf(total_df_out.format(fname), "total_{}".format(fname), format='table')


if __name__ == "__main__":
    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)
