#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Sep 26 15:33:11 EDT 2024
if __name__ == '__main__':
    import sys,os
    try:
        from misc.timer import Timer
        tt = Timer(f'[{os.getcwd()}] start ' + ' '.join(sys.argv))
    except:
        pass
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 geoplots import mapplot
#
if __name__ == '__main__':
    try:
        tt.check('end import')
    except:
        pass
#
#start from here
def wyplot(ax=None, labels_on=False):
    if ax is None: fig,ax = plt.subplots()
    ax.set_xticks(range(0,361,30))
    ax.set_yticks(range(-30,31,10))
    if labels_on:
        ax.set_facecolor('azure')
        mapplot(fill_continents=True)

    #NS index
    lon0, lon1 = 190, 240
    lat0, lat1 = 5, 10
    color = 'C0'
    #boxes
    ax.plot([lon0, lon1, lon1, lon0, lon0], [lat0, lat0, lat1, lat1, lat0], color=color, ls='-')
    ax.plot([lon0, lon1, lon1, lon0, lon0], [-lat0, -lat0, lat0, lat0, -lat0], color=color, ls='--')
    ax.plot([lon0, lon1, lon1, lon0, lon0], [-lat1, -lat1, -lat0, -lat0, -lat1], color=color, ls='-')
    if labels_on:
        #labels
        ax.text((lon0+lon1)/2, (lat0+lat1)/2, 'N', color=color, ha='center', va='center')
        ax.text((lon0+lon1)/2, 0, 'EQ', color=color, ha='center', va='center')
        ax.text((lon0+lon1)/2, -(lat0+lat1)/2, 'S', color=color, ha='center', va='center')
        ax.text(lon1, -lat1, f'{lon0}-{lon1}', color=color, ha='right', va='bottom', fontsize='small')
        #definition
        ax.text((lon0+lon1)/2, -lat1, '\n NS index\n (SST$_N$ + SST$_S$)/2 $-$ SST$_{EQ}$', color=color, ha='center', va='top')

    #EW index
    lon0, lon1 = 140, 170
    lon2, lon3 = 190, 270 # old ones for testing purpose from Jingyi's notebook: 200, 260
    lat0 = 5
    color = 'C2'
    #boxes
    ax.plot([lon0, lon1, lon1, lon0, lon0], [-lat0, -lat0, lat0, lat0, -lat0], color=color, ls='-')
    ax.plot([lon2, lon3, lon3, lon2, lon2], [-lat0, -lat0, lat0, lat0, -lat0], color=color, ls='--')
    if labels_on:
        #labels
        ax.text((lon0+lon1)/2, 0, 'W', color=color, ha='center', va='center')
        ax.text(lon1, -lat0, f'{lon0}-{lon1}', color=color, ha='right', va='bottom', fontsize='small')
        ax.text((lon2+lon3)/2, 0, 'E', color=color, ha='center', va='center')
        ax.text(lon3, -lat0, f'{lon2}-{lon3}', color=color, ha='right', va='bottom', fontsize='small')
        #definition
        ax.text((lon0+lon1)/2, lat0, 'EW index\n SST$_W$ $-$ SST$_E$\n', color=color, ha='center', va='bottom')

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #fig,ax = plt.subplots()
    #wyplot(labels_on=False, ax=ax)
    wyplot(labels_on=True)

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    try:
        tt.check(f'**Done**')
    except:
        pass
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()