#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jan  4 10:21:22 EST 2023
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('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
from geoplots.cartopy.api import cartoproj, cartofeature
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
lmr = 'seasonal'
ifile = '/projects/w/wenchang/data/LMR/seasonal/tos_mean.nc'
years = slice('1451', '1550')
years_ref = slice('1351', '1450')
#years = slice('1351', '1450')
#years_ref = slice('1251', '1350')
long_name=f'LMR{lmr} SST change from {years.start}-{years.stop} to {years_ref.start}-{years_ref.stop}'

#ofile after anomaly and ens/time mean
ofile = f'LMR{lmr}_ssta_{years.start}-{years.stop}_ref{years_ref.start}-{years_ref.stop}.nc'
if os.path.exists(ofile) and 'od' not in sys.argv:
    print('[exists]:', ofile)
    daa = xr.open_dataarray(ofile)
else:
    da = xr.open_dataarray(ifile).load()
    daa = da.sel(time=years).mean('time') - da.sel(time=years_ref).mean('time')
    daa = daa.assign_attrs(units='K', long_name=long_name)
    daa.to_dataset(name='sst').to_netcdf(ofile)
    print('[saved]:', ofile)

def wyplot(ax=None, **kws):
    dproj = cartoproj('cyl')
    proj = cartoproj('robin', central_longitude=210)
    if ax is None:
        fig,ax = plt.subplots(subplot_kw=dict(projection=proj), figsize=(6,3))
    da = daa
    da0 = da.isel(lon=slice(0,1))
    da0 = da0.assign_coords(lon=da0.lon.values+360)
    da = xr.concat([da,da0], dim='lon')
    da.assign_attrs(long_name='SST change').plot.contourf(ax=ax, transform=dproj, levels=21, **kws)
    #ds.pvalue.where(ds.pvalue<0.05).plot.contourf(ax=ax, transform=dproj, colors='none', hatches=['...'], add_colorbar=False)
    #mapplot(fill_continents=True)
    ax.add_feature(cartofeature('land'), color='k')
    ax.set_global()

    #MDR
    mdr_color = 'C0'
    latmin, latmax, lonmin, lonmax = 10, 25, 280, 340
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], transform=dproj, color=mdr_color, lw=1)
    lat0, lon0 = (latmin+latmax)/2, (lonmin+lonmax)/2
    ax.text(lon0, lat0, 'MDR', ha='center', va='center', transform=dproj, color=mdr_color)
    #trop
    latmin, latmax, lonmin, lonmax = -30, 30, 0+30, 360+30
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], transform=dproj, color='gray', ls='--', lw=1)
    #nino34
    nino34_color = 'C2'
    latmin, latmax, lonmin, lonmax = -5, 5, 360-170, 360-120
    ax.plot([lonmin, lonmax, lonmax, lonmin, lonmin], [latmin, latmin, latmax, latmax, latmin], transform=dproj, color=nino34_color, lw=1)
    lat0, lon0 = (latmin+latmax)/2, (lonmin+lonmax)/2
    ax.text(lon0, lat0, 'Nino3.4', ha='center', va='center', transform=dproj, color=nino34_color)

    #ax = plt.gca()
    ax.set_title(long_name, loc='center')
    ax.set_xlabel('')
    ax.set_ylabel('')

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    #from geoplots import mapplot
    #daa.plot.contourf(levels=21)
    #mapplot(fill_continents=True)
    wyplot()
    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        tag = long_name.replace(' ', '_')
        figname = __file__.replace('.py', f'__{tag}.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:
        pass
    else:
        plt.show()