#!/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, xlearn, geoxarray
import cartopy.crs as ccrs, cartopy.feature as cfeature
from shared import ess_in_corr
from mystats import p2r
#
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)
# FLOR tc
ifile = '/tigress/wenchang/analysis/TC/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.TS17.0001-2000.yearly.nc'
tc = xr.open_dataset(ifile)['NA'].load()
# sst
#ifile = '/tigress/gvecchi/DATA/LMR_2019/sst_MCruns_ensemble_mean.nc'
#ds_sst = xr.open_dataset(ifile)
#FLOR ts
ifile = '/tigress/wenchang/MODEL_OUT/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/analysis/sst.0001-2000.yearly.nc'
ds_sst = xr.open_dataset(ifile)

year0, year1 = 201, 2000
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)
    if 'b' not in ds.data_vars:
        da = ds_sst['sst'] \
            .filter.lowpass(1.0/n_window, dim='year', padtype=padtype) \
            .sel(year=slice(year0, year1))
        da_rsst = da.pipe(lambda x: x-x.sel(lat=slice(-30, 30)).geo.fldmean())
        ts = tc.filter.lowpass(1/n_window, dim='year', padtype=padtype) \
            .sel(year=slice(year0, year1))
        da_b = da.learn.regress(ts, normalize_x=True).coef_da
        da_b_rsst = da_rsst.learn.regress(ts, normalize_x=True).coef_da
        #append to data file
        ds.close()
        ds['b'] = da_b
        ds['b_rsst'] = da_b_rsst
        ds[['b', 'b_rsst']].to_netcdf(datafile, mode='a')
        print('[appended]:', 'b and b_rsst to', datafile)
except: # calculate in place if not exists
    alpha = 0.05
    da = ds_sst['sst'] \
        .filter.lowpass(1.0/n_window, dim='year', padtype=padtype) \
        .sel(year=slice(year0, year1))
    da_rsst = da.pipe(lambda x: x-x.sel(lat=slice(-30, 30)).geo.fldmean())
    ts = tc.filter.lowpass(1/n_window, dim='year', padtype=padtype) \
        .sel(year=slice(year0, year1))
    da_r = da.learn.regress(ts, normalize_xy=True).coef_da
    da_r_rsst = da_rsst.learn.regress(ts, normalize_xy=True).coef_da
    da_b = da.learn.regress(ts, normalize_x=True).coef_da
    da_b_rsst = da_rsst.learn.regress(ts, normalize_x=True).coef_da
    nlat, nlon = da.lat.size, da.lon.size
    zz = np.zeros((nlat, nlon)) # corr of TC and SST 
    zz_n = np.zeros_like(zz) # effective sample size
    zz_rsst = np.zeros_like(zz) # corr of TC and rSST
    zz_n_rsst = np.zeros_like(zz) # effective sample size
    if __name__ == '__main__':
        tt.check('start calculating ...')
    for ilat in range(nlat):
        for ilon in range(nlon):
            n = ess_in_corr(da.isel(lat=ilat, lon=ilon).values, ts.values) #effective sampel size for the correlation
            zz[ilat, ilon] = p2r(alpha, n-2)
            zz_n[ilat, ilon] = n
            #rsst
            n = ess_in_corr(da_rsst.isel(lat=ilat, lon=ilon).values, ts.values) #effective sampel size for the correlation
            zz_rsst[ilat, ilon] = p2r(alpha, n-2)
            zz_n_rsst[ilat, ilon] = n
    da_rsig = xr.DataArray(zz, 
        dims=('lat', 'lon'), 
        coords=[da.lat.values, da.lon.values], 
        attrs=dict(long_name=f'correlation coefficient threshold of significance at {alpha} level'))
    da_n = xr.DataArray(zz_n, 
        dims=('lat', 'lon'), 
        coords=[da.lat.values, da.lon.values], 
        attrs=dict(long_name=f'effective sample size'))
    da_rsig_rsst = xr.DataArray(zz_rsst, 
        dims=('lat', 'lon'), 
        coords=[da.lat.values, da.lon.values], 
        attrs=dict(long_name=f'correlation coefficient threshold of significance at {alpha} level for rSST'))
    da_n_rsst = xr.DataArray(zz_n_rsst, 
        dims=('lat', 'lon'), 
        coords=[da.lat.values, da.lon.values], 
        attrs=dict(long_name=f'effective sample size for the rSST case'))
    if __name__ == '__main__':
        tt.check('end calculating')
    ds = xr.Dataset(dict(r=da_r, r_sig=da_rsig, n=da_n,
        r_rsst=da_r_rsst, r_sig_rsst=da_rsig_rsst, n_rsst=da_n_rsst,
        b=da_b, b_rsst=da_b_rsst))
    ds.to_netcdf(datafile)
    print('[saved]:', datafile)
    

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), 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
    #levels = np.arange(-0.6, 0.61, 0.1)
    levels = np.arange(-1, 1.01, 0.1)
    
    #
    ax = axes[0,0]
    im = ds.r.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=levels,
                )
    #r.pvalue_da.pipe(lambda x: x.where(x<0.05)*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
    ds.r.where(np.abs(ds.r)>=ds.r_sig).pipe(lambda x: x*0).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=levels,
                )
    ds.r_rsst.where(np.abs(ds.r_rsst)>=ds.r_sig_rsst).pipe(lambda x: x*0).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)')
    
    levels = np.arange(-0.12, 0.121, 0.02)
    ax = axes[1,0]
    im10 = ds.b.plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
                 cmap=cmap, levels=levels,
                )
    #r.pvalue_da.pipe(lambda x: x.where(x<0.05)*0).plot.contourf(ax=ax, transform=dproj, add_colorbar=False,
    ds.r.where(np.abs(ds.r)>=ds.r_sig).pipe(lambda x: x*0).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=levels,
                )
    ds.r_rsst.where(np.abs(ds.r_rsst)>=ds.r_sig_rsst).pipe(lambda x: x*0).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.02
    #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=10, shrink=0.6, pad=0.02, orientation='horizontal')
    cbar = plt.colorbar(im11, ax=[axes[1,0], axes[1,1]], aspect=10, shrink=0.6, pad=0.02, 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'_{tt.today()}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        plt.savefig(figname, dpi=128)
        print('[saved]:', figname)
    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