from qgis.PyQt.QtCore import QSize, QT_VERSION_STR

if QT_VERSION_STR.startswith("5"):
    from matplotlib.backends.backend_qt5agg import (
        NavigationToolbar2QT as NavigationToolbar,
    )
else:
    from matplotlib.backends.backend_qtagg import (
        NavigationToolbar2QT as NavigationToolbar,
    )


class MyNavigationToolbar(NavigationToolbar):
    """Navigation bar for plot with home button, pan and zoom."""

    toolitems = [
        t for t in NavigationToolbar.toolitems if t[0] in ("Home", "Pan", "Zoom")
    ]

    def __init__(self, *args, **kwargs):
        super(MyNavigationToolbar, self).__init__(*args, **kwargs)
        self.layout().takeAt(3)  # 3 = Amount of tools we need
        self.layout().setSpacing(5)
        self.setIconSize(QSize(20, 20))


def zoom_with_wheel(FigCanvas, ax, zoomScale):
    """Implements zooming in plot whith mouse wheel. The function itself has
    to be returned and saved to a variable, otherwise it could be garbage
    collected."""

    def zoomFunc(event):
        # Get the current x and y limits
        cur_xlim = ax.get_xlim()
        cur_ylim = ax.get_ylim()
        xdata = event.xdata  # Get event x location
        ydata = event.ydata  # Get event y location
        if event.button == "up":
            # Deal with zoom in
            scale_factor = 1 / zoomScale
        elif event.button == "down":
            # Deal with zoom out
            scale_factor = zoomScale
        else:
            scale_factor = 1
        # Set new limits
        ax.set_xlim(
            [
                xdata - (xdata - cur_xlim[0]) * scale_factor,
                xdata + (cur_xlim[1] - xdata) * scale_factor,
            ]
        )
        ax.set_ylim(
            [
                ydata - (ydata - cur_ylim[0]) * scale_factor,
                ydata + (cur_ylim[1] - ydata) * scale_factor,
            ]
        )
        FigCanvas.draw()

    # Get the figure of interest
    fig = ax.get_figure()
    # Attach the call back
    fig.canvas.mpl_connect("scroll_event", zoomFunc)
    return zoomFunc
