#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue May  2 14:23:32 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
import geoxarray
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
detrend_on = True if 'detrend' in sys.argv else False
dTname = 'gmst'
ifile = '/tigress/cw55/work/2021_Constrain_ECS/github_ECS_TCR/data/CMIP6_dtas_dR_data_map.no_detrend.nc'
if detrend_on: ifile = ifile.replace('no_detrend', 'detrend_L2') #replace with detrend_L2 file
ds = xr.open_dataset(ifile).load()
models = [s.split()[0] for s in ds.model.values]
ds['model'] = models
#T
dT = ds['dtas_abrupt_4xCO2'].geo.fldmean()
"""
if dTname == 'gmst':
elif dTname == 'tropmst':
    dT = ds['dtas_abrupt_4xCO2'].sel(lat=slice(-30,30)).geo.fldmean()
"""
#N
dN = ds['dR_abrupt_4xCO2'].geo.fldmean()

def cal_ratio(dN, dT):
    #years 1-20
    yearspan = slice(0,20)
    reg = dN.isel(year=yearspan).linregress.on(dT.isel(year=yearspan), dim='year')
    lambda_early = reg.slope
    #years 21-150
    yearspan = slice(20,150)
    reg = dN.isel(year=yearspan).linregress.on(dT.isel(year=yearspan), dim='year')
    lambda_late = reg.slope
    lambda_ratio = lambda_late/lambda_early
    
    return lambda_ratio, lambda_late

lambda_ratio_gmst, lambda_late_gmst = cal_ratio(dN, dT)
#tropmst
dT = ds['dtas_abrupt_4xCO2'].sel(lat=slice(-30,30)).geo.fldmean()
lambda_ratio_tropmst, lambda_late_tropmst = cal_ratio(dN, dT)


#wrap data
ds = xr.Dataset(dict(
    lambda_ratio_gmst=lambda_ratio_gmst.assign_attrs(long_name='$\lambda$(21-150)/$\lambda$(1-20) for GMST'),
    lambda_ratio_tropmst=lambda_ratio_tropmst.assign_attrs(long_name='$\lambda$(21-150)/$\lambda$(1-20) for TROPMST'),
    lambda_late_gmst=lambda_late_gmst.assign_attrs(long_name='$\lambda$(21-150) for GMST'),
    lambda_late_tropmst=lambda_late_tropmst.assign_attrs(long_name='$\lambda$(21-150) for TROPMST'),
    ))

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figsize = 7,6
    fig,ax = plt.subplots(figsize=figsize)
    color = 'C0'
    ds.plot.scatter(x='lambda_ratio_gmst', y='lambda_ratio_tropmst', c='none', edgecolors=color, label='CMIP6 models')
    ax.scatter(0.66, 0.97, c='none', edgecolors='C1', label='FLOR2xCO2')
    #label some models
    for model in ds.model.values:
        x = ds['lambda_ratio_gmst'].sel(model=model).item()
        y = ds['lambda_ratio_tropmst'].sel(model=model).item()
        if x>0.96 or y-x>0.1425 or y<0.4: #pick some best (abs(y-1)<0.05) or worst (x>1) models
            if x>0.97 and x<0.98:
                ax.text(x, y, model, ha='center', va='top')
            else:
                ax.text(x, y, model, ha='center', va='bottom')

    
    """
    for ii,model in enumerate(ds.model.values, start=1):
        ax.text(ds.lambda_ratio_gmst.sel(model=model), ds.lambda_ratio_tropmst.sel(model=model), 
            f'{ii}', color='gray', ha='center', va='center')
    """
        

    ax.axhline(1, color='gray', ls='--')
    ax.axvline(1, color='gray', ls='--')
    ax.axline((1,1), slope=1, color='gray', ls='--')
    N = (np.abs(ds.lambda_ratio_tropmst-1)<np.abs(ds.lambda_ratio_gmst-1)).sum().item()
    ax.set_title(f'{N} of {ds.model.size} CMIP6 models show less change of $\lambda$(1-20) to $\lambda$(21-150)')

    ax.legend()
    ax.set_aspect('equal')
    ax.set_ylim(None, 1.8)
    ax.set_xlim(None, 1.8)
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if detrend_on: figname = figname.replace('.png', '_withDetrend.png')
        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()