#!/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
today = datetime.date.today().strftime('%Y-%m-%d')
print()

note = ''' 
model = 'HIRAM'
expname = 'amipHadISST_tigercpu_intelmpi_18_540PE'
basin = 'NA'
# MDR lons and lats
# ventilation index
ifile = 'data.VI.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)**2)).geo.fldmean() for vi0 in vi0_vec]
p = xr.concat(p, dim=pd.Index(vi0_vec, name='VI0'))
p.name = 'p'
'''
ifile = 'data.p2.NA.nc'
p = xr.open_dataarray(ifile)

p.mean('en').groupby('time.month').mean('time').plot(hue='VI0')
title = f'{model}.{expname}: 1/(1+(VI/VI0)$^2$)'
plt.title(title)

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

'''
ofile = f'data.p2.NA.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()