#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Jun  4 11:50:59 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
import xarray as xr, numpy as np, pandas as pd
import matplotlib.pyplot as plt
#more imports
import xaddon
maindir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if maindir not in sys.path:
    sys.path.append(maindir)
from HiRAM_plus2K_2xCO2.data_ntc_cycle import get_cycle as cycle_ntc
from HiRAM_plus2K_2xCO2.data_pTang_cycle import get_cycle as cycle_p
from HiRAM_plus2K_2xCO2.data_nseed_cycle import get_cycle as cycle_nseed
from HiRAM_plus2K_2xCO2.data_nseedxp_cycle import get_cycle as cycle_nseedxp
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
basin = 'NA'
years = None
life = 12
def scatterplot(x=None, y=None, data=None, ax=None, xlabel=None, ylabel=None, title=None, alpha=0.5, tag=None, **kws):
    if ax is None:
        ax = plt.gca()
    if tag is None:
        tag = ''
    ax.scatter(x=x, y=y, data=data, alpha=0.5, s=20, **kws)
    ax.set_xlim(0, None)
    ax.set_ylim(0, None)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    if title is not None:
        ax.set_title(title, loc='left')

    # linear regression
    xx, yy = data[x], data[y]
    rg = yy.linregress.on(xx)
    ax.plot(xx, rg.predicted)
    # text
    if rg.intercept.item()>=0:
        s = f'{tag}y={rg.slope.item():.1f}x $+$ {rg.intercept.item():.2f}\nr$^2$={rg.r.item()**2:.2f}'
    else:
        s = f'{tag}y={rg.slope.item():.1f}x $-$ {-rg.intercept.item():.2f}\nr$^2$={rg.r.item()**2:.2f}'
    ax.text(0.02, 1-0.02, s, ha='left', va='top', transform=ax.transAxes)
    #print(rg)

def get_scatter_data(basin='NA', years=None):
    if years is None:
        years = slice(None, None)
    ntc = cycle_ntc(basin=basin, years=years)['mclim']
    p = cycle_p(basin=basin, years=years)['mclim']
    nseed = cycle_nseed(basin=basin, years=years)['mclim'].sel(life=life, drop=True)
    nseedxp = cycle_nseedxp(basin=basin, years=years)['mclim'].sel(life=life, drop=True)
    ds = xr.Dataset(dict(
        ntc=ntc, p=p, nseed=nseed, nseedxp=nseedxp
        ))
    return ds
    
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    ds = get_scatter_data(basin=basin, years=years)
    ofile = __file__.replace('.py', '.nc')
    if not os.path.exists(ofile):
        ds.to_netcdf(ofile)
        print('[newly saved]:', ofile)
    else:
        print('[already saved]:', ofile)
    fig, axes = plt.subplots(1, 3, sharey=True, figsize=(8,3))
    suptitle = 'HiRAM_plus2K_2xCO2'
    figname = __file__.replace('.py', f'_{tt.today()}.png')

    ax = axes[0]
    scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='N_TC', x='p', y='ntc', data=ds)

    ax = axes[1]
    scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds)

    ax = axes[2]
    scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds)

    fig.suptitle(suptitle)

    plt.savefig(figname)
    print('[saved]:', figname)

    tt.check(f'**Done**')
    plt.show()