#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Feb  7 23:17:14 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 xaddon, xfilter
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figsize = (8, 10)
    fig, axes = plt.subplots(4, 1, figsize=figsize, sharex=True)
    ax = axes[0]
    plt.sca(ax)
    #data
    ifile = '/tigress/gvecchi/ANALYSIS/LMR_Paper1/normTCcounts_v1219.txt'
    tc_liz = pd.read_csv(ifile, sep='\s+', index_col=0)
    da = xr.Dataset(tc_liz)['normTCcounts'].rename(years='year').sel(year=slice(850, 2000))
    #fig
    da.plot()
    plt.axhline(da.sel(year=slice(850,1850)).mean(), color='gray', ls='--')
    peaks = da.go.find_peaks()
    for year,peak in zip(peaks.year.values, peaks.peaks.values):
        plt.text(year, peak, f'{year:.0f}', va='bottom', ha='center', color='C1')
    peaks = da.pipe(lambda x: -x).go.find_peaks()
    for year,peak in zip(peaks.year.values, -peaks.peaks.values):
        plt.text(year, peak, f'{year:.0f}', va='top', ha='center', color='C2')

    ax = axes[1]
    plt.sca(ax)
    #data
    ifile = os.path.join(os.path.dirname(__file__), 'data', 'LMR2019_TC.nc')
    lmr = ifile.split('/')[-1].split('_')[0]
    ds = xr.open_dataset(ifile)
    da = ds.HU.filter.lowpass(1/40, dim='year', padtype='even').mean('MCrun').sel(year=slice(850, 2000))
    #fig
    da.plot()
    plt.axhline(da.sel(year=slice(850,1850)).mean(), color='gray', ls='--')
    peaks = da.go.find_peaks()
    for year,peak in zip(peaks.year.values, peaks.peaks.values):
        if int(year) == 1686:
            plt.text(year, peak, f'  {year:.0f}', va='bottom', ha='left', color='C1', rotation=90)
        else:
            plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=45)
    peaks = da.pipe(lambda x: -x).go.find_peaks()
    for year,peak in zip(peaks.year.values, -peaks.peaks.values):
        if int(year) == 996:
            plt.text(year, peak, f'{year:.0f}', va='top', ha='left', color='C2', rotation=-90)
        elif int(year) == 1838:
            plt.text(year, peak, f'\n{year:.0f}', va='top', ha='left', color='C2', rotation=-90)
        else:
            plt.text(year, peak, f'{year:.0f}', va='top', ha='left', color='C2', rotation=-45)
    plt.ylabel(f'{lmr} HU')

    ax = axes[2]
    plt.sca(ax)
    #LMR2018 TC peaks
    import xaddon, xfilter
    #data
    ifile = os.path.join(os.path.dirname(__file__), 'data', 'LMR2018_TC.nc')
    lmr = ifile.split('/')[-1].split('_')[0]
    ds = xr.open_dataset(ifile)
    da = ds.HU.filter.lowpass(1/40, dim='year', padtype='even').mean('MCrun').sel(year=slice(850, 2000))
    #fig
    da.plot()
    plt.axhline(da.sel(year=slice(850,1850)).mean(), color='gray', ls='--')
    peaks = da.go.find_peaks()
    for year,peak in zip(peaks.year.values, peaks.peaks.values):
        plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=45)
    peaks = da.pipe(lambda x: -x).go.find_peaks()
    for year,peak in zip(peaks.year.values, -peaks.peaks.values):
        plt.text(year, peak, f'{year:.0f}', va='top', ha='left', color='C2', rotation=-45)
    plt.ylabel(f'{lmr} HU')

    ax = axes[3]
    plt.sca(ax)
    #data
    ifile = os.path.join(os.path.dirname(__file__), 'data', 'LME_TC.fullForcing.13ens.fullForcingRef.nc')
    case = 'LME fullForcing'
    ds = xr.open_dataset(ifile)
    da = ds.HU.filter.lowpass(1/40, dim='year', padtype='even').mean('en').sel(year=slice(850, 2000))
    #fig
    da.plot()
    plt.axhline(da.sel(year=slice(850,1850)).mean(), color='gray', ls='--')
    peaks = da.go.find_peaks()
    for year,peak in zip(peaks.year.values, peaks.peaks.values):
        if int(year) == 1096:
            plt.text(year, peak, f'     {year:.0f}', va='bottom', ha='left', color='C1', rotation=90)
        elif int(year) == 1109:
            plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=75)
        elif int(year) == 1368:
            plt.text(year, peak, f'    {year:.0f}', va='bottom', ha='left', color='C1', rotation=60)
        elif int(year) == 1595:
            plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=75)
        elif int(year) == 1678:
            plt.text(year, peak, f'    {year:.0f}', va='bottom', ha='left', color='C1', rotation=90)
        elif int(year) == 1784:
            plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=75)
        else:
            plt.text(year, peak, f'{year:.0f}', va='bottom', ha='left', color='C1', rotation=45)
    peaks = da.pipe(lambda x: -x).go.find_peaks()
    for year,peak in zip(peaks.year.values, -peaks.peaks.values):
        if int(year) == 1020:
            plt.text(year, peak, f'    {year:.0f}', va='top', ha='left', color='C2', rotation=-90)
        elif int(year) == 1049:
            plt.text(year, peak, f'    {year:.0f}', va='top', ha='left', color='C2', rotation=-75)
        elif int(year) == 1100:
            plt.text(year, peak, f'    {year:.0f}', va='top', ha='left', color='C2', rotation=-60)
        elif int(year) == 1144:
            plt.text(year, peak, f'    {year:.0f}', va='top', ha='left', color='C2', rotation=-15)
        elif int(year) == 1360:
            plt.text(year, peak, f'    {year:.0f}', va='top', ha='left', color='C2', rotation=-60)
        else:
            plt.text(year, peak, f'{year:.0f}', va='top', ha='left', color='C2', rotation=-45)
    plt.ylabel(f'{case} HU')

    ax.set_xticks(range(2000,850,-100))
    ax.set_xlim(850,2000)
    for ax in axes[0:-1]:
        ax.set_xlabel('')


    
    figname = __file__.replace('.py', f'.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()