#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun 13 11:36:47 EDT 2024
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 xesmf as xe
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#nino34 from nudgeLongAll
ifile = '/tigress/wenchang/analysis/FLOR/ens/nino34/t_surf_FLOR_nudgeLongAll_newdiag_tigercpu_intelmpi_18_576PE_6ens_1871-2016_nino34.nc'
da = xr.open_dataarray(ifile)
nino34 = da - 273.15
#from nudgeAll
ifile = '/tigress/wenchang/analysis/FLOR/ens/nino34/t_surf_FLOR_nudgeAll_e1_newdiag_tigercpu_intelmpi_18_576PE_1979-2016_nino34.nc'
da = xr.open_dataarray(ifile)
nino34_e1 = da - 273.15
ifile = '/tigress/wenchang/analysis/FLOR/ens/nino34/t_surf_FLOR_nudgeAll_e2_newdiag_tigercpu_intelmpi_18_576PE_1979-2016_nino34.nc'
da = xr.open_dataarray(ifile)
nino34_e2 = da - 273.15

#obs nino34
ifile = '/home/wenchang//sGEOCLIM/CM2.5_data/input/SST_TARGETS/temp_restore_HadISST_1870-2017.nc'
da_obs = xr.open_dataarray(ifile, decode_times=False) \
    .rename(TIME='time', LONGITUDE='lon', LATITUDE='lat') \
    .assign_coords(time=pd.date_range('1870-01', '2017-06',freq='MS')) \
    .sel(time=slice('1871', '2016')) \
    .assign_coords(time=nino34['time']) #set the same time axis as the model
#nino34_obs = da_obs.sel(lon=slice(-170,-120), lat=slice(-5,5)).geo.fldmean()
#regrid obs data onto model grid
ds_o = xr.open_dataset('/tigress/wenchang/analysis/FLOR/ens/raw/t_surf_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_ens01_1860-2100_atmos_month.nc') #target grid
regridder = xe.Regridder(da_obs, ds_o, 'bilinear')
da_regrid = regridder(da_obs)
nino34_obs = da_regrid.sel(lon=slice(360-170,360-120), lat=slice(-5,5)).geo.fldmean()

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots(figsize=(8,3))
    #tslice = slice('1871', '1900')
    tslice = slice('1979', '2016')
    seltime = lambda x: x.sel(time=tslice)
    nino34_obs.pipe(seltime).plot(color='k', label='HadISST')
    for ens in nino34.ens.values:
        label = 'nudgeLongAll' if ens == 1 else None
        nino34.sel(ens=ens).pipe(seltime).plot(alpha=0.5, label=label, lw=1, color='gray')
    #nudgeAll
    nino34_e1.pipe(seltime).plot(label='nudgeAll_e1', lw=1, color='C0', alpha=0.5)
    nino34_e2.pipe(seltime).plot(label='nudgeAll_e2', lw=1, color='C1', alpha=0.5)

    ax.legend(ncol=4, loc='lower right', borderpad=0)
    ax.set_title('Nino3.4 from obs and nudged model historical experiment')
    ax.set_xlabel('')
    ax.set_ylabel('degC')

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