#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun  2 09:50:52 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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
model = 'HIRAM'
expname = 'gv_amipHadISSTlongChancorr_tigercpu_intelmpi_18_540PE'
idir = os.path.join('/tigress/wenchang/MODEL_OUT', model, expname)
dsname = 'GPI'
daname = 'GPI'
years = range(1871,2018+1)
ens = range(1,5+1)
ofile = f'{daname}.{model}.{expname}.{len(ens)}ens.{years[0]}-{years[-1]}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
das_ens = []
for en in ens:
    das_years = []
    for year in years:
        ifile = os.path.join(idir, f'en{en:02d}', 'analysis_wy', 'TCI', f'{year:04d}0101.atmos_month.{dsname}.nc')
        print(f'{en = }; {ifile = }')
        da = xr.open_dataset(ifile)[daname].load()
        das_years.append(da)
    da = xr.concat(das_years, dim='time')
    das_ens.append(da)
da = xr.concat(das_ens, dim=pd.Index(ens, name='ens'))

ds = da.to_dataset(name=daname)
ds.to_netcdf(ofile)
print('[saved]:', ofile)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()