#!/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 misc.landmask import flagland
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
years_clim = slice('1980', '2018')
datafile = __file__.replace('.py', '.nc')
if os.path.exists(datafile):
    da = xr.open_dataarray(datafile)
    print('[loaded]:', datafile)
else: # compute and save if the data file does not exist
    basin = 'NA'
    bs = tc_basins()
    omega = xr.open_dataarray('data.omega.nc').load()
    eta = xr.open_dataarray('data.eta.nc').load()
    landflag = flagland(omega)# 1 over land and 0 over ocean
    
    cases = []
    das = []
    
    case = 'ctl'
    print(case)
    omgfstar = (-omega) * eta
    da = (omgfstar.where(landflag<0.5)
        .where( bs.mask(omgfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    das.append(da)
    cases.append(case)
    
    case = 'basinMeanFirst'
    print(case)
    # basin mean \omega and f* first
    omega_ = (omega.where(landflag<0.5)
        .where( bs.mask(omega) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    eta_ = (eta.where(landflag<0.5)
        .where( bs.mask(eta) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    da = (-omega_) * eta_
    das.append(da)
    cases.append(case)
    
    case = 'aclimfstar'
    print(case)
    # use annual mean climatology of fstar
    func_aclim = lambda x: x.sel(time=years_clim).mean('time')
    omgfstar = (-omega) * eta.pipe(func_aclim)
    da = (omgfstar.where(landflag<0.5)
        .where( bs.mask(omgfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    das.append(da)
    cases.append(case)
    
    case = 'aclimOmega'
    print(case)
    # use annual mean climatology of omega
    func_aclim = lambda x: x.sel(time=years_clim).mean('time')
    omgfstar = (-omega.pipe(func_aclim)) * eta
    da = (omgfstar.where(landflag<0.5)
        .where( bs.mask(omgfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    das.append(da)
    cases.append(case)
    
    case = 'upOnly'
    print(case)
    omgfstar = (-omega.where(omega<0)) * eta
    da = (omgfstar.where(landflag<0.5)
        .where( bs.mask(omgfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    das.append(da)
    cases.append(case)
    
    """
    case = 'sm10deglon'
    print(case)
    # smooth by rolling average along 10x10deg box
    npad = 20
    nwindow = npad*2+1
    func = lambda x: x.pad(lon=npad).rolling(lon=nwindow, center=True).mean().isel(lon=slice(npad, -npad))
    omgfstar = (-omega.pipe(func)) * eta.pipe(func)
    da = (omgfstar.where(landflag<0.5)
        .where( bs.mask(omgfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldint()
    )
    das.append(da)
    cases.append(case)
    """
    
    print('concatenating...')
    da = xr.concat(das, dim=pd.Index(cases, name='case'))
    da.name = 'omgfstar'
    da.attrs['long_name'] = f'-\omega*(f + \zeta) integrated over {basin}'
    da.attrs['units'] = 'Pa s**-2 m**2'
    print('saving to dataset ...')
    ds = da.to_dataset()
    ds.to_netcdf(datafile)
    print('[saved]:', datafile)

if __name__ == '__main__':
    from wyconfig import *
    plt.figure()
    figname = __file__.replace('.py', f'_all_{today_s}.png')
    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$^{-2}$ m$^2$')
    plt.title(r'$-\omega\times$(f + $\zeta$) integrated over NA')
    plt.savefig(figname)
    print('[saved]:', 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()