#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jul 12 15:44:51 EDT 2023
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.landmask import whereland
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
season = """
annual
JJA""".split()[-1]
lmr = """
2018
2019
""".split()[-1]
ifile = f'precip_AM2.1_amipLMR{lmr}SST0850_ens_tigercpu_intelmpi_18_30PE_5ens_1551-1650_{season}mean.nc'
da = xr.open_dataarray(ifile)
da_0850 = da

ifile = ifile.replace('0850', '1551')
da = xr.open_dataarray(ifile)
da = da.assign_coords(ens=range(6,11))
da_1551 = da

da = xr.concat([da_0850, da_1551], dim='ens')
da = da.sel(lon=slice(95, 140), lat=slice(25, 50))
da = da - da.mean('time').mean('ens')
da = da.pipe(whereland)

def plot_lonlatbox(region='North China', ax=None, **kws):
    if ax is None: fig, ax = plt.subplots()
    if region == 'North China':
        lonlatbox = 106, 120, 33, 40
    elif region == 'Northeast China':
        lonlatbox = 122, 130, 40, 46
    elif region == 'Mongolia':
        lonlatbox = 107, 122, 40, 47
    lonmin, lonmax, latmin, latmax = lonlatbox
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], ls='--', **kws)
    ax.text((lonmin+lonmax)/2, (latmin+latmax)/2, region, ha='center', va='center', color='gray')
    
def wyplot(ax=None, tspan=None): 
    from geoplots import mapplot
    if ax is None: fig, ax = plt.subplots()
    if tspan is None: tspan = slice('1621', '1640')
    daa = da.sel(time=tspan).mean('time').mean('ens').assign_attrs(units='mm/day', long_name='precip anom')
    if season in ('JJA',):
        vmin, vmax = -0.3, 0.3
        levels = 21
    elif season in ('annual',):
        vmin, vmax = -0.15, 0.15
        levels = 31
    daa.plot(cmap='BrBG', levels=levels, ax=ax, vmin=vmin, vmax=vmax)
    mapplot()
    ax.set_prop_cycle(None)
    plot_lonlatbox('North China', ax=ax)
    plot_lonlatbox('Northeast China', ax=ax)
    plot_lonlatbox('Mongolia', ax=ax)

    ax.set_title(f'{season} precip anomaly during {tspan.start}-{tspan.stop}, LMR{lmr}SST')
    ax.set_xlabel('')
    ax.set_ylabel('')
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #tspan = slice('1621', '1640')
    #tspan = slice('1624', '1643')
    tspan = slice('1625', '1644')
    #tspan = slice('1615', '1644')
    #tspan = slice('1595', '1644')
    wyplot(tspan=tspan)
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__{tspan.start}-{tspan.stop}{season}_LMR{lmr}SST.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()