import numpy as np
from PyQt5 import QtWidgets

import matplotlib
from matplotlib.figure import Figure
import matplotlib.cm as cm
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from matplotlib.axes._axes import Axes

matplotlib.use("Qt5Agg")


class MplCanvas(FigureCanvasQTAgg):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.axe = self.fig.add_subplot(1, 1, 1, label='good')
        super().__init__(self.fig)


class MainWindow(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()

        layout = QtWidgets.QVBoxLayout()

        self.canvas = MplCanvas(self, width=5, height=4, dpi=100)
        self.axe = self.canvas.axe

        layout.addWidget(self.canvas)

        self.pushButton_addColorBar = QtWidgets.QPushButton('Add ColorBars')
        layout.addWidget(self.pushButton_addColorBar)
        self.pushButton_remove = QtWidgets.QPushButton('remove')
        layout.addWidget(self.pushButton_remove)

        # Create a placeholder widget to hold our toolbar and canvas.
        widget = QtWidgets.QWidget()
        widget.setLayout(layout)
        self.setCentralWidget(widget)

        self.pushButton_remove.clicked.connect(self.removeImage)
        self.pushButton_addColorBar.clicked.connect(self.createColorBars)

        self.pcolormesh_test()

    def pcolormesh_test(self):
        """add two images"""
        delta = 0.01
        x = y = np.arange(-3.0, 3.0, delta)
        X, Y = np.meshgrid(x, y)
        Z1 = np.exp(-X ** 2 - Y ** 2)
        Z2 = np.exp(-(X - 1) ** 2 - (Y - 1) ** 2)
        Z = (Z1 - Z2) * 2

        im = self.axe.pcolormesh(X, Y, Z, cmap=cm.viridis, shading='auto')  # , extent=[-3, 3, -3, 3]
        im.set_clim(vmax=np.amax(Z), vmin=np.amin(Z))

        Zx = (Z1 + Z2) * 2
        imx = self.axe.pcolormesh(X, Y, Zx, cmap=cm.Blues, shading='auto')  # , extent=[-3, 3, -3, 3], , cmap=mpl.cm.Blues
        imx.set_clim(vmax=np.amax(Zx), vmin=np.amin(Zx))

    def createColorBars(self):
        """ to create a color bar for an image. """
        axe = self.axe
        fig = axe.get_figure()
        images = self.getImages(axe)

        for image in images:
            if not image.colorbar:  # color bar doesn't exist
                inset_axe = axe.inset_axes([1.115, 0, 0.05, 1], transform=axe.transAxes)  # [x0, y0, width, height]
                fig.colorbar(image, ax=axe, cax=inset_axe)
                break  # each trigger create one colorbar for one image

        self.reDraw()

    def removeImage(self):
        """
        Usage:
            * each trigger remove one image

        """
        images = self.getImages(self.axe)
        print(f'images={images}')
        if images:
            image = images[-1]
            color_bar = image.colorbar
            if color_bar:
                color_bar.remove()
                del color_bar

            # remove image
            image.remove()  # remove line from the list of ax.lines
            del image  # remove the line from the memory if no other variable reference to the line

        self.reDraw()

    def getImages(self, axe: Axes):
        """to obtain the image list in the axe"""
        images = []
        images.extend(axe.images)
        images.extend(axe.collections)

        return images

    def reDraw(self):
        self.canvas.draw_idle()
        self.canvas.flush_events()


if __name__ == '__main__':
    import sys
    app = QtWidgets.QApplication(sys.argv)

    window = MainWindow()
    window.show()
    sys.exit(app.exec_())