#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sat Dec 18 21:22:10 EST 2021
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    s = ' '
    tt = Timer(f'start {s.join(sys.argv)}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#from misc.landmask import flagland
import xlinregress
from scipy import stats
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
same_ens = False if len(sys.argv)>1 and 'all_ens' in sys.argv[1:] else True #whether to use all availale ens or use the same ens across experiments
paired_ttest = False if len(sys.argv)>1 and 'ind_ttest' in sys.argv[1:] else True #whether to use paired ttest

len_season = 3
def wyplot_hist(dsname, ax=None, alpha=0.5, bins=10, density=True, same_ens=True, paired_ttest=False):
    if ax is None: ax = plt.gca()
    #datasets
    ds_hist = xr.open_dataset(f'{dsname}.global.tcseason3.historical.noIO.nc')
    ds_nat = xr.open_dataset(f'{dsname}.global.tcseason3.hist-nat.noIO.nc')
    ds_ghg = xr.open_dataset(f'{dsname}.global.tcseason3.hist-GHG.noIO.nc')
    ds_era5 = xr.open_dataset(f'{dsname}.global.tcseason3.ERA5.wy.noIO.nc')
    ds_merra2 = xr.open_dataset(f'{dsname}.global.tcseason3.MERRA2.wy.noIO.nc')
    #global ds1, model_members
    #ds1 = ds_hist
    #subset of model_members to ensure the same across experiments
    if same_ens:
        model_members = [model_member for model_member in ds_nat.model_member.values
            if model_member in ds_ghg.model_member.values
            and model_member in ds_hist.model_member.values]
        ds_hist = ds_hist.sel(model_member=model_members)
        ds_nat = ds_nat.sel(model_member=model_members)
        ds_ghg = ds_ghg.sel(model_member=model_members)
        #save subset data
        ofile_hist = f'{dsname}.global.tcseason3.historical.{ds_hist.model_member.size}ens.noIO.nc'
        if not os.path.exists(ofile_hist):
            ds_hist.drop(['model', 'members']).to_netcdf(ofile_hist)
            print('[saved]:', ofile_hist)
        ofile_nat = f'{dsname}.global.tcseason3.hist-nat.{ds_nat.model_member.size}ens.noIO.nc'
        if not os.path.exists(ofile_nat):
            ds_nat.drop(['model', 'members']).to_netcdf(ofile_nat)
            print('[saved]:', ofile_nat)
        ofile_ghg = f'{dsname}.global.tcseason3.hist-GHG.{ds_ghg.model_member.size}ens.noIO.nc'
        if not os.path.exists(ofile_ghg):
            ds_ghg.drop(['model', 'members']).to_netcdf(ofile_ghg)
            print('[saved]:', ofile_ghg)
    #attrs
    units = ds_nat.timeseries.attrs['units']
    long_name = ds_nat.timeseries.attrs['long_name']
    trend_units = ds_nat.slope.attrs['units']

    #pdfs
    ds_hist.slope.plot.hist(ax=ax, density=density, alpha=alpha, color='gray')
    ds_nat.slope.plot.hist(ax=ax, density=density, alpha=alpha)
    ds_ghg.slope.plot.hist(ax=ax, density=density, alpha=alpha)

    #norm fit
    if dsname == 'PI' and 'MCM-UA-1-0' in ds_hist.model:#model MCM-UA-1-0 has problems for PI
        mu_hist, sigma_hist = stats.norm.fit(ds_hist.slope.dropna('model_member')) 
    else:
        mu_hist, sigma_hist = stats.norm.fit(ds_hist.slope)
    mu_nat, sigma_nat = stats.norm.fit(ds_nat.slope)
    mu_ghg, sigma_ghg = stats.norm.fit(ds_ghg.slope)
    xmin = min(ds_hist.slope.min(), ds_nat.slope.min(), ds_ghg.slope.min()).item()
    xmax = max(ds_hist.slope.max(), ds_nat.slope.max(), ds_ghg.slope.max()).item()
    xran = xmax - xmin
    x = np.linspace(xmin-xran*0.1, xmax+xran*0.1, 100)
    ax.plot(x, stats.norm.pdf(x, mu_hist, sigma_hist), color='gray', label='historical')
    ax.plot(x, stats.norm.pdf(x, mu_nat, sigma_nat), color='C0', label='hist-nat')
    ax.plot(x, stats.norm.pdf(x, mu_ghg, sigma_ghg), color='C1', label='hist-GHG')

    #ERA5/MERRA2
    """
    slope_obs = ( ds_era5.slope.item() + ds_merra2.slope.item() )/2
    #ax.axvline(slope_obs, color='k', label='ERA5/MERRA2 mean')
    ylim = ax.get_ylim()
    y = ylim[0] + (ylim[1] - ylim[0])/40
    #ax.scatter(slope_obs, y, marker='*', s=50, color='k', label='ERA5/MERRA2 mean', zorder=10)
    #ax.text(slope_obs, y, f'{slope_obs:.1E}({trend_units})', va='bottom', ha='right')
    ax.scatter(slope_obs, y, marker='*', s=50, color='k', label='Obs.', zorder=10)
    ax.text(slope_obs, y*2, f'{slope_obs:.1E}', va='bottom', ha='center')
    """
    #ERA5
    slope_obs = ds_era5.slope.item()
    #ax.axvline(slope_obs, color='k', label='ERA5/MERRA2 mean')
    ylim = ax.get_ylim()
    y = ylim[0] + (ylim[1] - ylim[0])/40
    #ax.scatter(slope_obs, y, marker='*', s=50, color='k', label='ERA5/MERRA2 mean', zorder=10)
    #ax.text(slope_obs, y, f'{slope_obs:.1E}({trend_units})', va='bottom', ha='right')
    ax.scatter(slope_obs, y, marker='*', s=50, color='k', label='wyERA5', zorder=10)
    ax.text(slope_obs, y*2, f'{slope_obs:.1E}', va='bottom', ha='left', color='gray', rotation=30)
    #MERRA2
    slope_obs = ds_merra2.slope.item()
    #ax.axvline(slope_obs, color='k', label='ERA5/MERRA2 mean')
    ylim = ax.get_ylim()
    y = ylim[0] + (ylim[1] - ylim[0])/40
    #ax.scatter(slope_obs, y, marker='*', s=50, color='k', label='ERA5/MERRA2 mean', zorder=10)
    #ax.text(slope_obs, y, f'{slope_obs:.1E}({trend_units})', va='bottom', ha='right')
    ax.scatter(slope_obs, y, marker='*', s=50, color='k', facecolor='none', label='wyMERRA2', zorder=10)
    ha = 'right' if dsname=='Vshear' else 'left'
    rotation = -30 if dsname=='Vshear' else 30
    ax.text(slope_obs, y*2, f'{slope_obs:.1E}', va='bottom', ha=ha, color='gray', rotation=rotation)

    #ttest_ind, kstest
    if same_ens and paired_ttest:
        ttest_nat = stats.ttest_rel(ds_nat.slope, ds_hist.slope)
        ttest_ghg = stats.ttest_rel(ds_ghg.slope, ds_hist.slope)
    else:
        ttest_nat = stats.ttest_ind(ds_nat.slope, ds_hist.slope)
        ttest_ghg = stats.ttest_ind(ds_ghg.slope, ds_hist.slope)
    kstest_nat = stats.kstest(ds_nat.slope, ds_hist.slope)
    kstest_ghg = stats.kstest(ds_ghg.slope, ds_hist.slope)
    #ttest = stats.ttest_ind(ds_nat.slope, ds_ghg.slope)
    #kstest = stats.kstest(ds_nat.slope, ds_ghg.slope)

    #
    if dsname == 'SST':
        ax.legend(ncol=2, loc='upper right')
    ax.set_title(f'{long_name}, nat(GHG) vs. historical t-test p={ttest_nat.pvalue:.1E}({ttest_ghg.pvalue:.1E}), KS-test p={kstest_nat.pvalue:.1E}(p={kstest_ghg.pvalue:.1E})', loc='left')
    ax.set_xlabel(trend_units)
    ax.set_ylabel('Probability Density')

def wyplot_series(dsname, ax=None, same_ens=True):
    if ax is None: ax = plt.gca()
    #datasets
    #global ds_hist, models, ds1, ds_nat, ds_ghg
    ds_hist = xr.open_dataset(f'{dsname}.global.tcseason3.historical.noIO.nc')
    ds_nat = xr.open_dataset(f'{dsname}.global.tcseason3.hist-nat.noIO.nc')
    ds_ghg = xr.open_dataset(f'{dsname}.global.tcseason3.hist-GHG.noIO.nc')
    ds_era5 = xr.open_dataset(f'{dsname}.global.tcseason3.ERA5.wy.noIO.nc')
    ds_merra2 = xr.open_dataset(f'{dsname}.global.tcseason3.MERRA2.wy.noIO.nc')
    #subset of model_members to ensure the same across experiments
    #ds1 = ds_hist
    if same_ens:
        model_members = [model_member for model_member in ds_nat.model_member.values
            if model_member in ds_ghg.model_member.values
            and model_member in ds_hist.model_member.values]
        ds_hist = ds_hist.sel(model_member=model_members)
        ds_nat = ds_nat.sel(model_member=model_members)
        ds_ghg = ds_ghg.sel(model_member=model_members)
        models = [model_member.split('_')[0] for model_member in ds_hist.model_member.values]
        models = xr.DataArray(models, dims='model_member', coords=[ds_hist.model_member.values])
        ds_hist = ds_hist.drop('model').assign_coords(model=models)
        ds_nat = ds_nat.drop('model').assign_coords(model=models)
        ds_ghg = ds_ghg.drop('model').assign_coords(model=models)
    #attrs
    units = ds_nat.timeseries.attrs['units']
    long_name = ds_nat.timeseries.attrs['long_name']
    trend_units = ds_nat.slope.attrs['units']

    #time series
    #historical
    da = ds_hist.timeseries.groupby(ds_hist.model).mean('model_member').mean('model')
    rg = da.linregress.on(da.year)
    #save regression results
    ofile = f'{dsname}.global.tcseason3.historical.wgtEnsMean.noIO.nc'
    if not os.path.exists(ofile):
        ds = rg
        ds['timeseries'] = da
        ds.to_netcdf(ofile)
        print('[saved]:', ofile)
    #da.plot(ax=ax, color='gray', marker='o', fillstyle='none', label='historical')
    da.plot(ax=ax, color='gray', marker='o', fillstyle='none', label=f'historical({rg.slope.item():.1E} {trend_units})')
    #rg.predicted.plot(ax=ax, color='gray', ls='--', label=f'historical trend ({rg.slope.item():.2g} {trend_units})')
    rg.predicted.plot(ax=ax, color='gray', ls='--')#, label=f'historical trend ({rg.slope.item():.2g} {trend_units})')
    #hist-nat
    da = ds_nat.timeseries.groupby(ds_nat.model).mean('model_member').mean('model')
    rg = da.linregress.on(da.year)
    #save regression results
    ofile = f'{dsname}.global.tcseason3.hist-nat.wgtEnsMean.noIO.nc'
    if not os.path.exists(ofile):
        ds = rg
        ds['timeseries'] = da
        ds.to_netcdf(ofile)
        print('[saved]:', ofile)
    #da.plot(ax=ax, color='C0', marker='o', fillstyle='none', label='hist-nat')
    da.plot(ax=ax, color='C0', marker='o', fillstyle='none', label=f'hist-nat({rg.slope.item():.1E} {trend_units})')
    #rg.predicted.plot(ax=ax, color='C0', ls='--', label=f'hist-nat trend ({rg.slope.item():.2g} {trend_units})')
    rg.predicted.plot(ax=ax, color='C0', ls='--')#, label=f'hist-nat trend ({rg.slope.item():.2g} {trend_units})')
    #hist-GHG
    da = ds_ghg.timeseries.groupby(ds_ghg.model).mean('model_member').mean('model')
    rg = da.linregress.on(da.year)
    #save regression results
    ofile = f'{dsname}.global.tcseason3.hist-GHG.wgtEnsMean.noIO.nc'
    if not os.path.exists(ofile):
        ds = rg
        ds['timeseries'] = da
        ds.to_netcdf(ofile)
        print('[saved]:', ofile)
    #da.plot(ax=ax, color='C1', marker='o', fillstyle='none', label='hist-GHG')
    da.plot(ax=ax, color='C1', marker='o', fillstyle='none', label=f'hist-GHG({rg.slope.item():.1E} {trend_units})')
    #rg.predicted.plot(ax=ax, color='C1', ls='--', label=f'hist-GHG trend ({rg.slope.item():.2g} {trend_units})')
    rg.predicted.plot(ax=ax, color='C1', ls='--')#, label=f'hist-GHG trend ({rg.slope.item():.2g} {trend_units})')
    ##obs
    #da = (ds_era5.timeseries + ds_merra2.timeseries)/2
    #da.plot(ax=ax, color='k', marker='o', fillstyle='none', label='ERA5/MERRA2 mean')
    #rg = da.linregress.on(da.year)
    #rg.predicted.plot(ax=ax, color='k', ls='--', label=f'ERA5/MERRA2 mean trend ({rg.slope.item():.2g} {trend_units})')

    ax.legend()
    ax.set_ylabel(units)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close('all')
    fig,axes = plt.subplots(4, 2, figsize=(10,10))
    dsnames = ('SST', 'PI', 'RH', 'Vshear')
    for ii,dsname in enumerate(dsnames):
        ax = axes[ii, 0]
        wyplot_hist(dsname, ax=ax, same_ens=same_ens, paired_ttest=paired_ttest)
        ax = axes[ii, 1]
        wyplot_series(dsname, ax=ax, same_ens=same_ens)
    for ax,label in zip(axes.flat, list('abcdefgh')):
        ax.text(0, 1,' '+label+')', fontsize='x-large', transform=ax.transAxes, ha='left', va='top')
    ax = axes[0, 1]
    ax.text(1, 1,' exclude 60-90E', transform=ax.transAxes, ha='right', va='bottom')
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_season{len_season}mon.png')
        if same_ens:
            figname = figname.replace('.png', '_sameEns.png')
            if paired_ttest:
                figname = figname.replace('.png', '_paired-ttest.png')
        else:
            figname = figname.replace('.png', '_allEns.png')
        wysavefig(figname)
    if len(sys.argv)>1 and 'savefigpdf' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_season{len_season}mon.pdf')
        if same_ens:
            figname = figname.replace('.pdf', '_sameEns.pdf')
            if paired_ttest:
                figname = figname.replace('.pdf', '_paired-ttest.pdf')
        else:
            figname = figname.replace('.pdf', '_allEns.pdf')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()