#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Feb 10 13:53:50 EST 2021
#WY2021-06-01: fig_lines_tc_all -> fig_lines_tc_all_poisson (for LMR HU/MHs)
#WY2021-06-03: fig_lines_tc_all_poisson -> fig_lines_tc_all_stde (for LMR HU/MHs); fig_lines_mdrssts -> fig_lines_mdrssts_stde
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
#from fig_lines_tc_liz_lmr2019 import plot_lines_liz, plot_hu_lmr, plot_mdr_lmr #replaced with beta1
#beta1
#from fig_lines_lizjk import wyplot as plot_lines_liz
from fig_lines_lizjk_HUcount import wyplot as plot_lines_liz #replace normalized HU acrtivity with HU count
#from fig_lines_tc_all import wyplot as plot_hu_lmr
#from fig_lines_tc_all_poisson import wyplot as plot_hu_lmr
#from fig_lines_tc_all_stde import wyplot as plot_hu_lmr
#from fig_lines_tc_all_poisson_HUcount import wyplot as plot_hu_lmr #replace normalized HU acrtivity with HU count
from fig_lines_tc_all_poisson_HUcount_tfrMDRTROP18702000 import wyplot as plot_hu_lmr #wy2021-10-18: use transformed rMDR/TROP SSTs
#from fig_lines_mdrssts import wyplot as plot_mdr_lmr
#from fig_lines_mdrssts_stde import wyplot as plot_mdr_lmr
from fig_lines_mdrssts_stde_tfrMDRTROP18702000 import wyplot as plot_mdr_lmr #wy2021-10-18: use transformed rMDR/TROP SSTs
from shared import labelax
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
#subplot images
fig_map = 'fig_map_sites.png'
fig_events = 'fig_dots_sites.png'
 

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import matplotlib.image as image, matplotlib.gridspec as gridspec
    from shared import scaleax
    #fig = plt.figure(figsize=(8*2,4*2+1), dpi=80)
    constrained_layout_off()
    #fig, axes = plt.subplots(3, 2, figsize=(8*2,4*2+1), dpi=80)
    fig, axes = plt.subplots(3, 2, figsize=(8*2,4*2+1), dpi=80)
    #norm_years = slice(1851,2000)
    #norm_years = slice(1901,2000)
    norm_years = slice(1870,2000) #wy2021-10-18: use the same period as the one in the transformation of rMDR/TROP SSTs
    #gs = fig.add_gridspec(120,120)
    afontsize = 'large'

    #ax = fig.add_subplot(gs[0:40, 60+dx:120]) #ax a
    ax = axes[0, 0]
    plot_lines_liz(ax, norm_years=norm_years)
    ax.set_xlim(850, 2000)
    ax.set_xticks(range(2000, 850, -100))
    ax.set_xlabel('')
    #ax.set_ylim(0, 0.07)
    ax.set_yticks(range(3,16,2))
    ax.set_title('Hurricanes from multi-site sediment reconstructions and modern records', loc='left')
    labelax(ax, 'A'+' '*4+'\n', fontsize=afontsize)

    #ax = fig.add_subplot(gs[40:80, 60+dx:120]) #ax b
    ax = axes[1, 0]
    plot_hu_lmr(ax, norm_years=norm_years)
    ax.set_xlim(850, 2000)
    ax.set_xticks(range(2000, 850, -100))
    #ax.set_ylim(0, 0.07)
    ax.set_yticks(range(3,16,2))
    ax.set_xlabel('')
    ax.set_title('Hurricanes from statistical model using LMR SSTs', loc='left')
    labelax(ax, 'B'+' '*4+'\n',fontsize=afontsize)

    #ax = fig.add_subplot(gs[80:, 60+dx:120]) #ax c
    ax = axes[2,0]
    plot_mdr_lmr(ax)
    ax.set_xlim(850, 2000)
    ax.set_xticks(range(2000, 850, -100))
    ax.set_xlabel('year')
    ax.set_title('MDR and tropical mean SST anomalies', loc='left')
    labelax(ax, 'C    \n', fontsize=afontsize)

    #ax = fig.add_subplot(gs[80:120, 0:60+dx]) #ax d
    ax = axes[0,1]
    img = image.imread(fig_map)
    ax.imshow(img)
    ax.axis('off')
    labelax(ax, 'D', fontsize=afontsize)

    #ax = fig.add_subplot(gs[40:80, 0:60+dx]) #ax f
    ax = axes[2,1]
    img = image.imread(fig_events)
    ax.imshow(img)
    ax.axis('off')
    labelax(ax, 'E', fontsize=afontsize)

    #gs.tight_layout(fig)
    #scaleax(axa, 0.2, 0.2, 0.2, 0.2)
    #fig.set_constrained_layout_pads(wspace=0.2)
    fig.tight_layout()
    scaleax(axes[0,0], right=0.1)
    scaleax(axes[1,0], right=0.1)
    scaleax(axes[2,0], right=0.1)
    scaleax(axes[0,1], right=0.4, left=0.4, down=0.8)
    axes[1,1].set_axis_off()
    scaleax(axes[2,1], right=0.3, left=0.3, up=0.6)
    
    figname = __file__.replace('.py', f'.png')
    if norm_years is not None:
        figname = figname.replace('.png', f'_norm{norm_years.start}-{norm_years.stop}.png')
    if len(sys.argv) > 1 and sys.argv[1] == 'savefig':
        wysavefig(figname, dpi=200)
    tt.check(f'**Done**')
    plt.show()
    constrained_layout_on()