#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed May 15 21:47:35 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 misc import roll_lon
from windspharm.xarray import VectorWind
import salem
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#shapefile based subset
shapefile = 'state_Rio_Grande_do_Sul/43UF2500G.shp'
shdf = salem.read_shapefile(shapefile)

timeslice = slice('2024-04-26', '2024-05-05')
lonslice = slice(-100, 20)
latslice = slice(20, -60)
rename = lambda x: x.rename(longitude='lon', latitude='lat')
subset = lambda x: x.sel(time=timeslice)
def getdata(daname_long):
    ifiles = f'/tigress/wenchang/data/era5/raw/{daname_long}/realtime202405/era5.{daname_long}.2024-0?.nc'
    ds = xr.open_mfdataset(ifiles).pipe(rename).sel(time=timeslice).load() \
        .resample(time='1D').mean(keep_attrs=True)

    #dataset to dataarray
    daname = list(ds.data_vars)[0]
    da = ds[daname]
    da = roll_lon(da)
    return da

daname_long = 'vertical_integral_of_eastward_water_vapour_flux'
viqu = getdata(daname_long)
daname_long = 'vertical_integral_of_northward_water_vapour_flux'
viqv = getdata(daname_long)
daname_long = 'mean_vertically_integrated_moisture_divergence'
vidiv = getdata(daname_long)
vidiv = vidiv * 24 *3600 #units convert from kg m**-2 s**-1 to mm/day
vidiv = vidiv.assign_attrs(units='mm/day', long_name=daname_long.replace('_', ' '))
daname_long = 'mean_total_precipitation_rate'
precip = getdata(daname_long)
precip = precip * 24 *3600 #units convert from kg m**-2 s**-1 to mm/day
precip = precip.assign_attrs(units='mm/day', long_name=daname_long.replace('_', ' '))
#calculated divergence (to compare to the provided divergence vidiv)
w = VectorWind(viqu, viqv)
vidiv_cal = w.divergence() * 24 * 3600 # to have the same units mm/day
vidiv_cal = vidiv_cal.assign_attrs(units='mm/day', long_name='cal '+vidiv.attrs['long_name'])

#wrap into dataset
ds = xr.Dataset(dict(
    viqu=viqu, 
    viqv=viqv, 
    vidiv=vidiv,
    precip=precip,
    vidiv_cal=vidiv_cal,
    ))
 

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    from misc import get_kws_from_argv
    #subset of ds
    ds_sub = ds.sel(lat=latslice, lon=lonslice)

    date = get_kws_from_argv('date', None)
    precip_on = True if 'precip' in sys.argv else False
    vidiv_cal_on = True if 'vidiv_cal' in sys.argv else False
    #plt.close()
    fig,ax = plt.subplots(figsize=(7,5))
    units = viqu.attrs['units']
    islice = slice(0, None, 12)
    ds_ = ds_sub.mean('time', keep_attrs=True) if date is None else ds_sub.sel(time=date)
    levels = [-2**ii for ii in range(8, 0, -1)] +[0,] + [2**ii for ii in range(1, 8+1)]
    if vidiv_cal_on:
        ds_['vidiv_cal'].plot.contourf(ax=ax, levels=levels, cmap='BrBG_r', cbar_kwargs={'shrink': 0.68}, extend='neither')
    elif precip_on:
        ds_['precip'].plot.contourf(ax=ax, levels=levels, cmap='BrBG', cbar_kwargs={'shrink': 0.68}, extend='neither')
    else:
        ds_['vidiv'].plot.contourf(ax=ax, levels=levels, cmap='BrBG_r', cbar_kwargs={'shrink': 0.68}, extend='neither')
    ds_ = ds_.isel(lon=islice, lat=islice)
    q = ds_.plot.quiver(x='lon', y='lat', u='viqu', v='viqv', scale=10000, ax=ax, add_guide=False, color='k')

    mapplot()
    shdf.boundary.plot(ax=ax, color='C3')
    ax.set_xlabel('')
    ax.set_ylabel('')
    U = 1000
    ax.quiverkey(q, 0.7, 1.03, U, f'{U} {units}', labelpos='E', labelcolor='k')
    title = f'ERA5 {timeslice.start} to {timeslice.stop}' if date is None else f'ERA5 {date}'
    ax.set_title(title)
    ax.set_aspect('equal')


    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if date is not None: figname = figname.replace(os.path.basename(figname), 'frames/'+os.path.basename(figname)).replace('.png', f'__{date}.png')
        if vidiv_cal_on: 
            figname = figname.replace('.png', '__vidivCal.png')
        elif precip_on:
            figname = figname.replace('.png', '__precip.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()