#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Mar  2 17:15:22 EST 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#print()

model = 'AM2.5'
expname = 'amipHadISST_tigercpu_intelmpi_18_540PE'
data_name = 'Vshear'
ens = range(1, 6)
years = range(1971, 2019)
ofile = f'data.{data_name}.nc'
ifiles = [[f'/tigress/wenchang/MODEL_OUT/{model}/{expname}/en{en:02d}/analysis_wy/TCI/{year:04d}0101.atmos_month.Vshear.nc'
    for year in years] for en in ens]

 
if __name__ == '__main__':
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.datetime.now()
    print('[start]:', t0.strftime(tformat))
    
    ds = xr.open_mfdataset(ifiles, concat_dim=['en', 'time'], combine='nested')
    ds = ds.rename({'grid_xt': 'lon', 'grid_yt': 'lat'})
    ds.attrs['model'] = model
    ds.attrs['expname'] = expname
    ds.load()

    encoding = {data_name: {'dtype': 'float32', 'zlib': True, 'complevel': 1}}
    ds.to_netcdf(ofile, encoding=encoding)
    print('[saved]:', ofile)

    t1 = datetime.datetime.now()
    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
    print('[end]:', t1.strftime(tformat))
    print()