from typing import Optional

import numpy as np
from pyqtgraph import PlotWidget, ImageItem
from scipy.interpolate import LinearNDInterpolator

from bivariatecolorrasterrendererapp.bivariatecolorrasterrenderer import BivariateColorRasterRenderer
from enmapbox.qgispluginsupport.qps.utils import SpatialExtent
from enmapbox.typeguard import typechecked
from enmapboxprocessing.rasterreader import RasterReader
from qgis.PyQt.QtGui import QColor
from qgis.PyQt.QtWidgets import QWidget, QToolButton, QCheckBox, QMainWindow, QComboBox, QMenu, QAction, QPushButton, \
    QLineEdit, QRadioButton
from qgis.PyQt.uic import loadUi
from qgis.core import QgsRasterLayer, QgsMapLayerProxyModel, QgsMapSettings
from qgis.gui import QgsMapCanvas, QgsRasterBandComboBox, QgsDoubleSpinBox, QgsMapLayerComboBox, QgsColorButton, \
    QgsSpinBox


@typechecked
class ColorPlanePlotWidget(PlotWidget):
    def __init__(self, *args):
        PlotWidget.__init__(self, parent=None, background='#f0f0f0')
        self.setDefaultPadding(padding=0.02)


@typechecked
class BivariateColorRasterRendererDialog(QMainWindow):
    mLayer: QgsMapLayerComboBox
    mBand1: QgsRasterBandComboBox
    mBand2: QgsRasterBandComboBox
    mMin1: QLineEdit
    mMin2: QLineEdit
    mMax1: QLineEdit
    mMax2: QLineEdit
    mColor1: QgsColorButton
    mColor2: QgsColorButton
    mColor3: QgsColorButton
    mColor4: QgsColorButton
    mUseColor2: QCheckBox
    mPlot: ColorPlanePlotWidget
    mMenu: QPushButton
    mMode: QComboBox
    mClasses: QgsSpinBox
    mMinMaxUser: QRadioButton
    mMinMaxPercentile: QRadioButton
    mP1: QgsDoubleSpinBox
    mP2: QgsDoubleSpinBox
    mExtent: QComboBox
    mAccuracy: QComboBox

    mLiveUpdate: QCheckBox
    mApply: QToolButton

    ContinuouseMode, EqualIntervalMode, QuantileMode = 0, 1, 2
    EstimatedAccuracy, ActualAccuracy = 0, 1
    WholeRasterExtent, CurrentCanvasExtent = 0, 1

    def __init__(self, parent=None):
        QWidget.__init__(self, parent)
        loadUi(__file__.replace('.py', '.ui'), self)

        from enmapbox.gui.enmapboxgui import EnMAPBox
        self.enmapBox = EnMAPBox.instance()
        self.colorPlane: np.ndarray

        self.mMapCanvas: Optional[QgsMapCanvas] = None
        self.mLayer.setProject(self.enmapBox.project())
        self.mLayer.setFilters(QgsMapLayerProxyModel.RasterLayer)
        self.mP1.setClearValue(self.mP1.value())
        self.mP2.setClearValue(self.mP2.value())

        self.predefinedColorPlanes = {
            'Katja Kowalski 1': ['#0198eb', '#830050', '#bebebf', '#fee413'],
            'Joshua Stevens 1': ['#64acbe', '#574249', '#e8e8e8', '#c85a5a'],
            'Joshua Stevens 2': ['#be64ac', '#3b4994', '#e8e8e8', '#5ac8c8'],
            'Joshua Stevens 3': ['#73ae80', '#2a5a5b', '#e8e8e8', '#6c83b5'],
            'Joshua Stevens 4': ['#9972af', '#804d36', '#e8e8e8', '#c8b35a'],
            'ArcMap BrBG': ['#a85e18', '#537440', '#f2f7f5', '#00856f'],
            'ArcMap BrBu_Dk': ['#c68023', '#463333', '#cccccc', '#5a84bc'],
            'ArcMap BrBu_Lt': ['#c58225', '#8d866c', '#cccccc', '#5785ba'],
            'ArcMap BuOr': ['#4a9fc2', '#364f4a', '#f3f3f3', '#f4b100'],
            'ArcMap PiGn': ['#d4188a', '#8d635a', '#f7f7f7', '#49aa25'],
            'ArcMap PuGn': ['#793293', '#405164', '#f6f6f6', '#018738'],
            'ArcMap PuOr': ['#a85e18', '#537440', '#f2f7f5', '#e85f03'],
            'ArcMap RdBu': ['#c90223', '#683769', '#f7f7f7', '#0470b2'],
            'ArcMap PiBu': ['#b94a8d', '#4e428b', '#cdc5da', '#0b94c9'],
            'ArcMap RdGy': ['#cb0022', '#993245', '#ffffff', '#656565'],
            'ArcMap RdYlBu': ['#d31820', '#85486a', '#feffbf', '#2c7ab8'],
            'ArcMap RdYlGn': ['#d81b1f', '#77572e', '#feffc1', '#199942']
        }
        menu = QMenu()
        for name in self.predefinedColorPlanes:
            menu.addAction(name, self.onColorPlaneSelected)
        self.mMenu.setMenu(menu)
        self.mLayer.layerChanged.connect(self.onLayerChanged)
        self.mBand1.bandChanged.connect(self.onLiveUpdate)
        self.mBand2.bandChanged.connect(self.onLiveUpdate)
        self.mColor1.colorChanged.connect(self.onColorChanged)
        self.mColor2.colorChanged.connect(self.onColorChanged)
        self.mColor3.colorChanged.connect(self.onColorChanged)
        self.mColor4.colorChanged.connect(self.onColorChanged)
        self.mUseColor2.stateChanged.connect(self.onColorChanged)
        self.mMode.currentIndexChanged.connect(self.onColorChanged)
        self.mClasses.valueChanged.connect(self.onColorChanged)
        self.mP1.valueChanged.connect(self.onLiveUpdate)
        self.mP2.valueChanged.connect(self.onLiveUpdate)
        self.mExtent.currentIndexChanged.connect(self.onLiveUpdate)
        self.mAccuracy.currentIndexChanged.connect(self.onLiveUpdate)

        self.mApply.clicked.connect(self.onApplyClicked)

        self.onColorChanged()

    def currentLayer(self) -> Optional[QgsRasterLayer]:
        return self.mLayer.currentLayer()

    def currentExtent(self) -> Optional[SpatialExtent]:
        layer = self.currentLayer()
        if layer is None:
            return None

        if self.mExtent.currentIndex() == self.WholeRasterExtent:
            return SpatialExtent(layer.crs(), layer.extent())
        elif self.mExtent.currentIndex() == self.CurrentCanvasExtent:
            mapSettings: QgsMapSettings = self.mMapCanvas.mapSettings()
            return SpatialExtent(mapSettings.destinationCrs(), self.mMapCanvas.extent()).toCrs(layer.crs())
        else:
            raise ValueError()

    def currentSampleSize(self) -> int:
        if self.mAccuracy.currentIndex() == self.EstimatedAccuracy:
            return int(QgsRasterLayer.SAMPLE_SIZE)
        elif self.mAccuracy.currentIndex() == self.ActualAccuracy:
            return 0  # use all pixel
        else:
            raise ValueError()

    def currentRenderer(self) -> Optional[BivariateColorRasterRenderer]:
        layer = self.currentLayer()
        if layer is None:
            return None

        band1 = self.mBand1.currentBand()
        band2 = self.mBand2.currentBand()
        if -1 in [band1, band2]:
            return None

        quantile_range = self.mP1.value(), self.mP2.value()

        extent = self.currentExtent()
        extent = extent.intersect(layer.extent())
        sampleSize = self.currentSampleSize()

        reader = RasterReader(layer)
        width, height = reader.samplingWidthAndHeight(band1, extent, sampleSize)
        array = reader.array(bandList=[band1, band2], width=width, height=height, boundingBox=extent)
        maskArray = np.all(reader.maskArray(array, [band1, band2]), axis=0)
        array1, array2 = array
        values1 = array1[maskArray]
        values2 = array2[maskArray]

        if len(values1) == 0:
            return None

        if self.mMinMaxPercentile.isChecked():
            min1, max1 = np.percentile(values1, quantile_range)
            min2, max2 = np.percentile(values2, quantile_range)
            self.mMin1.setText(str(min1))
            self.mMin2.setText(str(min2))
            self.mMax1.setText(str(max1))
            self.mMax2.setText(str(max2))

        min1 = float(self.mMin1.text())
        min2 = float(self.mMin2.text())
        max1 = float(self.mMax1.text())
        max2 = float(self.mMax2.text())

        mode = self.mMode.currentIndex()
        if mode in [self.ContinuouseMode, self.EqualIntervalMode]:
            binEdges1 = np.linspace(min1, max1, self.colorPlane.shape[0] + 1)
            binEdges2 = np.linspace(min2, max2, self.colorPlane.shape[1] + 1)
        else:
            binEdges1 = np.percentile(
                values1[(min1 <= values1) * (values1 <= max1)],
                np.linspace(0, 100, self.colorPlane.shape[0] + 1)
            )
            binEdges2 = np.percentile(
                values2[(min2 <= values2) * (values1 <= max2)],
                np.linspace(0, 100, self.colorPlane.shape[1] + 1)
            )

        # update plot
        ax = self.mPlot.getAxis('bottom')
        ay = self.mPlot.getAxis('left')
        if mode == self.ContinuouseMode:
            ax.setTicks([[(0, smartRound(binEdges1[0])), (len(binEdges1), smartRound(binEdges1[-1]))]])
            ay.setTicks([[(0, smartRound(binEdges2[0])), (len(binEdges2), smartRound(binEdges2[-1]))]])
        else:
            ax.setTicks([[(i, smartRound(value)) for i, value in enumerate(binEdges1)]])
            ay.setTicks([[(i, smartRound(value)) for i, value in enumerate(binEdges2)]])
        self.mPlot.setLabel(axis='bottom', text=reader.bandName(band1))
        self.mPlot.setLabel(axis='left', text=reader.bandName(band2))

        # make renderer
        renderer = BivariateColorRasterRenderer()
        renderer.setRange(min1, min2, max1, max2)
        renderer.setBands(band1, band2)
        renderer.setColorPlane(self.colorPlane)
        renderer.setBinEdges(binEdges1, binEdges2)

        return renderer

    def onLayerChanged(self):

        # get map canvas
        self.mMapCanvas = None
        layer = self.currentLayer()
        for mapDock in self.enmapBox.dockManager().mapDocks():
            if layer in mapDock.mapCanvas().layers():
                self.mMapCanvas = mapDock.mapCanvas()
                break

        # update gui
        self.mBand1.setLayer(layer)
        self.mBand2.setLayer(layer)
        self.mBand1.setBand(-1)
        self.mBand2.setBand(-1)

        if layer is None:
            return

        renderer = layer.renderer()
        if isinstance(renderer, BivariateColorRasterRenderer):
            self.mBand1.setBand(renderer.band1)
            self.mBand2.setBand(renderer.band2)
        else:
            self.mBand1.setBand(1)
            self.mBand2.setBand(min(2, layer.bandCount()))

    def onColorChanged(self):
        mode = self.mMode.currentIndex()

        colors = [mColor.color() for mColor in [self.mColor1, self.mColor2, self.mColor3, self.mColor4]]

        if not self.mUseColor2.isChecked():
            color2 = QColor(
                int((colors[0].red() + colors[3].red()) / 2),
                int((colors[0].green() + colors[3].green()) / 2),
                int((colors[0].blue() + colors[3].blue()) / 2),
            )
            colors[1] = color2

        colorsEdges = np.reshape([np.array([color.red(), color.green(), color.blue()]) for color in colors], (2, 2, 3))

        if mode == self.ContinuouseMode:
            self.mClasses.setEnabled(False)
            classes = 256
        else:
            self.mClasses.setEnabled(True)
            classes = self.mClasses.value()

        x, y = np.meshgrid([1, classes], [1, classes])
        Xnew, Ynew = np.meshgrid(np.linspace(1, classes, classes), np.linspace(1, classes, classes))
        points = list(zip(x.flat, y.flat))
        fRed = LinearNDInterpolator(points, list(colorsEdges[:, :, 0].flat))
        redPlane = fRed(Xnew, Ynew)
        fGreen = LinearNDInterpolator(points, list(colorsEdges[:, :, 1].flat))
        greenPlane = fGreen(Xnew, Ynew)
        fBlue = LinearNDInterpolator(points, list(colorsEdges[:, :, 2].flat))
        bluePlane = fBlue(Xnew, Ynew)
        self.colorPlane = np.transpose((redPlane, greenPlane, bluePlane), (1, 2, 0))

        self.mPlot.clear()

        imageItem = ImageItem()
        imageItem.setImage(
            self.colorPlane[::-1], autoHistogramRange=True, autoLevels=True, autoDownsample=True,
            axisOrder='row-major'
        )

        self.mPlot.addItem(imageItem)
        self.onApplyClicked()

    def onColorPlaneSelected(self):
        action: QAction = self.sender()
        name = action.text()
        hexcolors = self.predefinedColorPlanes[name]
        for hexcolor, mColor in zip(hexcolors, [self.mColor1, self.mColor2, self.mColor3, self.mColor4]):
            mColor.setColor(QColor(hexcolor))
        self.onColorChanged()

    def onApplyClicked(self):
        layer = self.currentLayer()
        if layer is None:
            return
        renderer = self.currentRenderer()
        if renderer is None:
            return
        layer.setRenderer(renderer)
        layer.triggerRepaint()

    def onLiveUpdate(self):
        if self.mLiveUpdate.isChecked():
            self.onApplyClicked()


@typechecked
def smartRound(value: float) -> str:
    if abs(value) < 0.1:
        return str(round(value, 4))
    elif abs(value) < 1:
        return str(round(value, 3))
    elif abs(value) < 10:
        return str(round(value, 2))
    elif abs(value) < 100:
        return str(round(value, 1))
    else:
        return str(int(round(value, 0)))
