#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Apr  1 16:07:25 EDT 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
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
daname = 'omega'
units = 'Pa/s'
level = 500 #hPa
ens = range(1,6)
years = range(1980,2019)
ofile = os.path.join(os.path.dirname(__file__), 'data.omega500.nc')
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile).load()
    print('[loaded]:', ofile)
else:
    das = []
    for en in ens:
        print(f'en = {en:02d}')
        ifiles = [f'/tigress/wenchang/MODEL_OUT/HIRAM/amipHadISST_tigercpu_intelmpi_18_540PE/en{en:02d}/POSTP/{year}0101.atmos_month.nc'
            for year in years]
        print('loading...')
        da = xr.open_mfdataset(ifiles)[daname].load().interp(pfull=level)
        da = da.groupby('time.month').mean('time')
        if 'grid_yt' in da.dims:
            da = da.rename(grid_yt='lat', grid_xt='lon')
        das.append(da)
    print('concatenating...')
    da = xr.concat(das, dim=pd.Index(ens, name='en'))
    da.attrs['units'] = units
    ds = da.to_dataset()
    ds = da.to_dataset(name=daname)
    print('saving...')
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)
 
sel_region = lambda x: x.where(x.lon+x.lat*35/20>295)

def wyplot(ax=None, month=10, lonmin=270, lonmax=360, latmin=0, latmax=30, **kws):
    from geoplots import mapplot
    if ax is None:
        fig, ax = plt.subplots()
    ds[daname].rename('omega500').mean('en').assign_attrs(units=units).sel(month=month, lat=slice(latmin,latmax), lon=slice(lonmin,lonmax)) \
        .pipe(sel_region) \
        .plot(ax=ax, robust=True, **kws)
    mapplot(xticks=range(270,361,30), yticks=range(0,31,10))
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot(vmax=0.1, center=0, levels=11, extend='both', cmap='RdBu_r')
    
    #savefig
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()