#!/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.mean('xt_ocean')#('lon')
#daname = ['precip', 'netrad_toa', 't_surf',][1]
#dsname = 'atmos_month'
daname = 'temp'
dsname = 'ocean'
volcs = ['CTL', 'Pinatubo', 'Agung', 'StMaria', 'Chichon', 'Chichon_CMIP5volc']
nens = 30
thisdir = os.path.dirname(__file__)

#ofile
ofile = os.path.join(thisdir, f'{daname}.x170W120Wmean.volcs.nens{nens}.40S-40N_0-300m.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.170W-120W_40S-40N_0-300m.nc'
print(ifile, '...')
da = xr.open_dataarray(ifile)
units_original = da.attrs['units']
da = da.pipe(func)
da_ctl = da
#volcs
das = [da,]
for volc in volcs[1:]:
    ifile = f'{daname}.{volc}_PI_ens_noleap.{dsname}.nens{nens}.170W-120W_40S-40N_0-300m.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'))
if daname in ('precip',):
    #convert units from kg/s/m^2 to mm/day
    da = da.pipe(lambda x: x*24*3600).assign_attrs(units='mm/day')
else:
    da = da.assign_attrs(units=units_original)
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()