#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Mar  3 13:21:20 EST 2023
#v2: sensitivity parameter from AM2.5 CTL1860_florSST
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
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'AM2.5'
case = '1pct4XCO2'
b_gmst_am2p1 = 2.63 #2.628766826842829 # 2.6 # % per K
b_fco2_am2p1 = -2.40 #-2.403719774284724 #-2.4 # % per 2xCO2 forcing
#AM2.5 values: see /tigress/wenchang/analysis/misc/AM2.5/wyshow_dprecip.log
b_gmst_am2p5 = 2.98 # % per K
b_fco2_am2p5 = -2.27 # % per 2xCO2 forcing
#HIRAM values: see /tigress/wenchang/analysis/misc/HIRAM/wyshow_dprecip.log
b_gmst_hiram = 2.97 # % per K
b_fco2_hiram = -2.32 # % per 2xCO2 forcing
#AM2.5C360 values: see /tigress/wenchang/analysis/misc/AM2.5/wyshow_dprecip.log
b_gmst_am2p5c360 = 3.29 # % per K
b_fco2_am2p5c360 = -2.23 # % per 2xCO2 forcing
#select the params
#b_gmst = b_gmst_am2p5
b_gmst = 3.28 #% per K
b_fco2 = b_fco2_am2p5

ifile = glob.glob('precip_AM2.5_CTL1860_florSST2001_tigercpu_intelmpi_18_540PE_0101-????_glbmean.nc')[0]
da = xr.open_dataarray(ifile)
da_ctl = da
da_ref = da.mean('time')

ifile = glob.glob('precip_AM2.5_CTL1860_florSST2001_1pct4xCO2_tigercpu_intelmpi_18_540PE_0101-????_glbmean.nc')[0]
da = xr.open_dataarray(ifile)
da_precip = da
 
zz = np.linspace(0, 2, 140*12+1)[1:da.time.size+1] # 140 years of monthly CO2 forcing, 1 at 2xCO2 and 2 at 4xCO2
#zz = np.hstack((zz, zz[-1::-1]))
da = xr.DataArray(zz, dims=['time',], coords=[da.time,]).assign_attrs(units='2xCO2', long_name='CO2 forcing')
da_fco2 = da
#da.plot(); plt.show(); sys.exit()
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    func_lp = lambda x: x.rolling(year=9, center=True, min_periods=1).mean('year')
    func_pct = lambda x: (x/da_ref*100-100).assign_attrs(units='%')
    fig, ax = plt.subplots()
    alpha = 0.3

    da_ = da_precip.groupby('time.year').mean('time')
    da_.pipe(func_pct).plot(color='k', alpha=alpha)
    da_.pipe(func_pct).pipe(func_lp).plot(label=f'{model} precip ({case})', color='k')
    ax.set_prop_cycle(None)
    y = da_.pipe(func_pct)

    da_ = da_fco2.groupby('time.year').mean('time')*b_fco2
    #da_.plot(color='C0', alpha=alpha)
    #da_.pipe(func_lp).plot(color='C0', label=f'modeled by 2xCO2({b_fco2:.1f}%/2xCO2)')
    da_.plot(color='C0', label=f'modeled by 2xCO2({b_fco2:.1f}%/2xCO2)')
    
    ax.legend()
    ax.set_ylabel('global mean precip change [%]')
    #ax.set_xticks(range(1000, 1300, 40))
    #ax.axvline(1140, color='gray', ls='--')
    #offset 100 years to start with year 1 
    xticks = np.arange(100, da_.year.max().item(), 10)
    xticklabels = [f'{x-100:g}' for x in xticks]
    ax.set_xticks(xticks)
    ax.set_xticklabels(xticklabels)
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_{model}_{case}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    """
    import seaborn as sns
    fig,ax = plt.subplots()
    da_res = y - y_hat_co2
    dgmst = da_ts.groupby('time.year').mean('time') - da_ref_ts
    reg = da
    sns.scatterplot(x=dgmst, y=da_res, color='.1')
    #sns.histplot(x=dgmst, y=da_res)
    sns.kdeplot(x=dgmst, y=da_res, fill=True)
    ax.axline((0,0), slope=b_gmst, color='C1')

    import xlinregress
    fig,ax = plt.subplots()
    dprecip = y
    dgmst = da_ts.groupby('time.year').mean('time') - da_ref_ts
    #dprecip = dprecip.isel(year=dgmst<1.35)
    #dgmst = dgmst.isel(year=dgmst<1.35)
    reg = dprecip.linregress.on(dgmst)
    slope, y0 = reg.slope, reg.intercept
    sns.scatterplot(x=dgmst, y=dprecip, color='.1')
    #sns.histplot(x=dgmst, y=dprecip)
    #sns.kdeplot(x=dgmst, y=dprecip, fill=True)
    ax.axline((0,y0), slope=slope, color='C1')
    """

    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()