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

# read historical forcing data to compute vertical shape function
nc = Dataset('CCSM4_volcanic_1850-2008_prototype1.nc')
colmass = nc.variables['colmass'][:]
mmrvolc = nc.variables['MMRVOLC'][:]
lev = nc.variables['lev'][:]
nc.close()

# Create vertical profile (from matlab code of EXPLAIN)
profile_lat = np.zeros((64,8))
for i in range(8):
    zlev = mmrvolc[:,i,:]/colmass
    profile_lat[:,i] = np.nanmean(zlev, axis=0)
profile_lat = profile_lat/np.nanmax(profile_lat)
    
# Load the CCSM data for horizontal shape 
nc = Dataset('volc_samples.nc')
month = nc.variables['month'][:]
St = nc.variables['Tropical_Volcanoe'][:,:]
Sh = nc.variables['Hemisphere_Volcanoe'][:,:]
lat = nc.variables['lat'][:]
nc.close()

# Normalise horizontal shapes by time maximum of global integral
area_earth = 4*np.pi*6371000**2
dlat = 2.79;
weights = np.cos(lat/180.*np.pi)
weights = weights/sum(weights)*area_earth
St = St/np.max(St) # convert from Tg to kg
Sh = Sh/np.max(Sh)
S = np.transpose(np.concatenate([np.zeros([1,64]),St[0:40,:],np.fliplr(Sh[0:30,:]),Sh[0:30,:]]))
# global averages
Gt = np.sum(St*weights.shape,axis=1)
Gh = np.sum(Sh*weights.shape,axis=1)
Gt = Gt/np.max(Gt)
Gh = Gh/np.max(Gh)
G = np.concatenate([[0],Gt[0:40],Gh[0:30],Gh[0:30]])
# vertical 
profile = np.sum(np.transpose(profile_lat)*weights,axis=1)
profile = profile / np.max(profile)

# prepare coordinates
latc = np.zeros(len(lat)+1)
latc[0] = -90
for i in range(len(lat)-1):
    latc[i+1] = 0.5*(lat[i]+lat[i+1])
latc[-1] = 90
tvec = np.arange(0,102)
[X,Y] = np.meshgrid(tvec[0:-1],lat)
[XC,YC] = np.meshgrid(tvec,latc)
[XP,YP] = np.meshgrid(lev,lat)

# plot 
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=2,ncols=2, sharex=False, sharey=False, figsize=(20,12),gridspec_kw={'width_ratios': [3, 1],'height_ratios': [1, 3]})
#
axes[0,0].plot(tvec[0:-1],G,color='k')
axes[0,0].spines["top"].set_visible(False)
axes[0,0].spines["right"].set_visible(False)
axes[0,0].spines["bottom"].set_visible(False)
axes[0,0].set_xlim(0,100)
axes[0,0].set_ylim(0,1)
axes[0,0].get_xaxis().set_visible(False)
axes[0,0].set_ylabel('Global load')
#
axes[0,1].axis('off')
#axes[0,1].plot(lev,profile,color='k')
#axes[0,1].get_xaxis().set_visible(False)
#axes[0,1].get_yaxis().set_visible(False)
#axes[0,1].spines["top"].set_visible(False)
#axes[0,0].spines["right"].set_visible(False)
#axes[0,1].spines["left"].set_visible(False)
#axes[0,1].spines["bottom"].set_visible(False)
#axes[0,1].set_ylim(0,1)
#axes[0,1].set_xlim(0,150)
#axes[0,1].yaxis.set_ticks_position('right')
#fig.colorbar(h,ax=axes[0,1],label='Normalized load',orientation='vertical',shrink=1)
#axes[0,1].set_label_position("right")
#axes[0,1].tick_right()
#
axes[1,0].contourf(X,Y,S,levels=np.arange(0,1+0.01,0.01),cmap=plt.get_cmap('afmhot_r'))
axes[1,0].plot([40,40],[-90,90],color='grey')
axes[1,0].plot([70,70],[-90,90],color='grey')
axes[1,0].set_xlim(0,100)
axes[1,0].set_ylim(-90,90)
axes[1,0].set_xticks([0,12,24,36,40,52,64,70,82,94])
axes[1,0].set_xticklabels([0,12,24,36,0,12,24,0,12,24])
axes[1,0].set_yticks([-90,-60,-30,0,30,60,90])
axes[1,0].set_xlabel('Months since eruption onset')
axes[1,0].set_ylabel('Latitude')
#
h=axes[1,1].contourf(XP,YP,profile_lat,levels=np.arange(0,1+0.01,0.01),cmap=plt.get_cmap('afmhot_r'))
#axes[1,1].spines["left"].set_visible(False)
axes[1,1].set_xlim(0,150)
axes[1,1].set_ylim(-90,90)
axes[1,1].set_xticks([0,50,100,150])
axes[1,1].set_yticks([-90,-60,-30,0,30,60,90])
axes[1,1].set_xlabel('Pressure [hPa]')
#
for axis in ['top','bottom','left','right']:
    axes[1,0].spines[axis].set_linewidth(1.5)
    axes[1,1].spines[axis].set_linewidth(1.5)
    axes[0,0].spines[axis].set_linewidth(1.5)

#
fig.colorbar(h,ax=axes[1,1],label='Normalized load',orientation='vertical',shrink=1,ticks=np.arange(0,1.1,0.1))
#
fig.subplots_adjust(hspace=0.05,wspace=0.1)
plt.savefig('forcing_shape.png',format='png',dpi=300,bbox_inches='tight')
plt.savefig('forcing_shape.pdf',format='pdf',bbox_inches='tight')
plt.savefig('forcing_shape.eps',format='eps',bbox_inches='tight')
plt.savefig('forcing_shape.jpg',format='jpeg',dpi=200,quality=80,bbox_inches='tight')

#plt.close('all')