#!/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
from xtc.mask import land_mask
from xtc import tc_basins
import geoxarray
#
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
basin = 'NA'
bs = tc_basins()
omega = xr.open_dataarray('data.omega.nc')
eta = xr.open_dataarray('data.eta.nc')
omfstar = (-omega) * eta
da = (omfstar.where(land_mask<0.5)
    .where( bs.mask(omfstar) == bs.map_keys(basin) ) 
    .sel(lat=slice(0, 30)).geo.fldint()
)
da.name = 'omgfstar'
da.attrs['long_name'] = f'-\omega*(f + \zeta) integrated over {basin}'
da.attrs['units'] = 'Pa s**-2 m**2'

#visualize
da.sel(time=slice('1981', '2018')).mean('en').groupby('time.month').mean('time').plot()
plt.tight_layout()
figname = f'fig.{da.name}.{basin}.png'
plt.savefig(figname)

# save nc file
ofile = f'data.{da.name}.{basin}.nc'
if os.path.exists(ofile):
    print('[exists]:', ofile)
else:
    da.to_dataset().to_netcdf(ofile)
    print('[saved]:', ofile)

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