import math
import json

from qgis.core import (
    QgsCoordinateTransform,
    QgsCoordinateTransformContext,
    QgsPointXY,
    QgsProcessing,
    QgsProcessingParameterBand,
    QgsProcessingParameterNumber,
    QgsProcessingParameterVectorLayer,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterString,
    QgsProcessingParameterFeatureSource,
    QgsGeometry,
)

from PreCourlis.processing.precourlis_algorithm import PreCourlisAlgorithm


class ImportLayerFromDemAlgorithm(PreCourlisAlgorithm):
    INPUT = "INPUT"
    LAYER_NAME = "LAYER_NAME"
    DEM = "DEM"
    MASK = "MASK"
    POINTS_FILTER = "POINTS_FILTER"
    BAND = "BAND"
    DEFAULT_ELEVATION = "DEFAULT_ELEVATION"

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterVectorLayer(
                self.INPUT,
                self.tr("Input"),
                types=[QgsProcessing.TypeVectorLine],
                defaultValue=None,
            )
        )
        self.addParameter(
            QgsProcessingParameterString(
                self.LAYER_NAME, self.tr("Layer name"), defaultValue=None
            )
        )
        self.addParameter(
            QgsProcessingParameterRasterLayer(
                self.DEM, self.tr("Digital Elevation Model"), defaultValue=None
            )
        )
        self.addParameter(
            QgsProcessingParameterFeatureSource(
                self.MASK,
                self.tr("Mask"),
                types=[QgsProcessing.TypeVectorPolygon],
                optional=True,
            )
        )
        self.addParameter(
            QgsProcessingParameterString(
                self.POINTS_FILTER,
                self.tr(
                    "Points filter (formatted as a json array or arrays : [ [sec_id, p_id], ... ])"
                ),
                optional=True,
            )
        )
        self.addParameter(
            QgsProcessingParameterNumber(
                self.DEFAULT_ELEVATION,
                self.tr(
                    "Default elevation (the default value when there is no Z value to extract)"
                ),
                QgsProcessingParameterNumber.Double,
                optional=True,
            )
        )
        self.addParameter(
            QgsProcessingParameterBand(
                self.BAND,
                self.tr("Band number"),
                1,
                "RASTER",
            )
        )

    def sample_raster(self, raster, point, band, default_elevation):
        value, ok = raster.dataProvider().sample(point, band)

        if math.isnan(value):
            # Try average value from neighboring pixels
            pixel_x = raster.rasterUnitsPerPixelX()
            pixel_y = raster.rasterUnitsPerPixelY()
            values = []
            for dx in (-pixel_x, 0, +pixel_x):
                for dy in (-pixel_y, 0, +pixel_y):
                    neighbourg = QgsPointXY(point.x() + dx, point.y() + dy)
                    value, ok = raster.dataProvider().sample(neighbourg, band)
                    if not math.isnan(value):
                        values.append(value)
            if len(values) > 0:
                value = math.fsum(values) / len(values)

        if math.isnan(value):
            if default_elevation is not None:
                value = default_elevation
            else:
                value = "NULL"

        return value

    def processAlgorithm(self, parameters, context, feedback):
        layer = self.parameterAsVectorLayer(parameters, self.INPUT, context)
        layer_name = self.parameterAsString(parameters, self.LAYER_NAME, context)
        raster = self.parameterAsRasterLayer(parameters, self.DEM, context)
        mask = self.parameterAsSource(parameters, self.MASK, context)
        points_filter = self.parameterAsString(parameters, self.POINTS_FILTER, context)
        band = self.parameterAsInt(parameters, self.BAND, context)
        if parameters.get(self.DEFAULT_ELEVATION, None) is not None:
            default_elevation = self.parameterAsDouble(
                parameters, self.DEFAULT_ELEVATION, context
            )
        else:
            default_elevation = None

        dest_field_index = layer.fields().indexFromName(layer_name)

        tr_context = QgsCoordinateTransformContext()
        transform = QgsCoordinateTransform(layer.crs(), raster.crs(), tr_context)

        has_filtering_points = len(points_filter) > 0
        if has_filtering_points:
            filtering_points = [
                [sec_id, p_id] for sec_id, p_id in json.loads(points_filter)
            ]
        else:
            filtering_points = []

        # This algorithm directly edit input layer in place using edit buffer
        layer.startEditing()
        layer.beginEditCommand("Set values from DEM")

        total = 100.0 / layer.featureCount() if layer.featureCount() else 0
        for current, f in enumerate(layer.getFeatures()):
            # Stop the algorithm if cancel button has been clicked
            if feedback.isCanceled():
                break

            line = next(f.geometry().constParts()).clone()
            values = []
            for p_id, current_value, point in zip(
                f.attribute("p_id").split(","),
                f.attribute(layer_name).split(","),
                line.points(),
            ):
                p_id = int(p_id)
                sec_id = f.attribute("sec_id")
                # Check if the point pass the filter
                is_point_matching = (
                    not has_filtering_points or [sec_id, p_id] in filtering_points
                )

                # Check if the point is inside the mask/polygon
                does_elevation_needs_to_be_computed = is_point_matching and (
                    mask is None
                    or any(
                        [
                            feature.geometry().contains(QgsGeometry(point.clone()))
                            for feature in mask.getFeatures()
                        ]
                    )
                )

                if does_elevation_needs_to_be_computed:
                    tr_point = transform.transform(QgsPointXY(point))
                    values.append(
                        self.sample_raster(raster, tr_point, band, default_elevation)
                    )
                else:
                    values.append(current_value)

            value = ",".join([str(v) for v in values])
            layer.changeAttributeValue(f.id(), dest_field_index, value)

            # Update the progress bar
            feedback.setProgress(int(current * total))

        layer.endEditCommand()

        return {}

    def name(self):
        return "import_layer_from_dem"

    def displayName(self):
        return self.tr("Import layer from DEM")

    def shortHelpString(self):
        return self.tr(
            "This algorithm import/extract the elevation value from a Digital Elevation Model"
            " (DEM).\n\n"
            + "We can specify the area of interest with a list of points and/or a mask."
        )

    def group(self):
        return self.tr("Import")

    def groupId(self):
        return "Import"

    def createInstance(self):
        return ImportLayerFromDemAlgorithm()
