# -*- coding: utf-8 -*-
"""
Created on Fri Oct 13 17:59:14 2023

@author: cenv1021

This script is used to compare ISIMIP3b and ML modeling results, including map, boxplot, and 2D kdeplot

"""

import os
import numpy as np
import pandas as pd
import glob
import xarray as xr
import seaborn as sns
import matplotlib.pyplot as plt
from shapely import Polygon
import geopandas as gpd
import multiprocessing as mp
import sys,re
import matplotlib as mpl
import matplotlib.ticker as ticker
from matplotlib.lines import Line2D
import scienceplots
import cartopy.feature as cf
import cartopy.crs as ccrs
import cartopy
import warnings
warnings.simplefilter('ignore')
import matplotlib.font_manager as font_manager
font_manager.fontManager.addfont(os.environ['DATA']+'/fonts/Helvetica/Helvetica.ttf')

plt.style.use(['science','nature','no-latex']) # require install SciencePlots
plt.rcParams.update({"font.size":12, 'font.family':'Helvetica'}) 

def add_histogram(ax, vals, extent, vmin, vmax, vind, cmap):
    ax2 = ax.inset_axes(extent)
    cm = plt.cm.get_cmap(cmap)
    bins = [-np.inf] + np.arange(vmin, vmax+vind, vind).tolist() + [np.inf]
    n, bins, patches = ax2.hist(vals, bins=bins)
    bins[0] = vmin - vind
    bins[-1] = vmax + vind
    vcenter=(vmin+vmax)/2
    norm = mpl.colors.TwoSlopeNorm(vmin=vmin, vmax=vmax, vcenter=vcenter)
    bin_center = 0.5 * (bins[:-1] + bins[1:])
    for c, p in zip(bin_center, patches):
        plt.setp(p, 'facecolor', cm(norm(c)))
    ax2.get_yaxis().set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.tick_params(top=False, which='both', labelsize  = 8)
    ax2.xaxis.label.set_visible(False)
    ax2.xaxis.set_minor_locator(ticker.NullLocator())
    ax2.set_xticks([vmin, vcenter, vmax])
    ax2.spines['left'].set_visible(False)
    ax2.spines['top'].set_visible(False)
    ax2.spines['right'].set_visible(False)
    ax2.axvline(x=vcenter, color='k')

def plot_map(ax, lons, lats, vals, vmin, vmax, cmap, title, label, norm = None, vind = 1):
    ax.set_global()
    ax.set_ylim([-6525154.6651, 8625154.6651]) 
    ax.set_xlim([-12662826, 15924484]) 
    ax.spines['geo'].set_linewidth(0)
    ax.coastlines(linewidth = .2, color = '#707070')
    ax.add_feature(cf.BORDERS, linewidth = .2, color = '#707070')
    # plot scatter setting
    if norm is None:
        norm = mpl.colors.Normalize(vmin = vmin, vmax = vmax)
    # add colorbar
    ras = ax.scatter(lons, lats, c = vals, norm = norm, cmap = cmap, s = 2, transform = ccrs.PlateCarree(), zorder = 3, linewidths = 0)
    ax2 = ax.inset_axes([.35, .03, 0.25, .03])
    fig.colorbar(ras, cax = ax2, orientation = 'horizontal', extend = 'both')
    ax2.tick_params(labelsize = 9)
    ax2.set_title(label, size = 9, pad = 2)
    # add histograms
    add_histogram(ax, vals, [.62, .2, .15, .2], vmin, vmax, vind, cmap)
    # set plot title
    ax.set_title(title, size = 10, pad = 4)
    return ax

########################################################################################
# read machine learning results
fnames = glob.glob('../results/*/train_test_*_log_500km2_addLonLat_2001-2014.csv')
df_ml = []
for fname in fnames:
    df_ml0 = pd.read_csv(fname)
    df_ml0 = df_ml0[['reach_id','climate','RL10','pred','ohdb_catchment_area','AnnSum','Pvol','Rx1day','Rx5day','rsds','sfcWind','tasmax','tasmin']]
    model = os.path.basename(fname).split('_')[2]
    df_ml0['model'] = model
    df_ml.append(df_ml0)
df_ml = pd.concat(df_ml).dropna()
df_ml['RL10_mmday'] = df_ml['RL10'] / df_ml['ohdb_catchment_area'] * 86400 / 1000
df_ml['pred_mmday'] = df_ml['pred'] / df_ml['ohdb_catchment_area'] * 86400 / 1000
reach_ids = df_ml.reach_id.unique()
df_ml = df_ml.rename(columns = lambda x:x.lower()).drop(columns=['rl10','pred'])

# read climate observation
fnames = glob.glob('../data/GRIT_catch_mean_1981-2014_obs_climate_MSWEP_CRU_TerraClimate_*.csv')
def func(fname):
    df0 = pd.read_csv(fname, index_col = [0]).set_index(['year','name']).T
    df0 = df0.loc[df0.index.astype(int).isin(reach_ids),:]
    # calculate mean of climate observation in 2001-2014
    df0 = df0.T.reset_index()
    df0 = df0.loc[(df0.year>=2001)&(df0.year<=2014),:].drop(columns=['year'])
    df0 = df0.groupby('name').mean().T.reset_index().rename(columns={'index':'reach_id'})
    df0['reach_id'] = df0.reach_id.astype(int)
    print(fname)
    return(df0)
pool = mp.Pool(16)
df_obs = pool.map(func, fnames)
df_obs = pd.concat(df_obs)

df_obs = df_obs.rename(columns={'tmn':'tasmin','tmx':'tasmax','ws':'sfcwind','srad':'rsds'}).rename(columns=lambda x:x.lower())
df_obs = df_obs[['reach_id', 'annsum', 'pvol', 'rx1day', 'rx5day', 'rsds', 'sfcwind', 'tasmax', 'tasmin']]

# merge df_obs with df_ml
df = df_ml.merge(df_obs, on = 'reach_id', suffixes = ['_ml','_obs'])

# read ISIMIP3b climate input (5 GCMs)
fnames = glob.glob('../data/ISIMP3b/GRIT_catch_mean_ISIMIP3b_annual_indices_*_1981-2014.csv')
def func(fname):
    try:
        df0 = pd.read_csv(fname).set_index(['year','name']).T
        df0 = df0.loc[df0.index.astype(int).isin(reach_ids),:]
        # calculate mean of climate observation in 2001-2014
        df0 = df0.T.reset_index()
        df0 = df0.loc[(df0.year>=2001)&(df0.year<=2014),:].drop(columns=['year'])
        df0 = df0.groupby('name').mean().T.reset_index().rename(columns={'index':'reach_id'})
        df0['reach_id'] = df0.reach_id.astype(int)
        df0['model'] = os.path.basename(fname).split('_')[6].upper()
        df0 = df0[['reach_id', 'AnnSum', 'Pvol', 'Rx1day', 'Rx5day', 'rsds', 'sfcwind', 'tasmax', 'tasmin','model']]
        print(fname)
        return (df0)
    except:
        df0 = pd.DataFrame(columns=['reach_id', 'AnnSum', 'Pvol', 'Rx1day', 'Rx5day', 'rsds', 'sfcwind', 'tasmax', 'tasmin','model'])
        print(fname)
        return (df0)
pool = mp.Pool(48)
df_hm_input = pool.map(func, fnames)
df_hm_input = pd.concat(df_hm_input)
df_hm_input = df_hm_input.rename(columns=lambda x:x.lower())

# merge df_hm_input with df
df = df.merge(df_hm_input, on = ['reach_id','model'], how = 'left')
df = df.rename(columns={a:a+'_hm' for a in df.columns[-8:]})

# Kelvin tem to Celcius tem
df[['tasmin_ml','tasmax_ml']] = df[['tasmin_ml','tasmax_ml']] - 273.15

# calculate bias
for name in ['annsum', 'pvol', 'rx1day', 'rx5day', 'rsds', 'sfcwind', 'tasmax', 'tasmin']:
    df[name+'_hm_bias'] = (df[name+'_hm'] - df[name+'_obs']) / df[name+'_obs'] * 100
    df[name+'_ml_bias'] = (df[name+'_ml'] - df[name+'_obs']) / df[name+'_obs'] * 100
df['rl10_ml_bias'] = (df['pred_mmday'] - df['rl10_mmday']) / df['rl10_mmday'] * 100

# read ISIMP3b 10-year flood
fnames = glob.glob('../data/ISIMP3b/*2001-2014_annMax-RL-emp.csv')
df_hm = []
for fname in fnames:
    df0 = pd.read_csv(fname)
    cm = os.path.basename(fname).split('_')[1]
    hm = os.path.basename(fname).split('_')[0]
    df0['model'] = cm.upper()
    df0['hm'] = hm
    df_hm.append(df0)
df_hm = pd.concat(df_hm).dropna()

# connect reach_id with ISIMIP3b grid
grid = df_hm[['lon','lat']].drop_duplicates().reset_index().drop(columns={'index'})
grid['ID'] = np.arange(grid.shape[0])
geo = pd.Series([Polygon(((a-.25, b+.25), (a+.25, b+.25), (a+.25, b-.25), (a-.25, b-.25), (a-.25, b+.25))) for a,b in zip(grid.lon.values, grid.lat.values)])
gdf_hm = gpd.GeoDataFrame(data = grid, geometry = geo, crs="EPSG:4326")

points = pd.read_csv(glob.glob('../results/*/train_test_*_log_500km2_addLonLat_2001-2014.csv')[0])[['reach_id','lon','lat']]
gdf_ml = gpd.GeoDataFrame(points, geometry = gpd.points_from_xy(points.lon, points.lat), crs="EPSG:4326")
connect = gpd.sjoin(gdf_ml, gdf_hm, lsuffix = 'ml', rsuffix = 'hm', predicate = 'within')[['lon_hm','lat_hm','reach_id']].drop_duplicates()
connect = connect.rename(columns={'lon_hm':'lon','lat_hm':'lat'})

# calculate average bias and flood observation within each grid
df = df.merge(connect, on = 'reach_id').drop(columns=['reach_id'])
colNames = [s for s in df.columns if s not in ['lon','lat','model','ohdb_catchment_area','climate']]+['climate']

# choice 1: catchment-area-weighted average
# df_grid = df.groupby(['lon','lat','model']).apply(
#     lambda x: pd.Series((x.drop(columns=['climate','ohdb_catchment_area']).mul(x['ohdb_catchment_area'],axis=0).sum(axis=0)/x.ohdb_catchment_area.sum()).tolist()+[x.climate.mode().iloc[0]], 
#                         index = colNames),
# )

# choice 2: simple average
df_grid = df.groupby(['lon','lat','model']).apply(
    lambda x: pd.Series(x.drop(columns=['climate','ohdb_catchment_area']).mean().tolist()+[x.climate.mode().iloc[0]], 
                        index = colNames),
)
df_grid = df_grid.reset_index()

# transform ISIMIP3 flood into mmday using drainagea are calculated from flow direction dataset (https://files.isimip.org/ISIMIP3b/InputData/geo_conditions/river_routing/ddm30_flowdir_cru_neva.nc)
ds_darea = xr.open_dataset('../data/ISIMP3b/isimip3b_darea.nc')
df_darea = ds_darea.darea.to_dataframe().reset_index().dropna()
df_hm = df_hm.merge(df_darea, on = ['lon','lat'])
df_hm['rl10_mmday'] = df_hm['RL10'] / df_hm['darea'] * 86400 / 1000
df_hm = df_hm.drop(columns=['RL10','darea'])
df_hm = df_hm.groupby(['lon','lat','model'])['rl10_mmday'].mean().reset_index()

# merge df_grid with df_hm
df_grid = df_grid.merge(df_hm, on = ['lon','lat','model'], how = 'left')

# calculate ratio of absolute relative bias
df_grid = df_grid.rename(columns={'rl10_mmday_x':'rl10_mmday','rl10_mmday_y':'rl10_mmday_hm'})
df_grid['rl10_hm_bias'] = (df_grid['rl10_mmday_hm'] - df_grid['rl10_mmday']) / df_grid['rl10_mmday'] * 100
df_grid.to_csv('../results/compare_RF_ISIMIP3b_grid.csv')

# calculate ensemble-mean ratio of absolute relative bias for subplot 1 (map)
df1 = df_grid.groupby(['lon','lat'])[['rl10_hm_bias','rl10_ml_bias']].mean().reset_index()
df1['bias_diff'] = df1.rl10_ml_bias.abs() / df1.rl10_hm_bias.abs()

better_frc = df1.loc[df1.bias_diff<=1,:].shape[0] / df1.shape[0] * 100
bias_ml = df1.rl10_ml_bias.abs().agg(lambda x: [x.quantile(.025).round(2), x.mean().round(2), x.quantile(.975).round(2)])
bias_hm = df1.rl10_hm_bias.abs().agg(lambda x: [x.quantile(.025).round(2), x.mean().round(2), x.quantile(.975).round(2)])
print(better_frc, bias_ml, bias_hm)

# calculate relative bias for 5 GCMs for subplot 2 (boxplot)
df2 = df_grid.dropna()[['rl10_hm_bias','rl10_ml_bias','model']].melt(id_vars='model')
df2['value'] = df2['value'].abs()
df2.loc[df2.variable=='rl10_hm_bias','variable'] = 'HM'
df2.loc[df2.variable=='rl10_ml_bias','variable'] = 'ML'

# calculate ensemble-mean relative bias of Rx1day and RL10 for HM and ML, respectively, for subplot 3 (scatterplot)
df3 = df_grid.groupby(['lon','lat'])[['rl10_hm_bias','rl10_ml_bias','rx1day_hm_bias','rx1day_ml_bias']].mean()
df3 = df3.abs().reset_index()

print('finish calculation')

########################################################################################
# plot

# map for the ratio of absolute relative bias (RF / HM)
fig, ax = plt.subplots(figsize=(8, 6), subplot_kw={'projection':ccrs.EqualEarth()})
# df1 = df_hm_all.groupby(['lon','lat'])['RL10'].mean().reset_index()
lons = df1.lon.values
lats = df1.lat.values
vals = df1['bias_diff'].values
vmin, vmax = 0, 2
cmap = plt.cm.RdBu_r
norm = mpl.colors.TwoSlopeNorm(vmin = vmin, vmax = vmax, vcenter = 1)
title = None
label = 'Ratio of absolute relative bias'
plot_map(ax, lons, lats, vals, vmin, vmax, cmap, title, label, norm = norm, vind = .2)

# boxplot for each model
ax2 = ax.inset_axes([0.07, -.8, .42, .6])
palette = ['#E15D1E','#156AE6']
g = sns.boxplot(data = df2, x = 'model', y = 'value', hue = 'variable', ax = ax2, showfliers = False, width = .5, palette = palette, capprops = dict(linewidth=0))
sns.move_legend(g, loc = 'lower center', bbox_to_anchor = (.5, 1.01), ncol = 2)
g.legend_.set_title(None)
ax2.set_ylabel('Absolute relative bias (%)', fontsize = 8)
ax2.xaxis.label.set_visible(False)
ax2.xaxis.set_minor_locator(ticker.NullLocator())
ax2.tick_params(axis='x', rotation = 45)

# # scatterplot between Rx1day bias and flood bias
# ax = ax2.inset_axes([1.25, 0, 1, 1])
# ax.scatter(df3.rx1day_hm_bias.abs(), df3.rl10_hm_bias.abs(), label = 'HM', s = 1)
# ax.scatter(df3.rx1day_ml_bias.abs(), df3.rl10_ml_bias.abs(), label = 'ML', s = 1)
# ax.legend()
# ax.set_yscale('log')
# ax.set_xlim(None, 200)
# ax.set_xlabel('Relative bias of Rx1day (%)', fontsize = 8)
# ax.set_ylabel('Relative bias of 10yr flood magnitude (%)', fontsize = 8)

# two-dimensional KDE plot between Rx1day bias and flood bias
ax = ax2.inset_axes([1.25, 0, .96, 1])
ax_upper = ax.inset_axes([0, 1.01, 1, .1])
ax_right = ax.inset_axes([1.01, 0, .1, 1])
df3_ = pd.concat([
    df3[['rx1day_ml_bias', 'rl10_ml_bias']].rename(columns = lambda x:re.sub('_ml','',x)).assign(model='ML'),
    df3[['rx1day_hm_bias', 'rl10_hm_bias']].rename(columns = lambda x:re.sub('_hm','',x)).assign(model='HM'),
]).reset_index()
df3_.loc[df3_.rx1day_bias<=0.1,'rx1day_bias'] = 0.1
palette = ['#156AE6','#E15D1E']
g = sns.scatterplot(data = df3_, x="rx1day_bias", y="rl10_bias", hue="model", ax = ax, palette = palette, alpha = .2)
ax.set_ylim(0.01, None)
ax.set_yscale('log')
ax_right.set_yscale('log')
# create twin ax to add kdeplot (may not a good idea!)
ax_twin = ax.twinx()
ax_twin.axis('off')
ax_twin.set_yscale('log')
ax_twin.set_xlim(ax.get_xlim())
ax_twin.set_ylim(ax.get_ylim())
sns.kdeplot(data=df3_, x="rx1day_bias", y="rl10_bias", hue="model", ax = ax_twin, levels = 5, thres = .1, palette = palette)
sns.boxplot(data = df3_, x = 'rx1day_bias', ax=ax_upper, hue = 'model', legend = False, palette = palette, showfliers = False, capprops = dict(linewidth=0), width = 1)
sns.boxplot(data = df3_, y = 'rl10_bias', ax=ax_right, hue = 'model', legend = False, palette = palette, showfliers = False, capprops = dict(linewidth=0), width = 1)
ax.set_xlabel('Absolute relative bias of Rx1day (%)', fontsize = 8)
ax.set_ylabel('Absolute relative bias of 10yr flood magnitude (%)', fontsize = 8)
ax_right.set_ylim(ax.get_ylim())
ax_upper.set_xlim(ax.get_xlim())
ax_upper.xaxis.label.set_visible(False)
ax_right.yaxis.label.set_visible(False)
for side in ['top','right','bottom','left']:
    ax_right.spines[side].set_visible(False)
    ax_upper.spines[side].set_visible(False)
ax_right.tick_params(axis='both',which='both',labelleft=False,bottom=False,left=False,right=False,top=False)
ax_upper.tick_params(axis='both',which='both',labelbottom=False,bottom=False,left=False,right=False,top=False)

# revise legend
sns.move_legend(ax_twin, loc = 'upper right', title = None)
sns.move_legend(ax, loc = 'lower right', title = None, labels = ['Raw Rx1day', 'Bias-corrected Rx1day'])
for lh in g.legend_.legendHandles: 
    lh.set_alpha(.7)
    lh.set_markersize(5)

fig.savefig('../picture/fig2.png', dpi = 600)

# test significance of different Rx1day
from scipy import stats
df3_all = df_grid.copy()
df3_all = df3_all[['model','rl10_ml_bias','rl10_hm_bias','rx1day_ml_bias','rx1day_hm_bias']].dropna()
for model in np.unique(df3_all.model.values):
    df3 = df3_all.loc[df3_all.model==model,:]
    print(model, 'RL10', stats.ttest_ind(df3.rl10_ml_bias.values, df3.rl10_hm_bias.values))
    print(model, 'Rx1day', stats.ttest_ind(df3.rx1day_ml_bias.values, df3.rx1day_hm_bias.values))