#!/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__)))
maindir = '/tigress/wenchang/analysis/seedTC'
if maindir not in sys.path:
    sys.path.append(maindir)
from data_basin_areas import da as basin_areas
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
years = slice('1980', '2018')
life = 12
def scatterplot(x=None, y=None, data=None, ax=None,
    xlabel=None, ylabel=None, title=None, alpha=0.5,
    tag=None, scatter_on=True, rg_on=True, **kws):
    if ax is None:
        ax = plt.gca()
    if tag is None:
        tag = ''
    if scatter_on:
        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
    if rg_on:
        xx, yy = data[x], data[y]
        rg = yy.linregress.on(xx)
        ax.plot(xx, rg.predicted, color='k', alpha=.5)
        # 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)
if True: 
    #ds = get_scatter_data(basin=basin, years=years)
    basins = ['NA', 'EP', 'WP', 'NI', 'SI', 'AU', 'SP']
    basin_area_on = [False, True][1]
    basin_removed = None#basins[4] #'EP'#'NA', remove one basin to see the sensitivity
    dss = []
    for basin in basins:
        if basin == 'NA':
            ifile = os.path.join(maindir, 'AM2p5', 'fig_scatter_ntc_vs_all.nc')
        else:
            ifile = os.path.join(maindir, 'AM2p5', basin, 'fig_scatter_ntc_vs_all.nc')
        ds = xr.open_dataset(ifile)
        dss.append(ds)
    ds = xr.concat(dss, dim=pd.Index(basins, name='basin'))

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    if basin_area_on: #p weighted by basin area, NA area is 1
        ds['p'] = ds['p']*(basin_areas/basin_areas.sel(basin='NA')) # normalized by basin areas, NA area is 1
        #ds['nseedxp'] = ds['nseedxp']*(basin_areas/basin_areas.sel(basin='NA')) # normalized by basin areas, NA area is 1
    if basin_removed is not None:
        basins.remove(basin_removed)
    ds = ds.sel(basin=basins)
    ds = ds.stack(s=['basin', 'month'])
            
    fig, axes = plt.subplots(1, 3, sharey=True, figsize=(8,3))
    suptitle = f'AM2.5 {basins}'
    figname = __file__.replace('.py', f'_{tt.today()}.png')
    if basin_removed is not None:
        figname = figname.replace(f'_{tt.today()}', f'_no{basin_removed}_{tt.today()}')
    if basin_area_on:
        figname = figname.replace(f'_{tt.today()}', f'_areaWeighted_{tt.today()}')
    #figname = __file__.replace('.py', f'_EPWPSI_{tt.today()}.png')

    ax = axes[0]
    scatterplot(ax=ax, xlabel='p($\Lambda$)', ylabel='N_TC', x='p', y='ntc', data=ds, color='.99', label=None)
    for basin in basins:
        scatterplot(ax=ax, x='p', y='ntc', data=ds.sel(basin=basin), rg_on=False, label=basin)

    ax = axes[1]
    scatterplot(ax=ax, xlabel='N_SEED', x='nseed', y='ntc', data=ds, color='.99', label=None)
    for basin in basins:
        scatterplot(ax=ax, x='nseed', y='ntc', data=ds.sel(basin=basin), rg_on=False, label=basin)

    ax = axes[2]
    scatterplot(ax=ax, xlabel='N_SEED$\\times$p($\Lambda$)', x='nseedxp', y='ntc', data=ds, color='.99', label=None)
    for basin in basins:
        scatterplot(ax=ax, x='nseedxp', y='ntc', data=ds.sel(basin=basin), rg_on=False, label=basin)

    fig.suptitle(suptitle)
    axes[2].legend(loc='lower right', frameon=False)

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

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