#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Feb  2 22:16:22 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
import xfilter, xlinregress, geoxarray #xlearn (use xlinregress instead)
import cartopy.crs as ccrs, cartopy.feature as cfeature
#from shared import ess_in_corr
#from mystats import p2r
from xpyleoclim.correlation import correlation
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
# params
# data
# tc liz
#ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v0719.txt'
ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v1219.txt'
tc_liz = pd.read_csv(ifile, sep='\s+', index_col=0)
# sst
ifile = '/tigress/gvecchi/DATA/LMR_2019/sst_MCruns_ensemble_mean.nc'
ds_sst = xr.open_dataset(ifile)

year0, year1 = 850, 1850
n_window = 40 
padtype = 'even'
datafile = os.path.join(
    os.path.dirname(__file__), 
    'figdata', 
    os.path.basename(__file__).replace('.py', f'.{year0:04d}-{year1:04d}_lp{n_window}{padtype}.nc')
    )
try: # load dataset for figure if exists
    ds = xr.open_dataset(datafile)
    print('[loaded]:', datafile)
except: # calculate in place if not exists
    #from xpyleoclim.correlation import correlation
    alpha = 0.05
    da = ds_sst['sst'].sel(time=slice(f'{year0:04d}', f'{year1:04d}')).mean('MCrun') \
        .filter.lowpass(1.0/n_window, dim='time', padtype=padtype)
    da_rsst = da.pipe(lambda x: x-x.sel(lat=slice(-30, 30)).geo.fldmean())
    ts = xr.DataArray(tc_liz['normTCcounts'][year0:year1], dims='time', coords=[da.time])
    ts = (ts - ts.mean())/ts.std() #standardized
    #use xlinregress to calculate corrcoef and slope
    reg = da.linregress.on(ts)
    da_r = reg.r
    da_b = reg.slope
    reg_rsst = da_rsst.linregress.on(ts)
    da_r_rsst = reg_rsst.r
    da_b_rsst = reg_rsst.slope
    #use xpyleoclim.correlation to get the significant info
    rr = correlation(da, ts, seed=20211201)
    da_p = rr.p
    da_signif = rr.signif
    da_signif_fdr = rr.signif_fdr
    rr_rsst = correlation(da_rsst, ts, seed=20211201)
    da_p_rsst = rr_rsst.p
    da_signif_rsst = rr_rsst.signif
    da_signif_fdr_rsst = rr_rsst.signif_fdr

    ds = xr.Dataset(dict(r=da_r, signif=da_signif, signif_fdr=da_signif_fdr, p=da_p,
        r_rsst=da_r_rsst, signif_rsst=da_signif_rsst, signif_fdr_rsst=da_signif_fdr_rsst, p_rsst=da_p_rsst,
        b=da_b, b_rsst=da_b_rsst))
    ds.to_netcdf(datafile)
    print('[saved]:', datafile)
    
def plot_map_corr(ax, daname, cbar_kwargs=None, **kws):
    fdr_on = kws.pop('fdr_on', False)
    land_color = 'k'
    cmap = ['turbo','RdBu_r'][-1]
    hatches = ['...']
    hatch_color = 'gray'
    hatch_color_default = plt.rcParams['hatch.color']
    plt.rcParams['hatch.color'] = hatch_color
    center = 0
    if daname.startswith('r'):
        levels = 21#np.arange(-1,1.01,0.1)
        vmax = 1
    elif daname.startswith('b'):
        levels = 25#
        vmax = 0.12
    else:
        print('daname must be from "r", "r_rsst", "b" or "b_rsst"')
    proj = ccrs.Robinson(central_longitude=180)
    dproj = ccrs.PlateCarree()
    im0 = ds[daname].plot.contourf(ax=ax, transform=dproj, add_colorbar=True,
                 cmap=cmap, center=center, vmax=vmax, levels=levels, cbar_kwargs=cbar_kwargs, **kws
                )
    signame = 'signif'
    if fdr_on:
        signame += '_fdr'
    if daname.endswith('_rsst'):
        signame += '_rsst'
    ds[signame].where(ds[signame]==1).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 colors='none', hatches=hatches,
                )
    ax.add_feature(cfeature.LAND, color=land_color)
    ax.set_global()
    if daname == 'r':
        title = 'corr(SST, TC)'
    elif daname == 'r_rsst':
        title = 'corr(rSST, TC)'
    elif daname == 'b':
        title = 'b(SST, TC)'
    elif daname == 'b_rsst':
        title = 'b(rSST, TC)'
    else:
        print('daname must be from "r", "r_rsst", "b" or "b_rsst"')
        return
    ax.set_title(title)
    plt.rcParams['hatch.color'] = hatch_color_default
    

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    # fig
    #proj = ccrs.Robinson(central_longitude=220)
    proj = ccrs.Robinson(central_longitude=180)
    dproj = ccrs.PlateCarree()
    #fig, axes = plt.subplots(2, 2, figsize=(10, 7), subplot_kw=dict(projection=proj))
    #fig, axes = plt.subplots(1, 2, figsize=(7, 2.5), subplot_kw=dict(projection=proj))
    #fig, axes = plt.subplots(2, 2, figsize=(7, 5.4), subplot_kw=dict(projection=proj))
    fig, axes = plt.subplots(2, 2, figsize=(6, 4.8), subplot_kw=dict(projection=proj))
    land_color = 'k'
    cmap = ['turbo','RdBu_r'][-1]
    hatches = ['...']
    hatch_color = 'gray'
    hatch_color_default = plt.rcParams['hatch.color']
    plt.rcParams['hatch.color'] = hatch_color
    t_lia = slice('1400', '1700')
    t_mwp = slice('0950', '1250')
    mwp_minus_lia = lambda x: x.sel(time=t_mwp).mean('time') - x.sel(time=t_lia).mean('time')
    r_sst = lambda x: x-x.sel(lat=slice(-30, 30)).geo.fldmean()
    signame = 'signif'
    signame_rsst = f'{signame}_rsst'
    
    #
    ax = axes[0,0]
    im = ds.r.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=np.arange(-1,1.01,0.1),
                )
    #r.pvalue_da.pipe(lambda x: x.where(x<0.05)*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
    ds[signame].where(ds[signame]==1).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 colors='none', hatches=hatches,
                )
    ax.add_feature(cfeature.LAND, color=land_color)
    ax.set_global()
    ax.set_title('corr(SST, TC)')
    
    # 
    ax = axes[0,1]
    im0 = ds.r_rsst.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=np.arange(-1,1.01,0.1),
                )
    ds[signame_rsst].where(ds[signame_rsst]==1).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 colors='none', hatches=hatches,
                )
    ax.add_feature(cfeature.LAND, color=land_color)
    ax.set_global()
    ax.set_title('corr(rSST, TC)')
    
    ax = axes[1,0]
    im10 = ds.b.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=np.arange(-0.12,0.121,0.02),
                )
    #r.pvalue_da.pipe(lambda x: x.where(x<0.05)*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
    ds[signame].where(ds[signame]==1).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 colors='none', hatches=hatches,
                )
    ax.add_feature(cfeature.LAND, color=land_color)
    ax.set_global()
    ax.set_title('b(SST, TC)')
    
    ax = axes[1,1]
    im11 = ds.b_rsst.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=np.arange(-0.12,0.121,0.02),
                )
    ds[signame_rsst].where(ds[signame_rsst]==1).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 colors='none', hatches=hatches,
                )
    ax.add_feature(cfeature.LAND, color=land_color)
    ax.set_global()
    ax.set_title('b(rSST, TC)')
    
    
    pad = 0.05
    #plt.colorbar(im0, ax=[axes[0,0], axes[0,1]], pad=pad)
    #plt.colorbar(im0, ax=[axes[0], axes[1]], aspect=10, shrink=0.6, pad=0.02, orientation='horizontal')
    plt.colorbar(im0, ax=[axes[0,0], axes[0,1]], aspect=30, shrink=0.6, pad=pad, orientation='horizontal')
    cbar = plt.colorbar(im11, ax=[axes[1,0], axes[1,1]], aspect=30, shrink=0.6, pad=pad, orientation='horizontal')
    cbar.set_label('K')
    """
    cbar = plt.colorbar(im1, ax=[axes[1,0], axes[1,1]], pad=pad)
    cbar.set_label('K')
    cbar = plt.colorbar(im2, ax=[axes[2,0], axes[2,1]], pad=pad)
    cbar.set_label('K')
    """
    
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname, dpi=128)
    else:
        print(f'figure not saved. to save figure: python {__file__} savefig')
    tt.check(f'**Done**')
    plt.show()
    plt.rcParams['hatch.color'] = hatch_color_default