#!/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

    #derivative of eta to y and Z
    R = 6370*100# Earth's radius in unites of m
    U = 20
    y = R * eta.lat/180*np.pi #y coordinate to replace "lat"
    #use y coordinate (units of m) to do the differentiation and convert back to "lat" after that
    eta_y = eta.rename(lat='y').assign_coords(y=y.values).differentiate('y').rename(y='lat').assign_coords(lat=eta.lat.values)
    #rolling mean over lon and lat
    nx, ny = 10, 3
    eta_y = eta_y.pad(lon=nx, mode='wrap').rolling(lon=nx*2+1, center=True).mean().isel(lon=slice(nx,-nx)).rolling(lat=ny*2+1, center=True).mean()
    eta_y = eta_y.where(np.abs(eta_y)>0) # remove zero values
    #Z value in Tsung-Lin's paper: eta uses absolute value of cyclonic grids and zero for anticyclonic grids
    Z = eta.where(eta.lat>0, other=-eta).pipe(lambda x: x.where(x>0, other=0)) / np.sqrt( np.abs(eta_y)*U )
    
    cases = []
    das = []
    
    case = 'ctl'
    print(case)
    omgZ = (-omega) * Z
    da = (omgZ.where(landflag<0.5)
        .where( bs.mask(omgZ) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    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.fldmean()
    )
    Z_ = (Z.where(landflag<0.5)
        .where( bs.mask(Z) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    da = (-omega_) * Z_
    das.append(da)
    cases.append(case)
    
    case = 'aclimZ'
    print(case)
    # use annual mean climatology of fstar
    func_aclim = lambda x: x.sel(time=years_clim).mean('time')
    omgZ = (-omega) * Z.pipe(func_aclim)
    da = (omgZ.where(landflag<0.5)
        .where( bs.mask(omgZ) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    das.append(da)
    cases.append(case)
    
    case = 'aclimBmeanZ'
    print(case)
    # use basin mean annual mean climatology of Z
    func_aclim = lambda x: x.sel(time=years_clim).mean('time')
    func_bmean = lambda x: x.where(landflag<0.5).where(bs.mask(x) == bs.map_keys(basin)).sel(lat=slice(30,0)).geo.fldmean() 
    omgZ = (-omega) * Z.pipe(func_aclim).pipe(func_bmean)
    da = (omgZ.where(landflag<0.5)
        .where( bs.mask(omgZ) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    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')
    omgZ = (-omega.pipe(func_aclim)) * Z
    da = (omgZ.where(landflag<0.5)
        .where( bs.mask(omgZ) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    das.append(da)
    cases.append(case)
    
    case = 'upOnly'
    print(case)
    omgZ = (-omega.where(omega<0)) * Z
    da = (omgZ.where(landflag<0.5)
        .where( bs.mask(omgZ) == bs.map_keys(basin) ) 
        .sel(lat=slice(30, 0)).geo.fldmean()
    )
    das.append(da)
    cases.append(case)
    
    
    print('concatenating...')
    da = xr.concat(das, dim=pd.Index(cases, name='case'))
    da.name = 'omgZ'
    da.attrs['long_name'] = f'-\omega*Z averaged over {basin}'
    da.attrs['units'] = 'Pa s**-1'
    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$^{-1}$')
    plt.title(r'$-\omega\times$Z averaged 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()