import numpy as np
from qgis.core import (
    QgsCoordinateTransform,
    QgsFeature,
    QgsFields,
    QgsGeometry,
    QgsMemoryProviderUtils,
    QgsPoint,
    QgsPointXY,
    QgsProject,
    QgsRasterLayer,
    QgsWkbTypes,
)
from qgis.PyQt import QtCore
from qgis.PyQt.QtCore import QModelIndex, QObject, QVariant
from qgis.PyQt.QtGui import QStandardItem, QStandardItemModel
from qgis.utils import iface

from openlog.toolbelt import PlgLogger


class LocalGridLayerModel(QStandardItemModel):
    BASE_X_COL = 0
    BASE_Y_COL = 1
    BASE_Z_COL = 2
    DEST_X_COL = 3
    DEST_Y_COL = 4
    DEST_Z_COL = 5

    def __init__(self, parent: QObject = None):
        """
        QStandardItemModel for collar layer creation.

        Use a temporary memory layer for collar position definition.

        Layer features are synchronized with model values

        Args:
            parent: QObject parent
        """
        super(LocalGridLayerModel, self).__init__(parent)
        self.setHorizontalHeaderLabels(
            [
                "Source Easting",
                "Source Northing",
                "Source Elevation",
                "Destination Easting",
                "Destination Northing",
                "Destination Elevation",
            ]
        )
        self.log = PlgLogger().log

        # Create temporary layer
        self._layer = QgsMemoryProviderUtils.createMemoryLayer(
            name="Local grid creation",
            fields=QgsFields(),
            geometryType=QgsWkbTypes.PointZ,
            crs=iface.mapCanvas().mapSettings().destinationCrs(),
        )

        # Add to project : must be removed by widget using model before delete
        QgsProject().instance().addMapLayer(self._layer)

        # Start edition for immediate collar add
        self._layer.startEditing()
        iface.actionAddFeature().trigger()

        # Synchronize layer modification for model
        self._featureAdded_connect = lambda fid: self._feature_added(fid)
        self._layer.featureAdded.connect(self._featureAdded_connect)

        self._featuresDeleted_connect = lambda fids: self._features_deleted(fids)
        self._layer.featuresDeleted.connect(self._featuresDeleted_connect)

        self._geometryChanged_connect = lambda fid, geometry: self._geometry_changed(
            fid, geometry
        )
        self._layer.geometryChanged.connect(self._geometryChanged_connect)

        # Synchronize model change for layer
        self._itemChanged_connect = lambda item: self._item_changed(item)
        self.itemChanged.connect(self._itemChanged_connect)

        self._dtm_layer = QgsRasterLayer()

    def remove_temp_layer(self) -> None:
        """
        Remove temp layer

        Must be called when model is not used anymore

        """
        if self._layer:
            #  Changes must be committed for no warning when removing layer
            self._layer.commitChanges(True)
            QgsProject().instance().removeMapLayer(self._layer)

    def get_src_array(self) -> np.array:
        """
        Get source array

        Returns: collar list (warning no person and dataset are defined)

        """
        result = np.zeros((self.rowCount(), 3))
        for i in range(0, self.rowCount()):
            result[i][0] = super().data(self.index(i, self.BASE_X_COL))
            result[i][1] = super().data(self.index(i, self.BASE_Y_COL))
            result[i][2] = super().data(self.index(i, self.BASE_Z_COL))
        return result

    def get_dst_array(self) -> np.array:
        """
        Get source array

        Returns: collar list (warning no person and dataset are defined)

        """
        result = np.zeros((self.rowCount(), 3))
        for i in range(0, self.rowCount()):
            result[i][0] = super().data(self.index(i, self.DEST_X_COL))
            result[i][1] = super().data(self.index(i, self.DEST_Y_COL))
            result[i][2] = super().data(self.index(i, self.DEST_Z_COL))
        return result

    def data(
        self, index: QtCore.QModelIndex, role: int = QtCore.Qt.DisplayRole
    ) -> QVariant:
        """
        Override QStandardItemModel data() to display x/y/z without scientific representation

        Args:
            index: QModelIndex
            role: Qt role

        Returns: QVariant

        """
        result = super().data(index, role)
        if role == QtCore.Qt.DisplayRole:
            if isinstance(result, float):
                # Render float to 2 dp
                result = "%.2f" % result
        return result

    def _item_changed(self, item: QStandardItem) -> None:
        """
        Function called when item is changed in model.

        Used to synchronize temporary memory layer (only x, y, z)

        Args:
            item: QStandardItem changed
        """
        col = item.index().column()
        if col == self.BASE_X_COL or col == self.BASE_Y_COL or col == self.BASE_Z_COL:
            row = item.index().row()
            fid = super().data(self.index(row, self.BASE_X_COL), QtCore.Qt.UserRole)
            self._update_layer_feature_geom(
                fid,
                x=super().data(self.index(row, self.BASE_X_COL)),
                y=super().data(self.index(row, self.BASE_Y_COL)),
                z=super().data(self.index(row, self.BASE_Z_COL)),
            )

    def _update_layer_feature_geom(
        self, fid: int, x: float, y: float, z: float
    ) -> None:
        """
        Update temporary layer geometry from collar definition

        Args:
            fid: id of feature
            collar: collar
        """
        # Layer must be editable or we won't be able to change geometry
        if not self._layer.isEditable():
            self._layer.startEditing()
        iface.setActiveLayer(self._layer)

        # Disconnect geometry changes signal to avoid infinite loop
        self._layer.geometryChanged.disconnect(self._geometryChanged_connect)
        self._layer.beginEditCommand("Moved point from local grid creation model")
        self._layer.moveVertexV2(QgsPoint(x=x, y=y, z=z), fid, 0)
        self._layer.endEditCommand()
        self._layer.triggerRepaint()

        # Restore synchronization for geometry change
        self._layer.geometryChanged.connect(self._geometryChanged_connect)

    def removeRow(self, row: int, parent: QModelIndex = QModelIndex()) -> bool:
        """
        Override default implementation of QStandardItemModel to remove feature from associated layer

        Args:
            row: (int) row to remove
            parent: (QModelIndex) parent

        Returns: True if row was removed, False otherwise

        """
        # Layer must be editable or we won't be able to change geometry
        if not self._layer.isEditable():
            self._layer.startEditing()
        iface.setActiveLayer(self._layer)

        fid = super().data(self.index(row, self.BASE_X_COL), QtCore.Qt.UserRole)

        # Disconnect geometry changes signal to avoid infinite loop
        self._layer.featuresDeleted.disconnect(self._featuresDeleted_connect)

        self._layer.beginEditCommand("Deleted point from local grid creation model")
        self._layer.deleteFeature(fid)
        self._layer.endEditCommand()
        self._layer.triggerRepaint()

        # Restore synchronization for geometry delete
        self._layer.featuresDeleted.connect(self._featuresDeleted_connect)

        return super().removeRow(row, parent)

    def _features_deleted(self, fids: [int]) -> None:
        """
        Synchronize feature removal in temporary memory layer in model

        Args:
            fids: list of feature id removed
        """
        for i in reversed(range(0, self.rowCount())):
            fid = super().data(self.index(i, self.BASE_X_COL), QtCore.Qt.UserRole)
            if fid in fids:
                super().removeRow(i)

    def _geometry_changed(self, fid: int, geometry: QgsGeometry):
        """
        Synchronize feature geometry changes in temporary memory layer in model

        Args:
            fid: id of feature changed
            geometry: new feature geometry
        """

        # Update only if geometry is defined
        if geometry.constGet():
            for i in range(0, self.rowCount()):
                # Get feature id of row from UserRole
                fid_row = super().data(
                    self.index(i, self.BASE_X_COL), QtCore.Qt.UserRole
                )
                if fid_row == fid:
                    # Disconnect item changes signal to avoid infinite loop
                    self.itemChanged.disconnect(self._itemChanged_connect)
                    # Update data with new geometry
                    self.setData(
                        self.index(i, self.BASE_X_COL), geometry.constGet().x()
                    )
                    self.setData(
                        self.index(i, self.BASE_Y_COL), geometry.constGet().y()
                    )
                    self.setData(
                        self.index(i, self.BASE_Z_COL), geometry.constGet().z()
                    )
                    # Restore synchronization for item change
                    self.itemChanged.connect(self._itemChanged_connect)
                    break

    def _feature_added(self, fid: int):
        """
        Synchronize feature added in temporary memory layer in model

        Args:
            fid: new feature id
        """
        if self._layer:
            self._add_feature(self._layer.getFeature(fid))

    def _add_feature(self, feature: QgsFeature):
        """
        Add row in model from QgsFeature

        Args:
            feature: QgsFeature added
        """
        self.insertRow(self.rowCount())
        row = self.rowCount() - 1

        point = feature.geometry().asPoint()

        # Disconnect item changes signal to avoid infinite loop
        self.itemChanged.disconnect(self._itemChanged_connect)

        self.setData(self.index(row, self.BASE_X_COL), feature.id(), QtCore.Qt.UserRole)

        self.setData(self.index(row, self.BASE_X_COL), point.x())
        self.setData(self.index(row, self.BASE_Y_COL), point.y())

        # Get value from DTM here
        if self._dtm_layer and self._dtm_layer.isValid():
            z = self._get_z_value_from_dtm(point)
        else:
            z = 0.0

        self.setData(self.index(row, self.BASE_Z_COL), z)

        # Restore synchronization for item change
        self.itemChanged.connect(self._itemChanged_connect)

    def _get_z_value_from_dtm(self, point):
        tr = QgsCoordinateTransform(
            self._layer.crs(), self._dtm_layer.crs(), QgsProject.instance()
        )
        point_dtm = tr.transform(point)
        z_val, res = self._dtm_layer.dataProvider().sample(point_dtm, 1)
        if res:
            z = z_val
        else:
            self.log(
                "Can't define DTM value for point : {0}/{1}. z value used is 0.0.".format(
                    point.x(), point.y()
                ),
                push=True,
            )
            z = 0.0
        return z

    def set_dtm_layer(self, layer: QgsRasterLayer) -> None:
        self._dtm_layer = layer

        if self._dtm_layer:
            # Disconnect item changes signal to avoid infinite loop
            self.itemChanged.disconnect(self._itemChanged_connect)
            for i in range(0, self.rowCount()):
                x = super().data(self.index(i, self.BASE_X_COL))
                y = super().data(self.index(i, self.BASE_Y_COL))
                point = QgsPointXY(x, y)
                z = self._get_z_value_from_dtm(point)
                self.setData(self.index(i, self.BASE_Z_COL), z)
            # Restore synchronization for item change
            self.itemChanged.connect(self._itemChanged_connect)
