#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Oct  5 11:09:59 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('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.modelout import get_modelout_data, update_modelout_data
import xfilter
nwindow, dimlp = 9, 'year'
#lowpass = lambda x: x.filter.lowpass(1/nwindow, dim=dimlp, padtype='odd')
lowpass = lambda x: x.rolling(year=nwindow, center=False, min_periods=1).mean()
import geoxarray
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
dTname = os.path.basename(os.getcwd()) #name of current dir: gmst or tropmst
ifile = '/tigress/cw55/work/2021_Constrain_ECS/github_ECS_TCR/data/CMIP6_dtas_dR_data_map.detrend_L2.nc'
ds = xr.open_dataset(ifile).load()
models = [s.split()[0] for s in ds.model.values]
ds['model'] = models
#T
if dTname == 'gmst':
    dT = ds['dtas_abrupt_4xCO2'].geo.fldmean()
elif dTname == 'tropmst':
    dT = ds['dtas_abrupt_4xCO2'].sel(lat=slice(-30,30)).geo.fldmean()
#N
dN = ds['dR_abrupt_4xCO2'].geo.fldmean()

#wrap data
ds = xr.Dataset(dict(
    dT=(dT/2).assign_attrs(units='K', long_name='dT/2'),
    dN=(dN/2).assign_attrs(units='W/m^2', long_name='dN/2')
    ))


"""
def wyplot(**kws):
    ax = kws.pop('ax', plt.gca())
    #ax.scatter(das['gmst'], das['netradtoa'])
    ds.plot.scatter(x='gmst', y='netradtoa', c='none', edgecolors='C0', alpha=0.5)
"""
def wyplot_linregress(model, yearstart=None, yearstop=None, **kws):
    color = kws.pop('color', 'C0')
    if yearstart is None:
        yearstart = 1
    if yearstop is None:
        yearstop = ds.year.size
    ds_ = ds.isel(year=slice(yearstart-1, yearstop)).sel(model=model)
    ax = kws.pop('ax', plt.gca())

    x = ds_['dT']
    y = ds_['dN']
    #ax.plot(x, y, marker='o', ls='none', fillstyle='none', alpha=0.5)
    ds_.plot.scatter(x='dT', y='dN', c='none', edgecolors=color, alpha=0.5)
    reg = y.linregress.on(x, dim='year')
    y0 = reg.intercept.item()
    b = reg.slope.item()
    ax.axline((0,y0), slope=b, ls='--', color=color)
    s = f' years{yearstart}-{yearstop}: y = {b:.2g}x + {y0:.2g}; ECS = {-y0/b:.2g}K'
    ax.text(0, y0, s, color=color)
    ax.text(0.99, 0.98, model, transform=ax.transAxes, ha='right', va='top')
    ax.set_title('')

    
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    plt.close('all')
    #wyplot()
    for ii,model in enumerate(ds.model.values):
        fig,ax = plt.subplots()
        wyplot_linregress(model, 1, 20, color='C0', ax=ax)
        wyplot_linregress(model, 21, 150, color='C1', ax=ax)

        ax.set_xlim(0, None)
        ax.axhline(0, color='gray', ls='--')
    
        #savefig
        if 'savefig' in sys.argv or 's' in sys.argv:
            figname = __file__.replace('.py', f'_{dTname}_{model}.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:
        pass
    else:
        plt.show()