#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Mar 31 14:53:32 EDT 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    print()
    today = datetime.date.today()
    today_s = today.strftime('%Y-%m-%d')
    tformat = '%Y-%m-%dT%H:%M:%S'
    t_start = datetime.datetime.now()
    print('[start]:', t_start.strftime(tformat))
    
#start from here
model = 'ERA5'
expname = 'reanalysis'
daname_raw = 'w'
daname = 'omega'
ens = None#range(1, 6)
years = range(1979, 2019)
ofile = f'data.{daname}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
    sys.exit()
ifiles = [f'/tigress/wenchang/data/era5/raw/monthly/plevels/vertical_velocity/era5.vertical_velocity.monthly.{year:04d}.nc'
    for year in years]

# read data
print('reading data ...')
da = (
    xr.open_mfdataset(ifiles, combine='by_coords')[daname_raw]
    .sel(level=500)
    .rename(latitude='lat', longitude='lon')
    .load()
    )

# create dataset
ds = da.to_dataset(name=daname)
ds.attrs['model'] = model
ds.attrs['expname'] = expname

# save data
print('saving data to netcdf ...')
encoding = {daname: {'dtype': 'float32', 'zlib': True, 'complevel': 1}}
ds.to_netcdf(ofile, encoding=encoding)
print('[saved]:', ofile)

if __name__ == '__main__':
    t_end = datetime.datetime.now()
    print('[end]:', t_end.strftime(tformat))
    print('[total time used]:', f'{(t_end - t_start).seconds:,} seconds')
    print()