from PyQt5 import QtWidgets, QtGui, QtCore
from typing import Tuple, TYPE_CHECKING, List
from copy import deepcopy

from ..coordTransform import *
if TYPE_CHECKING:
    from ..sampleview import SampleView
    from ..instrumentcom.instrumentComBase import InstrumentComBase
    from ..dataset import DataSet


class CoordTransformUI(QtWidgets.QWidget):
    ImageMarkerAdded: QtCore.pyqtSignal = QtCore.pyqtSignal(tuple)
    ImageMarkerDeleted: QtCore.pyqtSignal = QtCore.pyqtSignal(int)
    CenterOnMarker: QtCore.pyqtSignal = QtCore.pyqtSignal(int)
    WindowClosed: QtCore.pyqtSignal = QtCore.pyqtSignal()

    def __init__(self, viewParent: 'SampleView'):
        super(CoordTransformUI, self).__init__()
        self.setWindowTitle('Recalculate Coordinate System')
        self._viewParent: 'SampleView' = viewParent
        self._dataset: 'DataSet' = viewParent.dataset
        assert self._viewParent is not None
        self._instrctrl: 'InstrumentComBase' = self._viewParent.instrctrl
        self.coordTransf: CoordTransfer = CoordTransfer()
        self._transformNotYetSaved: bool = True

        self._imgMarkersSrc: List[ImageMarker] = []
        self._imgMarkersDst: List[ImageMarker] = []  # in the new coord system
        self._trayMarkersSrc: List[TrayMarker] = []
        self._trayMarkersDst: List[TrayMarker] = []  # in the new coord system

        self._layout = QtWidgets.QVBoxLayout()
        self.setLayout(self._layout)

        selectionLayout = QtWidgets.QHBoxLayout()
        self._trayMarkersBtn: QtWidgets.QRadioButton = QtWidgets.QRadioButton('Use Tray Markers')
        self._trayMarkersBtn.pressed.connect(self._updateTrayMarkersGUI)
        self._trayMarkersBtn.setChecked(False)
        self._imgMarkersBtn: QtWidgets.QRadioButton = QtWidgets.QRadioButton('Use Image Markers')
        self._imgMarkersBtn.pressed.connect(self._updateImageMarkersGUI)
        self._imgMarkersBtn.setChecked(True)

        selectionLayout.addWidget(self._trayMarkersBtn)
        selectionLayout.addWidget(self._imgMarkersBtn)

        self._addMarkerBtn: QtWidgets.QPushButton = QtWidgets.QPushButton('Add New Markers')
        self._addMarkerBtn.setStyleSheet("background-color : lightgrey")
        self._formatAsToggleBtn(self._addMarkerBtn)

        self._transferResultLbl: QtWidgets.QLabel = QtWidgets.QLabel('No valid input')

        self._trayMarkerGroup: QtWidgets.QGroupBox = self._getTrayMarkerGroup()
        self._imgMarkerGroup: QtWidgets.QGroupBox = self._getImageMarkerGroup()

        self._invertGroup: 'InvertAxisGroup' = InvertAxisGroup()
        self._invertGroup.CheckBoxToggled.connect(self._calculateCoordinateTransfer)

        saveBtn: QtWidgets.QPushButton = QtWidgets.QPushButton('Save to Dataset')
        saveBtn.released.connect(self._saveToDataset)
        btnLayout: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        btnLayout.addWidget(saveBtn)
        btnLayout.addStretch()

        self._layout.addLayout(selectionLayout)
        self._layout.addWidget(self._trayMarkerGroup)
        self._layout.addWidget(self._invertGroup)
        self._layout.addWidget(self._transferResultLbl)
        self._layout.addLayout(btnLayout)

        self._updateFromDataset()
        self._updateImageMarkersGUI()

    def _updateFromDataset(self) -> None:
        self._imgMarkersSrc = deepcopy(self._dataset.imageMarkers)
        for imgMarker in self._imgMarkersSrc:
            newImgMarker: ImageMarker = ImageMarker()
            newImgMarker.index = imgMarker.index
            newImgMarker.imgCoordX = imgMarker.imgCoordX
            newImgMarker.imgCoordY = imgMarker.imgCoordY

            self._imgMarkersDst.append(newImgMarker)

        self._trayMarkersSrc = deepcopy(self._dataset.trayMarkers)
        for _ in range(len(self._trayMarkersSrc)):
            newTrayMarker: TrayMarker = TrayMarker()
            self._trayMarkersDst.append(newTrayMarker)

    def _formatAsToggleBtn(self, btn: QtWidgets.QPushButton) -> None:
        btn.setCheckable(True)
        btn.setChecked(False)
        btn.clicked.connect(lambda: self._setBtnColor(btn))

    def _getTrayMarkerGroup(self) -> QtWidgets.QGroupBox:
        layout = QtWidgets.QGridLayout()
        group: QtWidgets.QGroupBox = QtWidgets.QGroupBox('Sample Tray Markers')
        group.setLayout(layout)
        layout.addWidget(QtWidgets.QLabel("NOT YET IMPLEMENTED"))
        return group

    def _getImageMarkerGroup(self) -> QtWidgets.QGroupBox:
        def makeDeleteLambda(ind):
            return lambda: self._deleteImageMarker(ind)
        
        def makeReadNewPosLambda(ind):
            return lambda: self._readNewImageMarker(ind)

        def makeCenterOnLambda(ind):
            return lambda: self.CenterOnMarker.emit(ind)

        layout = QtWidgets.QGridLayout()
        group: QtWidgets.QGroupBox = QtWidgets.QGroupBox('Image Based Markers')
        group.setLayout(layout)
        layout.addWidget(self._addMarkerBtn)
        if len(self._imgMarkersSrc) == 0:
            layout.addWidget(QtWidgets.QLabel('No image markers found.'))
        else:
            layout.addWidget(QtWidgets.QLabel('Number'), 1, 0)
            layout.addWidget(QtWidgets.QLabel('Old Coordinates'), 1, 1)
            layout.addWidget(QtWidgets.QLabel('New Coordinates'), 1, 2)

            for index, imgMarkerSrc in enumerate(self._imgMarkersSrc):
                newMarkerDst: ImageMarker = self._imgMarkersDst[index]
                newX, newY, newZ = newMarkerDst.worldCoordX, newMarkerDst.worldCoordY, newMarkerDst.worldCoordZ
                x, y, z = imgMarkerSrc.worldCoordX, imgMarkerSrc.worldCoordY, imgMarkerSrc.worldCoordZ
                if np.any(np.isnan([x, y, z])):
                    x, y, z = 'Not defined', 'Not defined', 'Not defined'
                else:
                    x, y, z = round(x), round(y), round(z)

                row: int = 2+index
                updateBtn: QtWidgets.QPushButton = QtWidgets.QPushButton('Read Coordinates')
                updateBtn.clicked.connect(makeReadNewPosLambda(index))
                focusBtn: QtWidgets.QPushButton = QtWidgets.QPushButton('Center On')
                focusBtn.clicked.connect(makeCenterOnLambda(index))
                delBtn: QtWidgets.QPushButton = QtWidgets.QPushButton('Delete')
                delBtn.clicked.connect(makeDeleteLambda(index))
                layout.addWidget(QtWidgets.QLabel(str(imgMarkerSrc.index+1)), row, 0)
                layout.addWidget(QtWidgets.QLabel(f'x: {x} µm\ny: {y} µm\nz: {z} µm'), row, 1)
                if np.nan in [newX, newY, newZ]:
                    layout.addWidget(QtWidgets.QLabel('Coordinates not yet read.'), row, 2)
                else:
                    layout.addWidget(QtWidgets.QLabel(f'x: {round(newX)} µm\ny: {round(newY)} µm\nz: {round(newZ)} µm'), row, 2)
                layout.addWidget(updateBtn, row, 3)
                layout.addWidget(focusBtn, row, 4)
                layout.addWidget(delBtn, row, 5)
        return group

    def _resetMarkerIndices(self) -> None:
        for index, (imgMarker, newImgMarker) in enumerate(zip(self._imgMarkersSrc, self._imgMarkersDst)):
            imgMarker.index = index
            newImgMarker.index = index

    def _calculateCoordinateTransfer(self) -> None:
        dstPoints, srcPoints = self._getSourceAndDestinationPoints()
        invertX, invertY, invertZ = self._invertGroup.getInvert()

        if len(srcPoints) == len(dstPoints) and len(srcPoints) >= 3 and np.all(np.isfinite(dstPoints)):
            if not self._dataset.coordinateSystemConfigured:
                self._setOriginalCoordSystem()
                srcPoints = self._setAndGetCreatedImgMarkerSourcePoints()

            if np.any(np.isnan(srcPoints)):
                self._viewParent.getLogger().warning("Corrdinate transfer cannot be done, source markers contain NaN")
                return

            try:
                self.coordTransf: CoordTransfer = getTransform(srcPoints, dstPoints, invertX, invertY, invertZ)
                errorString: str = "Transform errors:"
                for i, error in enumerate(self.coordTransf.error, start=1):
                    errorString += f" Marker {i}: {round(error, 1)} µm,"
                self._transferResultLbl.setText(errorString)
                self._transformNotYetSaved = True
            except ValueError as e:
                QtWidgets.QMessageBox.critical(self, "Transform failed", f"The following error occured:\n\n{e}\n\n"
                                                                         f"Please try slightly different coordinates.")
                self._dataset.invalidateCoordinateSystem()
        else:
            self._transferResultLbl.setText('No valid inputs for coordinate transfer')

    def _getSourceAndDestinationPoints(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Retrieves the currently defined point coordinates for source and destination coordinate system.
        Returns tuple of srcPoints (Nx3) array and dstPoints (Nx3) array of (x, y, z) coordinates for N markers.
        """
        oldMarkers: list = self._trayMarkersSrc if self._trayMarkersBtn.isChecked() else self._imgMarkersSrc
        newMarkers: list = self._trayMarkersDst if self._trayMarkersBtn.isChecked() else self._imgMarkersDst
        srcPoints: np.ndarray = np.array(
            [[marker.worldCoordX, marker.worldCoordY, marker.worldCoordZ] for marker in oldMarkers])
        dstPoints: np.ndarray = np.array(
            [[marker.worldCoordX, marker.worldCoordY, marker.worldCoordZ] for marker in newMarkers])
        return dstPoints, srcPoints

    def _setOriginalCoordSystem(self) -> None:
        """
        In case no original coordinate system exists yet, a default one is created as a base for the coordinate
        transfer.
        """
        if not self._imgMarkersBtn.isChecked():
            QtWidgets.QMessageBox.critical(self, "Error", "A new coordinate system has to be created."
                                                          "This only works for image markers, please check the"
                                                          "according button!")
            return

        self._dataset.pshift = self._instrctrl.getSpectrumPositionShift()
        self._dataset.signx = 1.0
        self._dataset.signy = -1.0
        pxScale: float = getPixelScaleFromImageMarkers(self._imgMarkersDst)
        self._dataset.pixelscale_df = self._dataset.pixelscale_bf = pxScale
        self._dataset.calculateImageDimsFromZImg()

        dstPointsWorld: np.ndarray = np.array(
            [[marker.worldCoordX, marker.worldCoordY, marker.worldCoordZ] for marker in self._imgMarkersDst])

        self._dataset.fitpoints = dstPointsWorld
        self._dataset.calculateHeightmapParams(self._instrctrl.getSoftwareZ(), self._instrctrl.getUserZ())

        imgDim: np.ndarray = self._dataset.imagedim_df
        signx, signy = self._dataset.signx, self._dataset.signy
        self._dataset.lastpos = np.array([imgDim[0]/2 * signx, imgDim[1]/2 * signy])
        self._dataset.coordinatetransform = None

    def _setAndGetCreatedImgMarkerSourcePoints(self) -> np.ndarray:
        """
        Fake (but realistic in terms of scale) source coordinates are set for the image markers. We just assume
        world (0, 0, 0) to align with image (0, 0).
        All offsets and rotations will then be handled by the actual coordinate transfer.
        """
        for imgMarker in self._imgMarkersSrc:
            x, y, z = self._dataset.mapToLengthWithoutCoordTransfer((imgMarker.imgCoordX, imgMarker.imgCoordY))
            imgMarker.worldCoordX = x
            imgMarker.worldCoordY = y
            imgMarker.worldCoordZ = z

        srcPoints = np.array([[marker.worldCoordX, marker.worldCoordY, marker.worldCoordZ] for marker in self._imgMarkersSrc])
        assert not np.any(np.isnan(srcPoints))
        self._updateImageMarkersGUI(autoUpdateCoordTransfer=False)
        return srcPoints

    def _setBtnColor(self, btn: QtWidgets.QPushButton):
        """
        Sets color of the toggle button to indicate its status.
        """
        if btn.isChecked():
            btn.setStyleSheet("background-color : lightblue")
        else:
            btn.setStyleSheet("background-color : lightgrey")

    def _updateTrayMarkersGUI(self) -> None:
        """
        Use markers on the sample tray that are NOT in the microscope image.
        """
        self._imgMarkerGroup.setParent(None)
        self._trayMarkerGroup.setParent(None)
        self._trayMarkerGroup = self._getTrayMarkerGroup()
        self._layout.insertWidget(1, self._trayMarkerGroup)

    def _updateImageMarkersGUI(self, autoUpdateCoordTransfer: bool = True) -> None:
        """
        Use markers that are defined through characteristic spots within the microscope image.
        :param autoUpdateCoordTransfer: Also update the coordinate transfer
        """
        self._imgMarkerGroup.setParent(None)
        self._trayMarkerGroup.setParent(None)
        self._imgMarkerGroup = self._getImageMarkerGroup()
        self._layout.insertWidget(1, self._imgMarkerGroup)
        if autoUpdateCoordTransfer:
            self._calculateCoordinateTransfer()

    def mousePressed(self, event: QtGui.QMouseEvent) -> bool:
        wasConsumed: bool = False
        if self._addMarkerBtn.isChecked() and event.button() == QtCore.Qt.LeftButton:
            scenePos = self._viewParent.mapToScene(event.pos())
            self._addNewImageMarker((scenePos.x(), scenePos.y()))
            wasConsumed = True
        return wasConsumed

    def _deleteImageMarker(self, markerIndex: int) -> None:
        self._imgMarkersSrc.remove(self._imgMarkersSrc[markerIndex])
        self._imgMarkersDst.remove(self._imgMarkersDst[markerIndex])
        self.ImageMarkerDeleted.emit(markerIndex)
        self._resetMarkerIndices()
        self._updateImageMarkersGUI()

    def _addNewImageMarker(self, pos: tuple) -> None:
        newImageMarkerScr: ImageMarker = ImageMarker()
        newImageMarkerScr.index = len(self._imgMarkersSrc)
        newImageMarkerScr.imgCoordX = pos[0]
        newImageMarkerScr.imgCoordY = pos[1]

        try:
            x, y, z = self._dataset.mapToLengthWithoutCoordTransfer(pos, self._viewParent.microscopeMode)
            newImageMarkerScr.worldCoordX = x
            newImageMarkerScr.worldCoordY = y
            newImageMarkerScr.worldCoordZ = z
        except TypeError:
            pass  # getting world coordinates was not possible due to not yet configured coordinate system
        self._imgMarkersSrc.append(newImageMarkerScr)

        newImageMarkerDst: ImageMarker = ImageMarker()
        newImageMarkerDst.imgCoordX = pos[0]
        newImageMarkerDst.imgCoordY = pos[1]
        self._imgMarkersDst.append(newImageMarkerDst)

        self.ImageMarkerAdded.emit(pos)
        self._updateImageMarkersGUI()

    def _readNewImageMarker(self, index: int) -> None:
        """Reads current instrument coordinates and assignes them to the respective "new image marker" """
        x, y, z = self._instrctrl.getPosition()
        self._imgMarkersDst[index].worldCoordX = x
        self._imgMarkersDst[index].worldCoordY = y
        self._imgMarkersDst[index].worldCoordZ = z
        self._updateImageMarkersGUI()

    def _saveToDataset(self) -> None:
        if self._transformNotYetSaved:
            if self._imgMarkersBtn.isChecked():
                self._dataset.saveNewImageMarkers(self._imgMarkersSrc)
                self._viewParent.logger.debug(f"saved image markers to dataset")
            else:
                raise NotImplementedError

            save: bool = True
            if np.any(self.coordTransf.error > 10.0):
                reply = QtWidgets.QMessageBox.question(self, "Warning", f"Coordinate transfer residues are large ({np.round(self.coordTransf.error)} µm)! \n"
                                                                        "Really save the coordinte transfer result?",
                                                       QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No,
                                                       QtWidgets.QMessageBox.No)

                save = reply == QtWidgets.QMessageBox.Yes

            if save:
                self._dataset.coordinatetransform = self.coordTransf
                self._viewParent.logger.info(f"Saved coordinate transfer with rotMat {self.coordTransf.rotMatrix},"
                                             f"offset {self.coordTransf.offset} and center {self.coordTransf.rotationCenter}."
                                             f"Errors are: {self.coordTransf.error}")
                self._transformNotYetSaved = False

    def closeEvent(self, a0: QtGui.QCloseEvent) -> None:
        reply = QtWidgets.QMessageBox.question(self, 'Close without saving',
                                               'The window is about to be closed.\nSave markers prior to closing?',
                                               QtWidgets.QMessageBox.Yes | QtWidgets.QMessageBox.No,
                                               QtWidgets.QMessageBox.Yes)
        if reply == QtWidgets.QMessageBox.Yes:
            self._saveToDataset()

        self.WindowClosed.emit()

        del self


class InvertAxisGroup(QtWidgets.QGroupBox):
    CheckBoxToggled: QtCore.pyqtSignal = QtCore.pyqtSignal()

    def __init__(self):
        super(InvertAxisGroup, self).__init__()
        self._invertX: QtWidgets.QCheckBox = QtWidgets.QCheckBox()
        self._invertY: QtWidgets.QCheckBox = QtWidgets.QCheckBox()
        self._invertZ: QtWidgets.QCheckBox = QtWidgets.QCheckBox()

        for checkbox in [self._invertX, self._invertY, self._invertZ]:
            checkbox.stateChanged.connect(lambda: self.CheckBoxToggled.emit())

        layout: QtWidgets.QHBoxLayout = QtWidgets.QHBoxLayout()
        layout.addWidget(QtWidgets.QLabel("Invert X:"))
        layout.addWidget(self._invertX)
        layout.addWidget(QtWidgets.QLabel(", Invert Y:"))
        layout.addWidget(self._invertY)
        layout.addWidget(QtWidgets.QLabel(", Invert Z:"))
        layout.addWidget(self._invertZ)
        layout.addStretch()

        self.setLayout(layout)

    def getInvert(self) -> Tuple[bool, bool, bool]:
        return self._invertX.isChecked(), self._invertY.isChecked(), self._invertZ.isChecked()