#!/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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
func_yearlyglbmean = lambda da: da.load().geo.fldmean().groupby('time.year').mean('time').sel(year=slice(1979,2014)) #glbmean and yearly mean
def get_glbmean(daname, expname):
    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)
        da = da.pipe(func_yearlyglbmean)
        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):
    #amip
    daname, expname = 'pr', 'amip'
    da = get_glbmean(daname, expname)
    da_amip = da
    #amip-p4K
    daname, expname = 'pr', 'amip-p4K'
    da = get_glbmean(daname, expname)
    da_amip_p4K = da
    #amip ts
    daname, expname = 'ts', 'amip'
    da = get_glbmean(daname, expname)
    ts_amip = da
    #amip-p4K ts
    daname, expname = 'ts', 'amip-p4K'
    da = get_glbmean(daname, expname)
    ts_amip_p4K = da
    #amip-4xCO2
    daname, expname = 'pr', 'amip-4xCO2'
    da = get_glbmean(daname, expname)
    da_amip_4xco2 = da

    #eta and eta_sst
    dP = da_amip_p4K.mean('year')/da_amip.mean('year') * 100 - 100
    dTs = ts_amip_p4K.mean('year') - ts_amip.mean('year')
    eta = dP/dTs
    eta_sst = dP/4
    eta.attrs['units'] = '%/K'
    eta_sst.attrs['units'] = '%/K'
    #adjustment term
    A = da_amip_4xco2.mean('year')/da_amip.mean('year') * 100 - 100
    A.attrs['units'] = '%'
    #save data
    ds = xr.Dataset(dict(
        pr_amip=da_amip, pr_amip_p4K=da_amip_p4K, ts_amip=ts_amip, ts_amip_p4K=ts_amip_p4K,
        pr_amip_4xco2=da_amip_4xco2,
        eta=eta, eta_sst=eta_sst, A=A
        ))
    ds.to_netcdf(ncfile)
    print('[saved]:', ncfile)
else:
    ds = xr.open_dataset(ncfile)
    print('[loaded]:', ncfile)
    da_amip = ds['pr_amip']
    da_amip_p4K = ds['pr_amip_p4K']
    ts_amip = ds['ts_amip']
    ts_amip_p4K = ds['ts_amip_p4K']
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    """
    fig,ax = plt.subplots()
    #da_.plot(hue='model_member')
    for ii,mm in enumerate(da_amip.model_member.values):
        da_amip.sel(model_member=mm).plot(ls='-', color=f'C{ii}', label=mm)
        da_amip_p4K.sel(model_member=mm).plot(ls='--', color=f'C{ii}')
    ax = plt.gca()
    ax.legend()

    fig,ax = plt.subplots()
    da_ = da_amip_p4K/da_amip.mean('year') * 100 - 100
    da_.plot(hue='model_member')
    ax.set_ylabel('%')

    fig,ax = plt.subplots()
    da_ = da_amip_p4K.mean('year')/da_amip.mean('year') * 100 - 100
    da_ = da_/4
    da_.to_pandas().plot.bar()
    ax.set_ylabel('%/K')

    """
    fig,ax = plt.subplots()
    dP = da_amip_p4K.mean('year')/da_amip.mean('year') * 100 - 100
    dTs = ts_amip_p4K.mean('year') - ts_amip.mean('year')
    ax.scatter(dP/4, dP/dTs, label='CMIP6')
    ax.scatter(3.5, 3.3, label='FLOR/AM2.5')
    ax.axline((3,3), slope=1)
    
    ax.set_xlabel('dP/dSST [%/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()