#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Dec  4 01:28:10 PM EST 2025
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
if 'rmWarming' in sys.argv:
    expname = 'amipHadISSTlong_chancorr_rmWarming_tigercpu_intelmpi_18_1080PE'
    ens = range(6,10+1)
    nens = len(ens)
    years = range(1982, 2019+1)
elif 'future' in sys.argv:
    expname = 'amipHadISSTrcp45_tigercpu_intelmpi_18_1080PE'
    ens = [6,7,8,16,17,18] 
    nens = len(ens)
    years = range(2022, 2099+1)
else:
    expname = 'amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE'
    ens = range(1,10+1)
    nens = len(ens)
    years = range(1982, 2019+1)
print(expname, ens, years)

if 'vws' in sys.argv:
    dsname = 'vertical_wind_shear'
else:
    dsname = 'steering_wind'
print(dsname)
ofile = f'{dsname}_T11_{expname}_{nens}ens_{years[0]}-{years[-1]}.nc'
if os.path.exists(ofile): print('[exists]:', ofile); sys.exit()

ds_ens = []
for ii in ens:
    print(f'{ii:02d} of {nens}')
    ifiles = f'{expname}/{dsname}_T11_{expname}_ens{ii:02d}_????.nc'
    print(ifiles)
    ds = xr.open_mfdataset(ifiles)
    ds_ens.append(ds)
print('concat...')
ds = xr.concat(ds_ens, dim=pd.Index(ens, name='ens'))

#save
print('save...')
ds.to_netcdf(ofile)
print('[saved]:', ofile)
 
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()