#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jul  9 11:03:16 EDT 2024
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
#compare eta_amip vs. eta_cmip (for both cases of gmst and tmsst)
#FLOR results
ifile = 'fig_bars_eta__131-150.nc'
ds = xr.open_dataset(ifile) # will use eta_cmip and eta_cmip_tmsst in the dataset
eta_amip = 3.28 + xr.zeros_like(ds.eta_cmip)
eta_amip_sst = 3.52 + xr.zeros_like(ds.eta_cmip_tmsst)
ds['eta_amip'] = eta_amip
ds['eta_amip_sst'] = eta_amip_sst
ds_flor = ds
#cmip6 results
ifile = 'Fig05.nc' #will use pr_actual, gmst, tmsst, A, eta_amip and eta_amip_sst
ds = xr.open_dataset(ifile)
yearslice = slice(131,150) #last 20 years of standard 150-year output from CMIP6 4xCO2 experiment
eta_cmip = (ds.pr_actual.sel(year=yearslice).mean('year') - ds.A)/ds.gmst.sel(year=yearslice).mean('year')
eta_cmip_tmsst = (ds.pr_actual.sel(year=yearslice).mean('year') - ds.A)/ds.tmsst.sel(year=yearslice).mean('year')
ds['eta_cmip'] = eta_cmip
ds['eta_cmip_tmsst'] = eta_cmip_tmsst
ds_cmip6 = ds


 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots(figsize=(5,5))
    ds_flor.plot.scatter(x='eta_amip', y='eta_cmip', color='C0', facecolor='none', marker='^',  label='FLOR $\eta$')
    #for ii in range(ds_flor.label.size):
    #    ax.text(ds_flor.eta_amip.isel(label=ii), ds_flor.eta_cmip.isel(label=ii), '12345678'[ii], color='C0', ha='center', va='center')
    ds_flor.plot.scatter(x='eta_amip_sst', y='eta_cmip_tmsst', color='C1', facecolor='none', marker='^', label='FLOR $\eta^{TMSST}$')
    #for ii in range(ds_flor.label.size):
    #    ax.text(ds_flor.eta_amip_sst.isel(label=ii), ds_flor.eta_cmip_tmsst.isel(label=ii), '12345678'[ii], color='C1', ha='center', va='center')

    ds_cmip6.plot.scatter(x='eta_amip', y='eta_cmip', color='C0', facecolor='none', label='CMIP6 4xCO2 $\eta$')
    ds_cmip6.plot.scatter(x='eta_amip_sst', y='eta_cmip_tmsst', color='C1', facecolor='none', label='CMIP6 4xCO2 $\eta^{TMSST}$')

    ax.set_xlim(2,4)
    ax.set_ylim(2,4)
    ax.axline((3,3), slope=1, color='gray', ls='--')
    ax.set_aspect('equal')
    ax.set_xlabel('$\eta$ from AGCM uniform warming [%/K]')
    ax.set_ylabel('$\eta$ from AOGCM experiments [%/K]')
    ax.legend()
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #kde plot
    import seaborn as sns
    fig,ax = plt.subplots()
    sns.kdeplot(ds_cmip6.eta_cmip, color='C0', ls='-', label='4xCO2 $\eta$')
    sns.kdeplot(ds_cmip6.eta_amip, color='C0', ls='--', label='AMIP $+$4K $\eta$')
    sns.kdeplot(ds_cmip6.eta_cmip_tmsst, color='C1', ls='-', label='4xCO2 $\eta^{TMSST}$')
    sns.kdeplot(ds_cmip6.eta_amip_sst, color='C1', ls='--', label='AMIP $+$4K $\eta^{TMSST}$')

    ax.legend(loc='upper left', ncol=1)
    ax.set_xlabel('$\eta$ from CMIP6 [%/K]')
    ax.set_ylabel('PDF')
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__kde.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()