#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct 26 15:47:39 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
from misc.landmask import whereocean
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
func_ctl = lambda da: da.load().geo.fldmean().mean('time') #glb time mean
func_ctl_tmsst = lambda da: da.load().pipe(whereocean).sel(lat=slice(-30,30)).geo.fldmean().mean('time') #time mean tmsst
func_glbmean = lambda da: da.load().geo.fldmean().groupby('time.year').mean('time') \
    .isel(year=slice(0, 150)).assign_coords(year=range(1,151)) #glbmean and yearly mean
func_tmsst = lambda da: da.load().pipe(whereocean).sel(lat=slice(-30,30)).geo.fldmean().groupby('time.year').mean('time') \
    .isel(year=slice(0, 150)).assign_coords(year=range(1,151)) #glbmean and yearly mean
def get_glbmean(daname, expname, is_tmsst=False):
    print(daname, expname)
    ifiles = glob.glob(f'{daname}.{expname}.*.nc')
    ifiles.sort()
    das = []
    mms = [] #${model}_${member}s
    for ifile in ifiles:
        model, member = ifile.split('.')[2:4]
        mm = f'{model}_{member}'
        print(mm)
        mms.append(mm)
        da = xr.open_dataarray(ifile)
        if expname == 'piControl':
            if is_tmsst:
                da = da.pipe(func_ctl_tmsst)
            else:
                da = da.pipe(func_ctl)
        else:
            if is_tmsst:
                da = da.pipe(func_tmsst)
            else:
                da = da.pipe(func_glbmean)
        das.append(da)
    da = xr.concat(das, dim=pd.Index(mms, name='model_member'))
    if daname == 'pr': 
        da = (da*24*3600).assign_attrs(units='mm/day')
    elif daname == 'ts':
        da = da.assign_attrs(units='K')

    return da

ncfile = __file__.replace('.py', '.nc')
if not os.path.exists(ncfile):
    #piControl
    daname, expname = 'pr', 'piControl'
    da = get_glbmean(daname, expname)
    da_ctl = da
    #abrupt-4xCO2
    daname, expname = 'pr', 'abrupt-4xCO2'
    da = get_glbmean(daname, expname)
    da_4xco2 = da
    #piControl ts
    daname, expname = 'ts', 'piControl'
    da = get_glbmean(daname, expname)
    ts_ctl = da
    #abrupt-4xCO2 ts
    daname, expname = 'ts', 'abrupt-4xCO2'
    da = get_glbmean(daname, expname)
    ts_4xco2 = da
    #piControl ts tmsst
    daname, expname = 'ts', 'piControl'
    da = get_glbmean(daname, expname, is_tmsst=True)
    tmsst_ctl = da
    #abrupt-4xCO2 ts tmsst
    daname, expname = 'ts', 'abrupt-4xCO2'
    da = get_glbmean(daname, expname, is_tmsst=True)
    tmsst_4xco2 = da

    #eta and eta_sst
    da_4xco2_anom = da_4xco2.pipe(lambda da: da/da_ctl*100 - 100).assign_attrs(units='%')
    ts_4xco2_anom = ts_4xco2.pipe(lambda da: da - ts_ctl).assign_attrs(units='K')
    tmsst_4xco2_anom = tmsst_4xco2.pipe(lambda da: da - tmsst_ctl).assign_attrs(units='K')
    eta = da_4xco2_anom.linregress.on(ts_4xco2_anom, dim='year').slope
    eta_tmsst = da_4xco2_anom.linregress.on(tmsst_4xco2_anom, dim='year').slope
    #save data
    ds = xr.Dataset(dict(
        pr_4xco2_anom=da_4xco2_anom, pr_ctl=da_ctl, 
        ts_4xco2_anom=ts_4xco2_anom, ts_ctl=ts_ctl,
        tmsst_4xco2_anom=tmsst_4xco2_anom, tmsst_ctl=tmsst_ctl,
        eta=eta, eta_tmsst=eta_tmsst
        ))
    ds.to_netcdf(ncfile)
    print('[saved]:', ncfile)
else:
    ds = xr.open_dataset(ncfile)
    print('[loaded]:', ncfile)

iyearspan = slice(None, 20)
da = ds['pr_4xco2_anom'].isel(year=iyearspan)
ts = ds['ts_4xco2_anom'].isel(year=iyearspan)
tmsst = ds['tmsst_4xco2_anom'].isel(year=iyearspan)
eta = da.linregress.on(ts, dim='year').slope
eta_tmsst = da.linregress.on(tmsst, dim='year').slope
ds['eta'] = eta
ds['eta_tmsst'] = eta_tmsst

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots()
    ds.plot.scatter(x='eta_tmsst', y='eta')
    
    ax.set_xlabel('dP/dTMSST [%/K]')
    ax.set_ylabel('dP/dTs [%/K]')
    ax.legend()
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.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()