# import libraries
import numpy as np 
from netCDF4 import Dataset 
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.colors as mcolors

# Read peak stratospheric volcanic aerosol mass loads 
#   These are derived from sulphate depositions of Sigl et al. 2015 and
#   converted to stratospheric aerosol mass using the Gao et al. 2007 scaling
volc_data = np.genfromtxt('nature14565-s6_stratAerosolMass.txt')
year = volc_data[:,0]   # extract years
flag = volc_data[:,1]   # extract location flag (1=tropical, 2=NH, 3=SH)
mass = volc_data[:,2]   # extract stratospheric volcanic aerosol mass (Tg)

# define ensemble size and time period for synthetic forcing
NMEM = 60
YEAR1 = 2006
YEARN = 2099

# re-sample eruption events from reconstruction
np.random.seed(0)
P=1./2500./12.
synth_year = []
synth_month = []
synth_flag = []
synth_mass = []
for member in range(NMEM):
    synth_year.append([])
    synth_month.append([])
    synth_flag.append([])
    synth_mass.append([])
    dummy=np.random.rand()
    # loop over years 
    for y in range(YEAR1,YEARN+1):
        # loop over months 
        for m in range(1,12+1):
            # loop over eruptions in Sigl el al. 2015 
            for eruption in range(len(mass)):
                # check if eruption is triggered and add to column mass if yes
                if ( np.random.rand() < P ):
                    synth_year[member].append(y) 
                    synth_month[member].append(m) 
                    synth_flag[member].append(flag[eruption]) 
                    synth_mass[member].append(mass[eruption]) 
    dummy=np.random.rand()
    
# plot chronologies

SMALL_SIZE = 18
MEDIUM_SIZE = 22
BIGGER_SIZE = 25
#
plt.rc('font', size=SMALL_SIZE)          # controls default text sizes
plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title
plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels
plt.rc('xtick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('ytick', labelsize=SMALL_SIZE)    # fontsize of the tick labels
plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize
plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title
#
fig,axes = plt.subplots(nrows=4,ncols=1, sharex=False, sharey=False, figsize=(20,12))
#
axId = 0 
for e in range(len(year)):
    if flag[e] == 1:
        axes[axId].plot([year[e],year[e]],[0,mass[e]],color='k')
    elif flag[e] == 2:
        axes[axId].plot([year[e],year[e]],[0,mass[e]],color='b')
    elif flag[e] == 3:
        axes[axId].plot([year[e],year[e]],[0,mass[e]],color='r')
axes[axId].set_xlim(-500,2000) 
axes[axId].set_ylim(0,180)
axes[axId].set_xlabel('Year CE')
#axes[axId].set_ylabel('Load (Tg)')
axes[axId].set_yticks([0,25,50,75,100,125,150,175])
#
memRanges = [[1,25],[26,50],[51,60]]
for axId in [1,2,3]:
    for m in range(memRanges[axId-1][0]-1,memRanges[axId-1][1]):
        if (m+axId) % 2 == 0:
            rect = patches.Rectangle(((m-(axId-1)*25)*100 - 500 + 6,0),95,180,edgecolor='none',facecolor=mcolors.to_rgb('ivory'))
        else:
            rect = patches.Rectangle(((m-(axId-1)*25)*100 - 500 + 6,0),95,180,edgecolor='none',facecolor=mcolors.to_rgb('aliceblue'))        
        axes[axId].add_patch(rect)
        for e in range(len(synth_year[m])):
            syear = synth_year[m][e] + (m-(axId-1)*25)*100 - 2500
            smass = synth_mass[m][e]
            if synth_flag[m][e] == 1:
                axes[axId].plot([syear,syear],[0,smass],color='k')
            elif synth_flag[m][e] == 2:
                axes[axId].plot([syear,syear],[0,smass],color='b')
            elif synth_flag[m][e] == 3:
                axes[axId].plot([syear,syear],[0,smass],color='r')
    axes[axId].plot([-500,-500],[0,0],color='k',label='Tropical')
    axes[axId].plot([-500,-500],[0,0],color='b',label='Extratropical NH')
    axes[axId].plot([-500,-500],[0,0],color='r',label='Extratropical SH')
    axes[axId].set_xlim(-500,2000)
    axes[axId].set_ylim(0,180) 
    if axId == 3:
        axes[axId].set_xlabel('Member')
        axes[axId].legend()
    #axes[axId].set_ylabel('Load (Tg)')
    xticks = []
    xticklabels = []
    for m in range(memRanges[axId-1][0]-1,memRanges[axId-1][1]):
        xticks.append((m-(axId-1)*25)*100 - 500 + 6 + 95/2)
        xticklabels.append(m+1)
    axes[axId].set_xticks(xticks)
    axes[axId].set_xticklabels(xticklabels)
    axes[axId].set_yticks([0,25,50,75,100,125,150,175])
#
fig.subplots_adjust(hspace=0.45,wspace=0)
fig.add_subplot(111, frame_on=False)
plt.tick_params(labelcolor="none", bottom=False, left=False)
plt.ylabel("Stratospheric volcanic aerosol load [Tg]",labelpad=20)
#
plt.savefig('forcing_chronology.jpg',format='jpeg',dpi=200,quality=80,bbox_inches='tight')