#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Mar  3 13:21:20 EST 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
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'FLOR'
cold_correction = [
True,
False,
][-1]
#case_default = 'co2x4'
#case = ['co2x4', 'co2x2', 'p2p0solar', 'p6p0solar', 'p1p0solar', 'm2p0solar', 'm6p0solar', 'co2xp5', 'co2xp25'][-1]
case = """
p2p0solar
p6p0solar
m2p0solar
m6p0solar
co2x4
co2xp5
co2xp25
co2x2
""".split()[-1]
ifiles = {
    'co2x2': '../precip_FLOR_co2x2_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc',
    'co2x4': '../precip_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc',
    'co2xp5': '../precip_FLOR_co2xp5_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc',
    'co2xp25': '../precip_FLOR_co2xp25_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc',
    'p1p0solar': '../precip_FLOR_p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc',
    'p2p0solar': '../precip_FLOR_p2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_glbmean.nc',
    'p6p0solar': '../precip_FLOR_p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0600_glbmean.nc',
    'm2p0solar': '../precip_FLOR_m2p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0400_glbmean.nc',
    'm6p0solar': '../precip_FLOR_m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE_0101-0597_glbmean.nc',
}
#nonlinear_ratio = True
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 = 2.7 #%/K
#b_tropmst = 3.01 #%/K
#if cold_correction and case in ('co2xp25', 'co2xp5',): b_tropmst = 3.24
b_tropmst = 3.24 #%/K
b_fco2 = -2.13 #% per CO2 forcing
#b_fco2 = -2.16 #% per 2xCO2
#b_solar = -0.83 # % per 2% solar
b_solar = -0.85 # % per 2% solar from AM2.5 ctl1860_florSST

#precip ctl
ifile = '../precip_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0001-2301_glbmean.nc'
da = xr.open_dataarray(ifile)
da_ctl = da
da_ref = da.sel(time=slice('0101', '0200')).mean('time')

#precip pert
ifile = ifiles[case]
"""
ifile = 'precip_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc'
ifile = ifile.replace(case_default, case).replace('0101-0250', '????-????')
ifile = glob.glob(ifile)[0]
"""
print(ifile)
da = xr.open_dataarray(ifile)
da_precip = da

#Ts ctl
ifile = 't_surf_FLOR_CTL1860_newdiag_tigercpu_intelmpi_18_576PE_0101-0200_tropmean.nc'
da = xr.open_dataarray(ifile)
da_ctl_ts = da
da_ref_ts = da.sel(time=slice('0101', '0200')).mean('time')

#Ts pert
ifile = ifiles[case].replace('../precip', 't_surf').replace('glbmean', 'tropmean')
"""
ifile = 't_surf_FLOR_co2x4_CTL1860_tigercpu_intelmpi_18_576PE_0101-0250_glbmean.nc'
ifile = ifile.replace(case_default, case).replace('0101-0250', '????-????')
ifile = glob.glob(ifile)[0]
"""
print(ifile)
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')
if case == 'co2x2':
    da_fco2 = da_ts*0+1
    da_solar = da_ts*0 
elif case == 'co2x4':
    da_fco2 = da_ts*0+2
    da_solar = da_ts*0 
elif case == 'co2xp5':
    da_fco2 = da_ts*0-1
    da_solar = da_ts*0 
elif case == 'co2xp25':
    da_fco2 = da_ts*0-2
    da_solar = da_ts*0 
elif case == 'p2p0solar':
    da_fco2 = da_ts*0
    da_solar = da_ts*0+1
elif case == 'p6p0solar':
    da_fco2 = da_ts*0
    da_solar = da_ts*0+3
elif case == 'p1p0solar':
    da_fco2 = da_ts*0
    da_solar = da_ts*0+0.5
elif case == 'm2p0solar':
    da_fco2 = da_ts*0
    da_solar = da_ts*0-1
elif case == 'm6p0solar':
    da_fco2 = da_ts*0
    da_solar = da_ts*0-3

#da_fco2 = 2 #4xCO2 forcing (2xCO2 forcing is the units)
#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.2

    #precip
    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'precip ({model} {case})', color='k')
    ax.set_prop_cycle(None)
    y = da_.pipe(func_pct) #save data for later use

    #modeled
    ii = 0; cc = f'C{ii}'
    if case in ('p2p0solar', 'p6p0solar', 'p1p0solar', 'm2p0solar', 'm6p0solar'):
        da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_tropmst + da_solar.groupby('time.year').mean('time')*b_solar
        da_.plot(color=cc, alpha=alpha)
        da_.pipe(func_lp).plot(label=f'modeled by tropmst({b_tropmst:.1f}%/K) and solar({b_solar:.1f}%/2%solar)', color=cc)
    else:
        da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_tropmst + da_fco2.groupby('time.year').mean('time')*b_fco2
        da_.plot(color=cc, alpha=alpha)
        da_.pipe(func_lp).plot(label=f'modeled by tropmst({b_tropmst:.1f}%/K) and co2({b_fco2:.1f}%/2xCO2)', color=cc)
    y_hat = da_ #save

    #from tropmsta
    ii += 1; cc = f'C{ii}'
    da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_tropmst
    da_.plot(color=cc, alpha=alpha)
    da_.pipe(func_lp).plot(label='tropmst contribution', color=cc)
    
    #from direct rad forcing
    ii += 1; cc = f'C{ii}'
    if case in ('p2p0solar', 'p6p0solar', 'p1p0solar', 'm2p0solar', 'm6p0solar'):
        da_ = -da_solar.groupby('time.year').mean('time')*b_solar
        #da_.plot(color=cc, alpha=alpha)
        da_.plot(label='solar contribution (mirrored)', ls='--', color=cc)
    else:
        da_ = -da_fco2.groupby('time.year').mean('time')*b_fco2
        #da_.plot(color=cc, alpha=alpha)
        da_.plot(label='co2 contribution (mirrored)', ls='--', color=cc)

    #da_ = (da_ts.groupby('time.year').mean('time') - da_ref_ts)*b_tropmst_am2p5 + da_fco2.groupby('time.year').mean('time')*b_fco2_am2p5
    #da_.pipe(func_lp).plot(label=f'modeled if AM2.5 coeff: tropmst({b_tropmst_am2p5}%/K) and co2({b_fco2_am2p5}%/2xCO2)')
    
    #residual
    ii += 1; cc = f'C{ii}'
    da_res = y - y_hat
    da_res.plot(color=cc, alpha=alpha)
    da_res.pipe(func_lp).plot(label='residual', color=cc)
    
    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='--')
    ax.axhline(0, color='gray', ls='--')
    #if case in ('4xCO2',):
    #    plt.xscale('log')

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'_{model}_{case}.png')
        if cold_correction and case in ('co2xp25', 'co2xp5',): figname = figname.replace('.png', '.cold-correct.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)

    #
    import seaborn as sns
    import xlinregress
    fig,ax = plt.subplots()
    dprecip = y
    dgmst = da_ts.groupby('time.year').mean('time') - da_ref_ts
    years = slice(101, 250)
    #years = slice(101, 121)
    #years = slice(101, 151)
    #years = slice(121, 251)
    dprecip = dprecip.sel(year=years)
    dgmst = dgmst.sel(year=years)
    #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')
    ax.text(0, 1, f' {model} {case}  tropmst\n y0 = {y0.item():.1f}\n b = {slope.item():.1f}', transform=ax.transAxes, va='top')
    ax.set_xlabel('Tsa [K]')
    ax.set_ylabel('precip anom [%]')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{model}_{case}_linregress{years.start}-{years.stop}.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()