import numpy as np
import matplotlib.pyplot as plt
#
# Original colormap
#
#import cmocean
#
from matplotlib.gridspec import GridSpec
from matplotlib.colors import ListedColormap, LinearSegmentedColormap

data= np.load('data.npy') #data in shape: simulation, time, Nx, Ny
sims = np.load('sims.npy') #simulation labels for data 
times = np.load('times.npy') #time labels

#Set up plot parameters
#
# Original colormap
#
#First create a custom colormap (or just replace newcmap with desired colormap)
#cmap = cmocean.cm.curl
#cmap2 = cmocean.cm.balance
#vmin=-1
#vmax=1
#top = cmap2(np.linspace(0.05,0.5,128))
#bottom = cmap(np.linspace(0.5,0.9,128))
#newcolors = np.vstack((top, bottom))
#newcmap = ListedColormap(newcolors)

newcmap="Spectral"

#Font sizes
SMALL_SIZE=15
MEDIUM_SIZE=15
BIGGER_SIZE=15
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

label_loc = (40,650)
#
#levels= np.linspace(-1.1,1.1,31) #specify range of density colorbar
#
levels= np.linspace(-0.8,0.8,31) #specify range of density colorbar

#Set up plotting grid
gs=GridSpec(3,5, width_ratios=[1,1,1,1,0.1], wspace=0.075, hspace=0.075)
fig=plt.figure(figsize=(16,7))
ax1=fig.add_subplot(gs[0,0]) # First row, first column
ax2=fig.add_subplot(gs[0,1]) # First row, second column
ax3= fig.add_subplot(gs[0,2]) #First row, third column
ax4=fig.add_subplot(gs[0,3]) #First row, fourth column
axcbar=fig.add_subplot(gs[:,4]) 
ax5=fig.add_subplot(gs[1,0]) # First row, first column
ax6=fig.add_subplot(gs[1,1]) # First row, second column
ax7= fig.add_subplot(gs[1,2]) #First row, third column
ax8=fig.add_subplot(gs[1,3]) #First row, fourth column
ax9=fig.add_subplot(gs[2,0]) # First row, first column
ax10=fig.add_subplot(gs[2,1]) # First row, second column
ax11= fig.add_subplot(gs[2,2]) #First row, third column
ax12=fig.add_subplot(gs[2,3]) #First row, fourth column

cs = ax1.contourf(data[0,0,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(a)$',label_loc, color='w', size=15)
#
ax1.annotate('$(a)$',label_loc, color='k', size=15)
ax1.set_title(f'$T={times[0]:.2f}$')
ax1.set_xticks([])
ax1.set_yticks([])
ax1.set_rasterization_zorder(-1)
    
cs = ax5.contourf(data[1,0,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(e)$',label_loc, color='w', size=15)
#
ax5.annotate('$(e)$',label_loc, color='k', size=15)
ax5.set_xticks([])
ax5.set_yticks([])
ax5.set_rasterization_zorder(-1)

cs = ax9.contourf(data[2,0,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(i)$',label_loc, color='w', size=15)
#
ax9.annotate('$(i)$',label_loc, color='k', size=15)
ax9.set_xticks([])
ax9.set_yticks([])
ax9.set_rasterization_zorder(-1)

cs = ax2.contourf(data[0,1,:,:], cmap=newcmap, levels=levels, zorder=-9)
ax2.set_title(f'$T={times[1]:.2f}$')
#
#ax1.annotate('$(b)$',label_loc, color='w', size=15)
#
ax2.annotate('$(b)$',label_loc, color='k', size=15)
ax2.set_xticks([])
ax2.set_yticks([])
ax2.set_rasterization_zorder(-1)

cs = ax6.contourf(data[1,1,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(f)$',label_loc, color='w', size=15)
#
ax6.annotate('$(f)$',label_loc, color='k', size=15)
ax6.set_xticks([])
ax6.set_yticks([])
ax6.set_rasterization_zorder(-1)

cs = ax10.contourf(data[2,1,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(j)$',label_loc, color='w', size=15)
#
ax10.annotate('$(j)$',label_loc, color='k', size=15)
ax10.set_xticks([])
ax10.set_yticks([])
ax10.set_rasterization_zorder(-1)

cs = ax3.contourf(data[0,2,:,:], cmap=newcmap, levels=levels, zorder=-9)
ax3.set_title(f'$T={times[2]:.2f}$')
#
#ax1.annotate('$(c)$',label_loc, color='w', size=15)
#
ax3.annotate('$(c)$',label_loc, color='k', size=15)
ax3.set_xticks([])
ax3.set_yticks([])
ax3.set_rasterization_zorder(-1)

cs = ax7.contourf(data[1,2,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(g)$',label_loc, color='w', size=15)
#
ax7.annotate('$(g)$',label_loc, color='k', size=15)
ax7.set_xticks([])
ax7.set_yticks([])
ax7.set_rasterization_zorder(-1)

cs = ax11.contourf(data[2,2,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(k)$',label_loc, color='w', size=15)
#
ax11.annotate('$(k)$',label_loc, color='k', size=15)
ax11.set_xticks([])
ax11.set_yticks([])
ax11.set_rasterization_zorder(-1)

cs = ax4.contourf(data[0,3,:,:], cmap=newcmap, levels=levels, zorder=-9)
ax4.set_title(f'$T={times[3]:.2f}$')
#
#ax1.annotate('$(d)$',label_loc, color='w', size=15)
#
ax4.annotate('$(d)$',label_loc, color='k', size=15)
ax4.set_xticks([])
ax4.set_yticks([])
ax4.set_rasterization_zorder(-1)
    
cs = ax8.contourf(data[1,3,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(h)$',label_loc, color='w', size=15)
#
ax8.annotate('$(h)$',label_loc, color='k', size=15)
ax8.set_xticks([])
ax8.set_yticks([])
ax8.set_rasterization_zorder(-1)

cs = ax12.contourf(data[2,3,:,:], cmap=newcmap, levels=levels, zorder=-9)
#
#ax1.annotate('$(l)$',label_loc, color='w', size=15)
#
ax12.annotate('$(l)$',label_loc, color='k', size=15)
ax12.set_xticks([])
ax12.set_yticks([])
ax12.set_rasterization_zorder(-1)

cbar = fig.colorbar(cs, cax=axcbar)
cbar.set_label(label='$\\mathrm{Density}$ $\\rho$', labelpad=-5)
#
#cbar.set_ticks([-1,0,1])
#
cbar.set_ticks([-0.8,  -0.4,0,   0.4,  0.8])