#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu May 16 15:39:05 EDT 2024
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
from misc import get_kws_from_argv
from windspharm.xarray import VectorWind
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = '/tigress/wenchang/data/era5/raw/monthly/plevels/specific_humidity/era5.specific_humidity.monthly.2020.nc'
q = xr.open_dataarray(ifile).load()

ifile = '/tigress/wenchang/data/era5/raw/monthly/plevels/u_component_of_wind/era5.u_component_of_wind.monthly.2020.nc'
u = xr.open_dataarray(ifile).load()

ifile = '/tigress/wenchang/data/era5/raw/monthly/plevels/v_component_of_wind/era5.v_component_of_wind.monthly.2020.nc'
v = xr.open_dataarray(ifile).load()

qu = q * u
qv = q * v

#vertically integrate
viqu = qu.integrate('level')*100/9.8 #hPa->Pa and divided by g; units of kg m**-1 s**-1
viqv = qv.integrate('level')*100/9.8 #hPa->Pa and divided by g; units of kg m**-1 s**-1

#add attrs
viqu = viqu.assign_attrs(units='kg m**-1 s**-1', long_name='eastward vertically integrated moisture transport')
viqv = viqv.assign_attrs(units='kg m**-1 s**-1', long_name='northward vertically integrated moisture transport')

#divergence
w = VectorWind(viqu, viqv)
vimdiv = w.divergence() * 24 * 3600 #kg m**-2 s**-1 to mm/day
vimdiv = vimdiv.assign_attrs(units='mm/day', long_name='moisture divergence')

ds = xr.Dataset(dict(viqu=viqu, viqv=viqv, vimdiv=vimdiv))
ds = ds.rename(longitude='lon', latitude='lat')
units = viqu.attrs['units']
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    mon = get_kws_from_argv('mon', None)
    mon = 1 if mon is None else int(mon)
    fig,ax = plt.subplots(figsize=(8,3.62))
    ds_ = ds.isel(time=mon-1, lon=slice(None, None, 20), lat=slice(None, None, 20))
    levels = [-2**ii for ii in range(8, 0, -1)] +[0,] + [2**ii for ii in range(1, 8+1)]
    ds_.vimdiv.plot.contourf(levels=levels, cmap='BrBG_r', extend='neither')
    quiver = ds_.plot.quiver(x='lon', y='lat', u='viqu', v='viqv', add_guide=False, scale=10000)
    U = 500
    ax.quiverkey(quiver, 0.7, 1.03, U, f'{U} {units}', labelpos='E')

    mapplot()
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.set_aspect('equal')
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__mon{mon:02d}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()