#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Apr 22 11:04:47 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
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__':
   tt.check('end import')

#start from here
basin = 'NA'
#oname = 'omgfstar'
ofile = f'data.omgfstars.{basin}.nc'
if os.path.exists(ofile):
    ds = xr.open_dataset(ofile)
    print('[exists]:', ofile)
else:
    bs = tc_basins()
    omega = xr.open_dataarray('data.omega.nc')
    #omega = omega.where(omega<0) # only select values of upward motion
    omega_m = (omega.where(land_mask<0.5)
        .where( bs.mask(omfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(0, 30)).geo.fldmean()
    ) #basin mean
    #omega_a = omega - omega_m # anomaly from basin mean
    eta = xr.open_dataarray('data.eta.nc')
    omgfstar = (-omega) * eta
    da = (omgfstar.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'
    omgfstar = da

    # basin mean omega used
    da = (-omega_m) * eta
    da = (da.where(land_mask<0.5)
        .where( bs.mask(omfstar) == bs.map_keys(basin) ) 
        .sel(lat=slice(0, 30)).geo.fldint()
    )
    da.name = 'omg_mfstar' 
    da.attrs['long_name'] = f'-\omega*(f + \zeta) integrated over {basin}: basin mean \omega used'
    da.attrs['units'] = 'Pa s**-2 m**2'
    omg_mfstar = da

    # save nc file
    ds = xr.Dataset(dict(omgfstar=omgfstar, omg_mfstar=omg_mfstar))
    ds.to_netcdf(ofile)
    print('[saved]:', ofile)

if __name__ == '__main__':
    from wyconfig import *
    #visualize
    plt.figure()
    ds.omgfstar.sel(time=slice('1981', '2018')).mean('en').groupby('time.month').mean('time').plot(label='CTL', color='k', marker='o')
    ds.omg_mfstar.sel(time=slice('1981', '2018')).mean('en').groupby('time.month').mean('time').plot(label='basin-mean $\omega$ used', marker='o')
    (ds.omgfstar - ds.omg_mfstar).sel(time=slice('1981', '2018')).mean('en').groupby('time.month').mean('time').plot(label='anom from basin-mean $\omega$ used', marker='o')
    plt.legend()
    plt.xticks(range(1,13))
    plt.ylabel('Pa s$^{-1}$ m$^2$')
    plt.axhline(0, ls='--', color='gray')
    plt.title(f'$-\omega\\times f^*$ integrated over the {basin} basin of 0-30N')

    figname = __file__.replace('.py', f'_{basin}_{tt.today()}.png')
    plt.savefig(figname)
    plt.show()