import scipy.io as sio
import matplotlib.pyplot as plt
import numpy as np
import matplotlib.gridspec as gridspec
import matplotlib.cm as cm
from mpl_toolkits.axes_grid1 import make_axes_locatable

# Bo Map
mat_data_0=sio.loadmat('bo.mat')
tmap = mat_data_0['boMap']
print(np.min(tmap))
print(np.max(tmap))
vmin = -0.3
vmax = 0.3

def multiSlicePlot(data,nRow, nCol, nSlrep, vmin, vmax, sliceInit = 0):
    """""
    data = imageMatrix3D [sl,ph,rd]
    nRow, nCol dimensions of the multislicePlot
    nSlRep =  number of slices to represent
    sliceInit = slice in wihich we start representing
    """""
    images = []
    fig = plt.figure(figsize=(nRow, nCol), dpi=500)
    gs1 = gridspec.GridSpec(nRow, nCol)
    gs1.update(wspace=0.020, hspace=0.020)  # set the spacing between axes.
    for i in range(nSlrep):
        ii = i + sliceInit
        # print(ii)
        ax1 = plt.subplot(gs1[i])
        # plt.axis('off')
        ax1.set_xticklabels([])
        ax1.set_yticklabels([])
        ax1.set_aspect('equal')
        dataAux = data[int(ii), :, :]
        imgPlot = ax1.imshow(dataAux,vmin=vmin, vmax=vmax)
        images.append(imgPlot)
        ax1.axis('off')
    return fig

fig = multiSlicePlot(tmap*1e3,3,6,16,vmin,vmax,0)

cbar_ax = fig.add_axes([0.92, 0.15, 0.01, 0.7]) 
norm = plt.Normalize(vmin=vmin, vmax=vmax)
cbar = plt.colorbar(cm.ScalarMappable(norm=norm, cmap=cm.viridis), cax=cbar_ax)
cbar.set_label('Bo [mT]', fontsize=2)
cbar.ax.tick_params(labelsize=2)
plt.show()