#!/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 = 'w'
units = 'Pa/s'
level = 500 #hPa
years = range(1980,2019)
ifiles = [f'/tigress/wenchang/data/era5/raw/monthly/plevels/vertical_velocity/era5.vertical_velocity.monthly.{year}.nc'
    for year in years]
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:
    ds = xr.open_mfdataset(ifiles).sel(level=level).load()
    ds = ds.groupby('time.month').mean('time')
    if 'latitude' in ds.dims:
        ds = ds.rename(latitude='lat', longitude='lon')
    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') \
        .assign_attrs(units=units) \
        .sel(month=month, lat=slice(latmax,latmin), lon=slice(lonmin,lonmax)) \
        .pipe(sel_region) \
        .plot(ax=ax, robust=True, **kws)
    mapplot(xticks=range(270,361,30), yticks=range(0,31,10))
    #ax.set_xticks(range(270,361,30))
    
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    wyplot(vmax=0.1, extend='both', levels=11, center=0, 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()