#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu May 27 14:57:22 EDT 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, matplotlib.pyplot as plt
#more imports
import xlinregress#,xaddon
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'ssttc.histrcp85.1850-2100.ens39.nc'
ds = xr.open_dataset(ifile)
da = ds.HU
pdyears = slice(1861,2005)
futureyears= slice(2006,2100)
da_ = da.sel(year=pdyears).pipe(np.log)
lintrend_pd = da_.linregress.numba(da_.year, dim='year')
slope_pd = lintrend_pd.slope*100*100#fraction per year -> % per century
slope_pd = slope_pd.assign_attrs(units='% per century', long_name=f'linear trend over {pdyears.start}-{pdyears.stop}')
da_ = da.sel(year=futureyears).pipe(np.log)
lintrend_future = da_.linregress.numba(da_.year, dim='year')
slope_future = lintrend_future.slope*100*100#fraction per year -> % per century
slope_future = slope_future.assign_attrs(units='% per century', long_name=f'linear trend over {futureyears.start}-{futureyears.stop}')
ds = xr.Dataset(dict(pd=slope_pd, future=slope_future))

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,ax = plt.subplots(figsize=(6,4.5))
    ds.plot.scatter(x='pd', y='future', color='none', edgecolors='C0', ax=ax)
    #for ii,model in enumerate(ds.model.values, start=1):
    #    ax.text(ds.pd.sel(model=model), ds.future.sel(model=model), ' ' + str(ii), ha='left', va='center')

    #regression
    rg = ds.future.linregress.numba(ds.pd)
    x = ds.pd.values
    y = (rg.slope*ds.pd + rg.intercept).values
    dy = ( rg.predict_stderr * rg.t_alpha ).values
    ci = ( (y-dy)[np.argsort(x)], (y+dy)[np.argsort(x)]) 
    ax.fill_between(np.sort(x), ci[0], ci[1], alpha=1/4)
    #ax.plot(ds.pd, rg.slope*ds.pd + rg.intercept)
    ax.axline([x.min(), y.min()], [x.max(), y.max()], ls='--')

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()