#!/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 glob import glob


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 save_to_timeseries(df, fname, dest, fmt='feather'):
    df = df.to_dataframe()
    dest = os.path.join(dest, fname)
    if fmt == 'parquet':
        print('saving to parquet ...')
        df.to_parquet(
            '{}.parquet.gzip'.format(dest),
            compression='gzip')
    elif fmt == 'feather':
        print('reset index ...')
        var_df = df.astype(DTYPE).reset_index()
        print('saving to feather ...')
        var_df.to_feather(
            '{}.ft'.format(dest),
            )


def preprocess(ds, n=8):
    """keep only the first N timestep for each file"""
    return ds.isel(time=range(n))

def dropna_and_setindex(df, output_string, grp=None, column_name=[]):
    '''This is a helper function for ret_scores to do as title says for the groupbys'''
    if grp is None:
        stat = df.dropna(subset=[column_name]).set_index('time')
    else:
        stat = df.get_group(grp).dropna(subset=[column_name]).set_index('time')
    # print(output_string, stat.head(), stat.size, 'TIME', stat.columns, stat.index)
    return stat

def verify_df_size(df, 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'):
    decs = 2
    obs_stat = dropna_and_setindex(obs_df, 'OBS_STAT', grp, column_name=OBS[obs]['obs_var'])
    if obs_stat.size == 0:
        return verify_df_size(obs_stat, grpname)
    mod_stat = dropna_and_setindex(mod_df, 'MOD_STAT', grp, column_name=OBS[obs]['mod_var'])
    tot = obs_stat.merge(mod_stat, on=['time','station'])
    print("TOTAL DONE" )
    mvar, ovar = OBS[obs]['mod_var'], OBS[obs]['obs_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=None):
    """ Convert data from daily netCDF to time series or interpolated observations timeseries """

    station_dict = None

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

    # paths of all models
    mods_path = [os.path.join(MODELS[mod]['path'], 'netcdf', '{}.nc'.format(MODELS[mod]['template'])) for mod in model]

    if obs:
        obs_path = os.path.join(OBS[obs]['path'], 'netcdf', '{}.nc'.format(OBS[obs]['template']))

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

    # paths of all models per all months
    mod_paths = [["{}/{}*{}".format(os.path.dirname(mod_path), month, os.path.basename(mod_path)) for month in months] for mod_path in mods_path]

    obs_paths = [obs_path.format(OBS[obs]['obs_var'], month)
            if obs else ''
            for month in months]

    if 'flt_var' in OBS[obs]:
        flt_paths = [obs_path.format(OBS[obs]['flt_var'], month)
                if obs else ''
                for month in months]
        print("FLT", flt_paths)

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

    columns = ['station'] + list(MODELS.keys())
    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')
#         sites = pd.read_table(os.path.join(CURRENT_PATH, '../conf/',
#             OBS[obs]['sites']), delimiter=r"\s+",
#             engine='python').sort_values(by='AREA')

    models_ds = []
    for mpath in mod_paths:
        fnames = [x for mp in mpath for x in sorted(glob(mp))]
        try:
            #print(fnames)
            mod_ds = xr.open_mfdataset(fnames,
                                   concat_dim='time',
                                   combine='nested',
                                   join='override',
                                   preprocess=preprocess)
            models_ds.append(mod_ds)
        except Exception as err:
            print('Error', str(fnames), str(err))
            models_ds.append(None)
            continue

    try:
        # read filter variable
        flt_ds = None
        if 'flt_var' in OBS[obs]:
            flt_ds = xr.open_mfdataset(flt_paths,
                                   concat_dim='time',
                                   combine='nested',
                                   ) #.astype('float32')
            # select stations and remove others
            flt_ds = flt_ds.where(flt_ds['station_name'].isin(sites['SITE'].values.astype(np.bytes_)))
            flt_df = flt_ds.to_dataframe().reset_index().dropna(subset=['station_name'])
            flt_df = flt_df[['time', 'station', 'ndata', 'station_name', OBS[obs]['flt_var']]]

#             # re-numbering stations
#             old_stations = flt_df['station'].unique()
#             new_stations = np.arange(old_stations.size)
#             for old_st, new_st in zip(old_stations, new_stations):
#                 flt_df.loc[flt_df['station'] == old_st, 'station'] = new_st

        # read observations and transform to a dataframe with only needed stations
        obs_ds = xr.open_mfdataset(obs_paths,
                               concat_dim='time',
                               combine='nested',
                               ) #.astype('float32')
        if 'latitude' in obs_ds.variables:
            obs_lon_name = 'longitude'
            obs_lat_name = 'latitude'
        else:
            obs_lon_name = 'lon'
            obs_lat_name = 'lat'
        if 'sites' in OBS[obs]:
            print("Station selection")
            # select only stations in a given text file and remove others
            obs_ds = obs_ds.where(obs_ds['station_name'].isin(sites['SITE'].values.astype(np.bytes_)))
            obs_df = obs_ds.to_dataframe().reset_index().dropna(subset=['station_name'])
            obs_df = obs_df[['time', 'station', 'ndata', 'station_name', OBS[obs]['obs_var']]]

#             print("Station re-numbering")
#             # re-numbering stations
#             for old_st, new_st in zip(old_stations, new_stations):
#                 obs_df.loc[obs_df['station'] == old_st, 'station'] = new_st

            print("Station adding areas")
            # adding area to each station
            for site in sites['SITE']:
                obs_df.loc[obs_df['station_name'] == site.encode('utf-8'), 'area'] = sites.loc[sites['SITE'] == site, 'AREA'].values

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

            obs_df = obs_ds.to_dataframe().reset_index()

        # apply filtering on the whole dataframe
        if 'flt_var' in OBS[obs]:
            print("OBS DF")
            # print(obs_df.head())
            print("FLT DF")
            # print(flt_df.head())
            obs_df.loc[flt_df[OBS[obs]['flt_var']]>0.6, OBS[obs]['obs_var']] = np.nan

            # group by area
            obs_grps_area = obs_df.sort_values(['area']).groupby('area')

        else:
            obs_grps_area = [] 

    except Exception as err:
        print('Error', str(obs_paths), 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", obs_grps_area.groups.keys())
    for mod_idx, mod_ds in enumerate(models_ds):
        stations = []
        bias = []
        rmse = []
        corr = []
        frge = []
        totn = []
        curr_mod = list(model)[mod_idx] 
        print('IDX', mod_idx, curr_mod)
        if mod_ds 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

        for var in mod_ds.variables:
            if variable != var.upper():
                continue
            curvar = var
            mod_da = mod_ds[curvar].where((mod_ds[curvar] > 0) & (mod_ds[curvar] < 10.))
            if obs_ds:
                print('interpolating to observations ...')
                if 'longitude' in mod_ds.variables:
                    mod_interp = mod_da.interp(longitude=obs_lon, latitude=obs_lat)
                else:
                    mod_interp = mod_da.interp(lon=obs_lon, lat=obs_lat)
                print('converting to df  ...')
            else:
                mod_interp = mod_da
                print('converting dataset to df ...')

            mod_df = mod_interp.to_dataframe().reset_index()
            if OBS[obs]['mod_var'].lower() in mod_df.columns:
                mod_df.rename(columns={OBS[obs]['mod_var'].lower(): OBS[obs]['mod_var']}, inplace=True)

            if '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():
                    mod_df.loc[mod_df['station'] == st, 'area'] = obs_df.loc[obs_df['station'] == st, 'area'].values[0]

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

            print("MOD_GRPS", mod_grps_area.groups.keys())
            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')
                    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 mod_ds:
            mod_ds.close()

    if obs_ds:
        obs_ds.close()
    if flt_ds:
        flt_ds.close()

    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__":
    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)
