#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Dec 23 16:01:21 EST 2021
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    s = ' '
    tt = Timer(f'start {s.join(sys.argv)}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
da2020 = xr.open_dataarray('sedimentHU_smooth40yr_v20201219.nc')
da2021 = xr.open_dataarray('sedimentHU_smooth40yr_v20211203.nc')
years = slice(850, 2000)
yearsPI = slice(850, 1850)
sel_years = lambda da: da.sel(year=years)
sel_yearsPI = lambda da: da.sel(year=yearsPI)

def do_corrcoef(yearspan):
    da0 = da2020.sel(member='normTCcounts').sel(year=yearspan)
    da1 = da2021.sel(member='sediment_estimate').sel(year=yearspan)
    rg = da1.linregress.on(da0)
    return rg.r.item()

    
    
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import xlinregress
    
    plt.close()
    fig, ax = plt.subplots()
    da = da2020.sel(member=['smoothlower', 'smoothupper']).pipe(sel_years)
    ax.fill_between(da.year, *da.transpose(), color='C0', alpha=0.2)
    da = da2020.sel(member='normTCcounts').pipe(sel_years)
    da.plot(ax=ax, label='v2020')

    da = da2021.sel(member=['lower_estimate', 'upper_estimate']).pipe(sel_years)
    ax.fill_between(da.year, *da.transpose(), color='C1', alpha=0.2)
    da = da2021.sel(member='sediment_estimate').pipe(sel_years)
    da.plot(ax=ax, label='v2021')

    ax.legend(loc='upper left', ncol=2)
    cc850_2000 = do_corrcoef(slice(850, 2000)) #correlation coefficient
    cc850_1850 = do_corrcoef(slice(850, 1850)) #correlation coefficient
    ax.set_title(f'correlation coefficient: {cc850_2000:.2g} (850-2000); {cc850_1850:.2g} (850-1850)')
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()