#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Feb  8 22:29:10 EST 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd
import matplotlib.pyplot as plt
#more imports
from geoplots.cartopy.api import cartoplot
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#data
dc = dict()
"""
#LMR_2019
ifile = '/tigress/gvecchi/DATA/LMR_2019/sst_MCruns_ensemble_mean.nc' print(ifile, 'loading...')
da = xr.open_dataarray(ifile).load()
da = da.mean('MCrun') \
    .assign_coords(time=da.time.dt.year.values) \
    .rename(time='year')
dc['LMR_2019'] = da
#LME_fullForcing
ifile = 'data/b.e11.BLMTRC5CN.f19_g16.001-013.cam.h0.TS.0850-2005.nc'
print(ifile, 'loading...')
da = xr.open_dataarray(ifile).load()
da = da.mean('en')
dc['LME'] = da
"""
#LME_OA
ifile = 'data/b.e11.BLMTRC5CN.f19_g16.OZONE_AER.001-004.cam.h0.TS.0850-2005.nc'
print(ifile, 'loading...')
da = xr.open_dataarray(ifile).load()
da = da.mean('en')
dc['LME_OA'] = da

def plot_sst_anom(inputdata, yearc, n_window=40, yearsRef=None, rSST=False, **kws):
    """inputdata: 'LMR_2019', 'LME_fullForcing', or 'LME_GHG'.
       yearc: year center """
    if yearsRef is None:
        yearsRef = slice(850,1850)
    da = dc[inputdata]
    da_ = da.sel(year=slice(yearc-n_window//2, yearc+n_window//2)).mean('year') - da.sel(year=yearsRef).mean('year')
    long_name = f'SST anom'
    if rSST:
        da_ = da_ - da_.sel(lat=slice(-30,30)).geo.fldmean()
        long_name = long_name.replace('SST', 'rSST')
    da_.assign_attrs(units='K', long_name=long_name) \
        .plot.contourf(cmap='RdBu_r', **kws)
    plt.title(f'{inputdata} {yearc-n_window//2}-{yearc+n_window//2} minus {yearsRef.start}-{yearsRef.stop}')
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import cartopy.crs as ccrs, cartopy.feature as cfeature
    inputdata = 'LME_OA'
    proj = ccrs.Robinson(central_longitude=180)
    dproj = ccrs.PlateCarree()
    land_color = 'k'
    figsize = (9,2.1)
    fig, axes = plt.subplots(1, 2, figsize=figsize, subplot_kw=dict(projection=proj))
    """
    ax = axes[0,0]
    plt.sca(ax)
    plot_sst_anom(inputdata, 1277, ax=ax, transform=dproj, vmax=0.32)
    ax.add_feature(cfeature.LAND, color=land_color)
    ax = axes[0,1]
    plt.sca(ax)
    plot_sst_anom(inputdata, 1277, rSST=True, ax=ax, transform=dproj, vmax=0.32)
    ax.add_feature(cfeature.LAND, color=land_color)
    """
    yearsRef = slice(1850,1890)
    ax = axes[0]
    plt.sca(ax)
    plot_sst_anom(inputdata, 1980, yearsRef=yearsRef, ax=ax, transform=dproj, 
        robust=True, vmax=1, center=0, levels=21, cbar_kwargs=dict(aspect=80))
    ax.add_feature(cfeature.LAND, color=land_color)
    ax = axes[1]
    plt.sca(ax)
    plot_sst_anom(inputdata, 1980, rSST=True, yearsRef=yearsRef, ax=ax, transform=dproj, 
        robust=True, vmax=1, center=0, levels=21, cbar_kwargs=dict(aspect=80))
    ax.add_feature(cfeature.LAND, color=land_color)
    
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()