#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Feb  4 15:39:30 EST 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
from misc.landmask import flagland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
idaname = 'TS'
daname = 'sst'
ifiles = 'LME_TS/b.e11.B1850C5CN.f19_g16.0850cntl.001.cam.h0.TS.*.nc'
ofile = f'b.e11.B1850C5CN.f19_g16.0850cntl.001.cam.h0.TS.0850-2005.nc'
nc_encoding = {daname: dict(dtype='float32', zlib=True, complevel=1)}
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    ds_ = xr.open_mfdataset(ifiles)
    print('loading...')
    ds_.load()
    ds_['time'] = ds_.indexes['time'].shift(-1, 'MS') #cesm model output needs one month shift back
    print('[time shifted]:', 'one month back')
    da = ds_[idaname].groupby('time.year').mean('time').pipe(lambda x: x.where(flagland(x)<0.5))
    da.attrs['long_name'] = 'sea surface temperatrue'#ds_[idaname].attrs['long_name']
    da.attrs['units'] = ds_[idaname].attrs['units']
    ds = da.to_dataset(name=daname)
    ds.to_netcdf(ofile, encoding=nc_encoding)
    print('[saved]:', ofile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import geoxarray, xfilter
    da = ds[daname].geo.fldmean()
    da.plot()
    plt.figure()
    da.filter.bandpass([1/500, 1/40], dim='year', padtype='even').plot()
    
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname)
        print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()