#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun 17 11:20:02 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
maindir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if maindir not in sys.path:
    sys.path.append(maindir)
from ERA5.data_seedgenesis_map_noLP import da as seed
from ERA5.data_pTang_map import da as p
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
savedfile = os.path.abspath(__file__).replace('.py', '.nc')
data_name = 'seedxp'
long_name = 'seed genesis density $\\times$ p'

if os.path.exists(savedfile):
    print('**reading from saved file**:', savedfile)
    with xr.open_dataset(savedfile) as ds:
        da = ds[data_name]
else:
    p_ = p.interp(lon=seed.lon, lat=seed.lat).assign_coords(time=seed.time)
    seedxp = seed*p_
    da = seedxp
    da = da.assign_attrs(long_name=long_name)
    da.to_dataset(name=data_name).to_netcdf(savedfile)
    print('**saved**:', savedfile)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    years = slice('1980', '2018')
    figname = __file__.replace('.py', f'.aclim{years.start}-{years.stop}.{tt.today()}.png')
    da.sel(time=years).mean('time').pipe(lambda x: x*12) \
        .assign_attrs(long_name=da.attrs['long_name'])\
        .plot.contourf(cmap='OrRd', levels=2.**np.arange(-2-5, 3.1-5, 0.5), cbar_kwargs=dict(aspect=30))
    mapplot()
    ax.set_title(f'ERA5 {years.start}-{years.stop} seed genesis density $\\times$p')

    plt.savefig(figname)
    print('**saved**:', figname)
    
    tt.check(f'**Done**')
    plt.show()