#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb 22 12:14:06 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, matplotlib.pyplot as plt
#more imports
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#nino3.4 region
lons = slice(360-170, 360-120)
lats = slice(-5, 5)
daname = ['t_surf',][0]
dsname = 'atmos_month'
#nens, nyears = 30, 15
#for the first ensemble member (100 years available)
#nens, nyears = 1, 100
nens, nyears = 1, 1000
thisdir = os.path.dirname(__file__)

#ofile
ofile = os.path.join(thisdir, f'nino34.ctl.nyears{nyears}.nc')
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
#ctl
ifile = f'{daname}.CTL1860_noleap_tigercpu_intelmpi_18_576PE.{dsname}.nens{nens}.nyears{nyears}.nc'
print(ifile, '...')
da = xr.open_dataarray(ifile).sel(lon=lons, lat=lats).geo.fldmean()
da_ctl = da
#ensemble-mean monthly climatology
da_ctl_mclim = da_ctl.mean('en').groupby('time.month').mean('time')
#anomaly from the ensemble-mean monthly climatology
da_ctl = da_ctl.groupby('time.month') - da_ctl_mclim
"""
#volcs
das = [da_ctl,]
for volc in volcs[1:]:
    ifile = f'{daname}.{volc}_PI_ens_noleap.{dsname}.nens{nens}.nc'    
    print(ifile, '...')
    da = xr.open_dataarray(ifile).sel(lon=lons, lat=lats).geo.fldmean()
    da = da.assign_coords(time=da_ctl.time.values)
    #anomaly from the ensemble-mean monthly climatology of the ctl run
    da = da.groupby('time.month') - da_ctl_mclim
    das.append(da)
print('concatenating...')
da = xr.concat(das, dim=pd.Index(volcs, name='volc'))
"""
da = da_ctl.sel(en=1).drop('en')
da.attrs['note'] = f'Nino3.4 index (sst averaged over {lons.start}-{lons.stop} and {-lats.start}S-{lats.stop}N'
print('saving...')
encoding = {'nino34': dict(zlib=True, complevel=1)}
da.to_dataset(name='nino34').to_netcdf(ofile, encoding=encoding)
print('[saved]:', ofile)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        figname = __file__.replace('.py', '.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()