#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Sep 11 08:42:49 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 cftime import DatetimeNoLeap
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#LMR sst anomaly plus HadISST base climatology
lmr_version = '2018' if '2018' in sys.argv else '2019'
ofile = f'LMR{lmr_version}_nino34.nc'

if os.path.exists(ofile) and not 'o' in sys.argv:
    print('[loaded]:', ofile)
    da = xr.open_dataarray(ofile)
else:
    #lmr
    print('LMR', lmr_version)
    ifile = f'/tigress/gvecchi/DATA/LMR_{lmr_version}/sst_MCruns_ensemble_mean.nc'
    print(ifile)
    da = xr.open_dataarray(ifile)
    #da.isel(time=0, MCrun=0).plot(); plt.show(); sys.exit()
    #da = da.groupby('time.year').mean('time') #raw LMR sst is yearly data, rename 'time'->'year'

    #nino34
    da = da.sel(lon=slice(360-170,360-120), lat=slice(-5,5)).geo.fldmean()
    da = da.groupby('time.year').mean('time') #'time' -> 'year'
    
    #save
    print('saving...')
    da.to_dataset(name='nino34').to_netcdf(ofile)
    print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xfilter
    fig,ax = plt.subplots()
    lpwindow = 15
    func = lambda da: da.filter.lowpass(1/lpwindow, dim='year').sel(year=slice(850,None))
    da.mean('MCrun').pipe(func).plot()

    years = 1200,1400
    ax.axvspan(*years, color='gray', alpha=0.2)
    damean = da.sel(year=slice(*years)).mean(['MCrun', 'year'])
    ax.plot(years, [damean, damean], ls='--')

    years = 1450,1550
    ax.axvspan(*years, color='gray', alpha=0.2)
    damean = da.sel(year=slice(*years)).mean(['MCrun', 'year'])
    ax.plot(years, [damean, damean], ls='--')

    years = 1551,1650
    ax.axvspan(*years, color='gray', alpha=0.2)
    damean = da.sel(year=slice(*years)).mean(['MCrun', 'year'])
    ax.plot(years, [damean, damean], ls='--')

    ax.set_title(f'LMR{lmr_version} nino34, {lpwindow}-year lowpass')

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        #figname = __file__.replace('.py', f'.png')
        figname = ofile.replace('.nc', '.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()