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

qclim_on = True if 'qclim' in sys.argv else False
uvclim_on = True if 'uvclim' in sys.argv else False
qanom_on = True if 'qanom' in sys.argv else False
uvanom_on = True if 'uvanom' in sys.argv else False
ifile = '/tigress/wenchang/data/era5/analysis_wy/realtime202405/data_era5_moisture_transport.nc'
if qclim_on:
    ifile = ifile.replace('.nc', '_qclim.nc')
elif qanom_on:
    ifile = ifile.replace('.nc', '_qanom.nc')
if uvclim_on:
    ifile = ifile.replace('.nc', '_uvclim.nc')
elif uvanom_on:
    ifile = ifile.replace('.nc', '_uvanom.nc')

if qclim_on and uvclim_on:
    ds = xr.open_dataset(ifile).load()
else:
    ds = xr.open_dataset(ifile).sel(time=timeslice).load()


if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    from misc import get_kws_from_argv
    #subset of ds
    if qclim_on and uvclim_on:
        ds_ = ds.pipe(roll_lon).sel(lon=lonslice, lat=latslice)
    else:
        ds_ = ds.pipe(roll_lon).sel(lon=lonslice, lat=latslice).mean('time', keep_attrs=True)

    fig,ax = plt.subplots(figsize=(7,5))
    units = ds.viqu.attrs['units']
    #contouf
    levels = [-2**ii for ii in range(8, 0, -1)] +[0,] + [2**ii for ii in range(1, 8+1)]
    ds_['vimdiv'].plot.contourf(ax=ax, levels=levels, cmap='BrBG_r', cbar_kwargs={'shrink': 0.68}, extend='neither')
    #quiver
    islice = slice(0, None, 12)
    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', lw=1)
    ax.set_xlabel('')
    ax.set_ylabel('')
    U = 500
    ax.quiverkey(q, 0.72, 1.03, U, f'{U} {units}', labelpos='E', labelcolor='k')
    title = f'ERA5 moisture transport and divergence {timeslice.start} to {timeslice.stop}\n'
    if qclim_on:
        title = title + ' qclim'
    elif qanom_on:
        title = title + ' qanom'
    if uvclim_on:
        title = title + ' uvclim'
    elif uvanom_on:
        title = title + ' uvanom'
    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 qclim_on: 
            figname = figname.replace('.png', '__qclim.png')
        elif qanom_on: 
            figname = figname.replace('.png', '__qanom.png')
        if uvclim_on:
            figname = figname.replace('.png', '__uvclim.png')
        elif uvanom_on:
            figname = figname.replace('.png', '__uvanom.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()