#!/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
import xlinregress
from geoplots.cartopy.api import cartoproj, cartofeature
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
lmr = '2018' if '2018' in sys.argv else '2019'
lmr_version = '2.0' if lmr == '2018' else '2.1'
ifile = f'/tigress/gvecchi/DATA/LMR_{lmr}/sst_MCruns_ensemble_mean.nc'
years = slice('1351', '1550')
long_name=f'LMR{lmr_version} SST linear trend over {years.start}-{years.stop}'

#ofile after anomaly and ens/time mean
ofile = f'LMR{lmr_version}_sst_lintrend{years.start}-{years.stop}.nc'
if os.path.exists(ofile) and 'od' not in sys.argv:
    print('[exists]:', ofile)
    ds = xr.open_dataset(ofile)
else:
    da = xr.open_dataarray(ifile).load()
    da = da.sel(time=years).groupby('time.year').mean('time').mean('MCrun')
    rg = da.linregress.on(da.year, dim='year')
    rg['slope'] = (rg['slope'] * 100).assign_attrs(units='K/century', long_name='SST linear trend')
    rg.to_netcdf(ofile)
    print('[saved]:', ofile)
    ds = rg

def wyplot(ax=None, **kws):
    #from geoplots import mapplot
    dproj = cartoproj('cyl')
    if ax is None:
        proj = cartoproj('robin', central_longitude=210)
        fig,ax = plt.subplots(subplot_kw=dict(projection=proj), figsize=(6,3))
    da = ds['slope']
    da0 = da.isel(lon=slice(0,1))
    da0 = da0.assign_coords(lon=da0.lon.values+360)
    da = xr.concat([da,da0], dim='lon')
    da.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
    wyplot()

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__LMR{lmr_version}.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()