#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Mar  4 19:10:44 EST 2020
import sys, os.path, os, datetime
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
import geoxarray
from xtc import tc_basins
from wystart import *
today = datetime.date.today().strftime('%Y-%m-%d')
print()
 
model = 'ERA5'
basin = 'NA'
# MDR lons and lats
# ventilation index
ifile = f'data.VI.{model}.nc'
vi0_vec = [0.02, 0.04, 0.06]
bs = tc_basins()
vi = xr.open_dataset(ifile)['VI']
vi = (vi
    # NA basin
    .where( bs.mask(vi)==bs.map_keys(basin) )
    # lat range 10-30N
    .where(vi.lat>10).where(vi.lat<30)
    )
p = [vi.pipe(lambda x: 1./(1. + x/vi0)).geo.fldmean() for vi0 in vi0_vec]
p = xr.concat(p, dim=pd.Index(vi0_vec, name='VI0'))
p.name = 'p'

p.groupby('time.month').mean('time').plot(hue='VI0')
title = f'{model}'
plt.title(title)

plt.tight_layout()
figname = f'fig.p.{basin}.{model}.png'
if os.path.exists(figname):
    figname = fig.replace('.png', f'.{today}.png')
plt.savefig(figname)


ofile = f'data.p.NA.{model}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
else:
    p.to_dataset().to_netcdf(ofile)
    print('[saved]:', ofile)

#if __name__ == '__main__':
#    tformat = '%Y-%m-%d %H:%M:%S'
#    t0 = datetime.datetime.now()
#    print('[start]:', t0.strftime(tformat))
#    
#    t1 = datetime.datetime.now()
#    print('[total time used]:', f'{(t1-t0).seconds:,} seconds')
#    print('[end]:', t1.strftime(tformat))
#    print()