#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Oct 27 20:26:58 EDT 2023
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.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
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__':
    tt.check('end import')
#
#start from here
ncfile = __file__.replace('.py', '.nc')
if not os.path.exists(ncfile):
    #amip dat
    ifile = '/tigress/wenchang/data/cmip6/analysis/hydrological_sensitivity/wyplot_scatter_eta_amip.nc'
    ds_amip = xr.open_dataset(ifile)
    #cmip data
    ifile = '/tigress/wenchang/data/cmip6/analysis/hydrological_sensitivity/wyplot_scatter_eta_cmip.nc'
    ds_cmip = xr.open_dataset(ifile)
    #prediction
    pr_actual = ds_cmip['pr_4xco2_anom']
    gmst = ds_cmip['ts_4xco2_anom']
    tmsst = ds_cmip['tmsst_4xco2_anom']
    eta_amip = ds_amip['eta']
    eta_amip_sst = ds_amip['eta_sst']
    A = ds_amip['A']
    pr_predict_gmst = eta_amip * gmst + A
    pr_predict_gmst.attrs['units'] = '%'
    pr_predict_tmsst = eta_amip_sst * tmsst + A
    pr_predict_tmsst.attrs['units'] = '%'
    #wrap into a dataset
    ds = xr.Dataset(dict(
        pr_actual=pr_actual,
        pr_predict_gmst=pr_predict_gmst,
        pr_predict_tmsst=pr_predict_tmsst,
        gmst=gmst,
        tmsst=tmsst,
        A=A,
        eta_amip=eta_amip,
        eta_amip_sst=eta_amip_sst
        ))
    ds.to_netcdf(ncfile)
    print('[saved]:', ncfile)
else:
    ds = xr.open_dataset(ncfile)
    print('[loaded]:', ncfile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #lines plot
    fig,axes = plt.subplots(3, 3, figsize=(12, 8), sharex=True, sharey=True, dpi=100)
    abcs = list('abcdefghi')
    for ii,mm in enumerate(ds.model_member.values):
        ax = axes.flat[ii] 
        model = mm.split('_')[0]
        #mm = ds.model_member.values[0]
        ds_ = ds.sel(model_member=mm)
        ds_['pr_actual'].plot(color='k', label='actual Precip', ax=ax)
        ds_['pr_predict_gmst'].plot(color='C0', label='$\eta^{amip}\cdot\Delta T_s^{GMST} + A$', ax=ax)
        ds_['pr_predict_tmsst'].plot(color='C1', label='$\eta_{sst}^{amip}\cdot\Delta T_s^{TMSST} + A$', ax=ax)
        ax.legend(loc='lower right')
        eta_amip = ds_.eta_amip.item()
        eta_amip_sst = ds_.eta_amip_sst.item()
        A = ds_.A.item()
        s = f'({abcs[ii]}) {model}\n $\eta^{{amip}}={eta_amip:.2g}$%/K\n $\eta_{{sst}}^{{amip}}={eta_amip_sst:.2g}$%/K\n A={A:.2g}%'
        ax.text(0.01, 0.99, s, transform=ax.transAxes, va='top', ha='left')
        ax.set_title('')
        ax.set_xlabel('')
        ax.set_ylabel('')
        ax.set_xlim(0,150)
    for ax in axes[:,0]: ax.set_ylabel('$\Delta$P [%]')
    for ax in axes[-1,:]: ax.set_xlabel('year')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'pdf' in sys.argv: figname = figname.replace('.png', '.pdf')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #scatter plot
    fig,axes = plt.subplots(3, 3, figsize=(9, 8), sharex=True, sharey=True, dpi=100)
    abcs = list('abcdefghi')
    for ii,mm in enumerate(ds.model_member.values):
        ax = axes.flat[ii] 
        model = mm.split('_')[0]
        #mm = ds.model_member.values[0]
        ds_ = ds.sel(model_member=mm)
        ds_.plot.scatter(x='pr_actual', y='pr_predict_gmst', 
            label='$\eta^{amip}\cdot\Delta T_s^{GMST} + A$', c='none', edgecolor='C0', alpha=0.3, ax=ax)
        ds_.plot.scatter(x='pr_actual', y='pr_predict_tmsst', 
            label='$\eta_{sst}^{amip}\cdot\Delta T_s^{TMSST} + A$', c='none', edgecolor='C1', alpha=0.3, ax=ax)
        ax.axline((0,0), slope=1, color='gray', ls='--')
        ax.legend(loc='lower right')
        eta_amip = ds_.eta_amip.item()
        eta_amip_sst = ds_.eta_amip_sst.item()
        A = ds_.A.item()
        s = f'({abcs[ii]}) {model}\n $\eta^{{amip}}={eta_amip:.2g}$%/K\n $\eta_{{sst}}^{{amip}}={eta_amip_sst:.2g}$%/K\n A={A:.2g}%'
        ax.text(0.01, 0.99, s, transform=ax.transAxes, va='top', ha='left')
        ax.set_title('')
        ax.set_xlabel('')
        ax.set_ylabel('')
    for ax in axes[:,0]: ax.set_ylabel('predicted precip change [%]')
    for ax in axes[-1,:]: ax.set_xlabel('actual precip change [%]')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__scatter.png')
        if 'pdf' in sys.argv: figname = figname.replace('.png', '.pdf')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()