import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
import mikeio
from matplotlib.pyplot import MultipleLocator
import matplotlib.colors as mcolors

## define sequential colorbar
def create_colorbar(vmin, vmax, cmap):
    norm = plt.Normalize(vmin=vmin, vmax=vmax)
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    return sm

## define TwoSlopeNorm colorbar
def create_TwoSlopeNorm_colorbar(vmin, vmax, vcenter, cmap):
    norm = mcolors.TwoSlopeNorm(vmin=vmin, vcenter=vcenter, vmax=vmax)
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    return sm

gs = GridSpec(2, 1, wspace = 0.08, hspace = 0.3)   
fig = plt.figure(dpi=600, figsize=(6.4, 5))
ax1 = fig.add_subplot(gs[0, 0])
ax2 = fig.add_subplot(gs[1, 0])
file_paths = [r'E:\Model\MT-model.m21fm - Result Files\Case_1.dfsu', 
              r'E:\Model\MT-model.m21fm - Result Files\Case_2.dfsu']

cmap = 'bwr'
vmax = 2
vmin = -6

item = 'Total bed thickness change'
step = -1   ## last time step

texts = ['(a) Case 1', '(b) Case 2']
axes = [ax1, ax2]

## =============================================================================
## plot figure
ds = mikeio.read(file_paths[0])
ds[item][step].plot(ax=ax1, vmin = vmin, vmax = vmax, cmap = cmap, label = '', add_colorbar = True, plot_type = 'shaded')   ## original: colorbar 1

ds_2 = mikeio.read(file_paths[1])
ds_2[item][step].plot(ax=ax2, vmin = vmin, vmax = vmax, cmap = cmap, label = '', add_colorbar = True, plot_type = 'shaded')  ## original: colorbar 1

for ax, text in zip([ax1, ax2], texts):
    ax.set_title('')
    ax.tick_params(direction = 'in', axis = 'both', labelsize = ax.xaxis.get_label().get_fontsize())
    ax.set_xlabel('')
    ax.set_ylabel('')
    ax.text(0.0, 1.05, text, transform = ax.transAxes)
    
    x_major_locator = MultipleLocator(5000)
    ax.xaxis.set_major_locator(x_major_locator)

## =============================================================================
## add colorbar 2
colorbar_2 = fig.colorbar(mappable = create_colorbar(vmin=vmin, vmax=vmax, cmap=cmap), 
                        cax=fig.add_axes([0.9, 0.15, 0.02, 0.7]),
                        orientation='vertical',
                        extend = 'both')
colorbar_2.set_label(label = 'colorbar 1', fontsize = ax.xaxis.get_label().get_fontsize())
colorbar_2.ax.xaxis.set_tick_params(labelsize = ax.xaxis.get_label().get_fontsize())
colorbar_2.set_ticks([-6, -4, -2, 0, 2])
colorbar_2.set_ticklabels([-6, -4, -2, 0, 2])

## =============================================================================
## add colorbar 3
colorbar_3 = fig.colorbar(mappable = create_TwoSlopeNorm_colorbar(vmin=vmin, vmax=vmax, vcenter = 0, cmap=cmap), 
                        cax=fig.add_axes([1.03, 0.15, 0.02, 0.7]),
                        orientation='vertical',
                        extend = 'both')
colorbar_3.set_label(label = 'colorbar 2', fontsize = ax.xaxis.get_label().get_fontsize())
colorbar_3.ax.xaxis.set_tick_params(labelsize = ax.xaxis.get_label().get_fontsize())
colorbar_3.set_ticks([-6, -4, -2, 0, 2])
colorbar_3.set_ticklabels([-6, -4, -2, 0, 2])