#!/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 = 'FLOR'
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 = 'precip_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_2001-2100_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl = da
da_ref = da.mean('time')

ifile = 'precip_FLOR_CTL1860_newdiag_1pct4xCO2_tigercpu_intelmpi_18_576PE_2001-2140_glbmean.nc'
da = xr.open_dataarray(ifile)
da_precip = da

#Ts
ifile = glob.glob('t_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_????-????_glbmean.nc')[0]
print(ifile)
da = xr.open_dataarray(ifile).sel(time=slice('2001','2100'))
da_ctl_ts = da
da_ref_ts = da.mean('time')

ifile = 't_surf_FLOR_CTL1860_newdiag_1pct4xCO2_tigercpu_intelmpi_18_576PE_2001-2140_glbmean.nc'
da_ts = xr.open_dataarray(ifile)
 
zz = np.linspace(0, 2, 140*12+1)[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_ts.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_ts.groupby('time.year').mean('time') - da_ref_ts)*b_gmst + 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 gmst({b_gmst:.1f}%/K) and co2({b_fco2:.1f}%/2xCO2)')
    y_hat = da_
    
    da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_gmst
    da_.plot(color='C1', alpha=alpha)
    da_.pipe(func_lp).plot(color='C1', label='gmst contribution')
    
    da_ = -da_fco2.groupby('time.year').mean('time')*b_fco2
    da_.plot(color='C2', label='co2 contribution (mirrored)', ls='--')
    y_hat_co2 = da_

    #da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_gmst_am2p5 + da_fco2.groupby('time.year').mean('time')*b_fco2_am2p5
    #da_.pipe(func_lp).plot(label=f'modeled if AM2.5 coeff: gmst({b_gmst_am2p5}%/K) and co2({b_fco2_am2p5}%/2xCO2)')

    da_ = y - y_hat
    da_.plot(color='C3', alpha=alpha)
    da_.pipe(func_lp).plot(color='C3', label='residual')
    
    ax.legend(loc='upper left')
    ax.set_ylabel('global mean precip change [%]')
    #ax.set_xticks(range(1000, 1300, 40))
    #ax.axvline(1140, color='gray', ls='--')
    
    #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()