#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Apr 22 11:04:47 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
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from ERA5.data_omega500Z import da
#
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
if __name__ == '__main__':
    from wyconfig import *
    years_clim = slice('1980', '2018')

    plt.figure()
    #figname = __file__.replace('.py', f'_{today_s}.png')
    da_ = da.sel(case=['ctl', 'aclimZ', 'aclimOmega']) \
        .assign_coords(case=['CTL', 'use annual mean clim. Z', 'use annual mean clim. $\omega$'])
    da_.sel(time=years_clim).groupby('time.month').mean('time').plot(hue='case', marker='o', fillstyle='none')
    plt.axhline(0, color='gray', ls='--')
    plt.xticks(range(1,13))
    plt.ylabel(r'Pa s$^{-1}$')
    plt.title(r'$-\omega\times$Z averaged over NA from ERA5')
    
    figname = __file__.replace('.py', '.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname)

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