#!/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
func = lambda x: x.sel(lat=slice(-5, 5)).mean('lat')
daname = ['t_surf',][0]
dsname = 'atmos_month'
volcs = ['CTL', 'Pinatubo', 'Agung', 'StMaria', 'Chichon', 'Chichon_CMIP5volc']
nens = 30
thisdir = os.path.dirname(__file__)

#ofile
ofile = os.path.join(thisdir, f'{daname}.y5S5Nmean.volcs.nens{nens}.nc')
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
#ctl
ifile = f'{daname}.CTL1860_noleap_tigercpu_intelmpi_18_576PE.{dsname}.nens{nens}.nyears15.nc'
print(ifile, '...')
da = xr.open_dataarray(ifile).pipe(func)
da_ctl = da
#volcs
das = [da,]
for volc in volcs[1:]:
    ifile = f'{daname}.{volc}_PI_ens_noleap.{dsname}.nens{nens}.nc'    
    print(ifile, '...')
    da = xr.open_dataarray(ifile).pipe(func)
    da = da.assign_coords(time=da_ctl.time.values)
    das.append(da)
print('concatenating...')
da = xr.concat(das, dim=pd.Index(volcs, name='volc'))
print('saving...')
encoding = {daname: dict(zlib=True, complevel=1)}
da.to_dataset(name=daname).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()