# -*- coding: utf-8 -*-
"""
generate_rem.py

🌄 Generate Relative Elevation Model (REM)

Takes a projected DEM and a clean river line and produces a REM
showing height above the nearest river channel.

Pipeline:
    1. Points along geometry (regular interval along river)
    2. Sample DEM elevation at each point
    3. IDW interpolation at coarse resolution → river surface
    4. Warp river surface to DEM resolution/extent
    5. REM = DEM − river surface  (clamped at 0)
    6. Optional percentile clip (removes bridge/culvert spikes)
"""

import os
import tempfile
import uuid

from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (
    QgsProcessing,
    QgsProcessingAlgorithm,
    QgsProcessingException,
    QgsProcessingParameterRasterLayer,
    QgsProcessingParameterVectorLayer,
    QgsProcessingParameterNumber,
    QgsProcessingParameterRasterDestination,
    QgsProcessingParameterVectorDestination,
    QgsProcessingMultiStepFeedback,
    QgsProcessingLayerPostProcessorInterface,
    QgsVectorLayer,
    QgsRasterLayer,
)
import processing


# ------------------------------------------------------------------ #
#  Post-processor: apply the bundled REM colour ramp on load          #
# ------------------------------------------------------------------ #

class _REMStylePostProcessor(QgsProcessingLayerPostProcessorInterface):
    """
    Applies the bundled REM colour ramp (styles/rem_style.qml) to the
    output REM layer when it is added to the project, so the user gets
    the publication-ready 0-10 m ramp automatically.
    """
    instance = None

    def postProcessLayer(self, layer, context, feedback):
        if not isinstance(layer, QgsRasterLayer):
            return
        style_path = os.path.join(
            os.path.dirname(os.path.dirname(__file__)),
            "styles", "rem_style.qml",
        )
        if os.path.exists(style_path):
            layer.loadNamedStyle(style_path)
            layer.triggerRepaint()

    @staticmethod
    def create():
        _REMStylePostProcessor.instance = _REMStylePostProcessor()
        return _REMStylePostProcessor.instance


class GenerateREMAlgorithm(QgsProcessingAlgorithm):

    INPUT_DEM             = "INPUT_DEM"
    RIVER_LINE            = "RIVER_LINE"
    SAMPLE_INTERVAL       = "SAMPLE_INTERVAL"
    IDW_RESOLUTION_FACTOR = "IDW_RESOLUTION_FACTOR"
    IDW_NEIGHBOURS        = "IDW_NEIGHBOURS"
    IDW_POWER             = "IDW_POWER"
    CLIP_PERCENTILE       = "CLIP_PERCENTILE"
    OUTPUT_REM            = "OUTPUT_REM"
    OUTPUT_RIVER_POINTS   = "OUTPUT_RIVER_POINTS"

    def tr(self, string):
        return QCoreApplication.translate("GenerateREMAlgorithm", string)

    def createInstance(self):
        return GenerateREMAlgorithm()

    def name(self):
        return "generaterem"

    def displayName(self):
        return self.tr("🌄 Generate Relative Elevation Model (REM)")

    def group(self):
        return self.tr("REM Tools")

    def groupId(self):
        return "remtools"

    def shortHelpString(self):
        return self.tr("""
        <h3>🌄 Generate Relative Elevation Model (REM)</h3>
        <p>Produces a raster showing how high every terrain cell sits
        <b>above the nearest river channel</b>. Floodplain features
        (levees, palaeochannels, terraces) become clearly visible once
        the regional slope is removed.</p>

        <h4>Pipeline</h4>
        <ol>
          <li>Points along river geometry at regular interval</li>
          <li>Sample DEM elevation at each point</li>
          <li>IDW interpolation → smooth river surface raster</li>
          <li>Warp river surface to DEM resolution and extent</li>
          <li>REM = DEM − river surface (negative values clamped to 0)</li>
          <li>Optional percentile clip to remove bridge/culvert spikes</li>
        </ol>

        <h4>Tips</h4>
        <ul>
          <li><b>Sample interval</b>: ~5× DEM pixel size is a good start.
              Too large = blocky REM; too small = slow IDW.</li>
          <li><b>IDW resolution factor</b>: 10 runs IDW at 10× the DEM
              pixel size — much faster, barely any quality loss since
              the river surface is a smooth trend.</li>
          <li><b>Percentile clip</b>: 98 works well for most datasets.
              Set to 0 to disable.</li>
          <li>Apply <i>styles/rem_style.qml</i> for a publication-ready
              colour ramp.</li>
          <li>Use <b>〰 Create River Line from DEM</b> first if you need
              to auto-extract the river network.</li>
        </ul>
        """)

    def initAlgorithm(self, config=None):

        self.addParameter(QgsProcessingParameterRasterLayer(
            self.INPUT_DEM,
            self.tr("Input DEM (projected, metre units)"),
        ))

        self.addParameter(QgsProcessingParameterVectorLayer(
            self.RIVER_LINE,
            self.tr("River / channel line"),
            types=[QgsProcessing.TypeVectorLine],
        ))

        self.addParameter(QgsProcessingParameterNumber(
            self.SAMPLE_INTERVAL,
            self.tr("Sample interval along river (metres) — ~5× DEM pixel size recommended"),
            type=QgsProcessingParameterNumber.Double,
            defaultValue=50.0, minValue=1.0,
        ))

        p = QgsProcessingParameterNumber(
            self.IDW_RESOLUTION_FACTOR,
            self.tr("IDW resolution factor (10 = run IDW at 10× DEM pixel size, much faster)"),
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=10, minValue=1, maxValue=50,
        )
        p.setFlags(p.flags() | p.FlagAdvanced)
        self.addParameter(p)

        p = QgsProcessingParameterNumber(
            self.IDW_NEIGHBOURS,
            self.tr("IDW: number of nearest neighbours"),
            type=QgsProcessingParameterNumber.Integer,
            defaultValue=12, minValue=3, maxValue=64,
        )
        p.setFlags(p.flags() | p.FlagAdvanced)
        self.addParameter(p)

        p = QgsProcessingParameterNumber(
            self.IDW_POWER,
            self.tr("IDW: distance weighting power"),
            type=QgsProcessingParameterNumber.Double,
            defaultValue=2.0, minValue=0.5, maxValue=10.0,
        )
        p.setFlags(p.flags() | p.FlagAdvanced)
        self.addParameter(p)

        p = QgsProcessingParameterNumber(
            self.CLIP_PERCENTILE,
            self.tr("Clip REM above percentile (0 = no clip, 98 recommended for display)"),
            type=QgsProcessingParameterNumber.Double,
            defaultValue=98.0, minValue=0.0, maxValue=100.0,
        )
        p.setFlags(p.flags() | p.FlagAdvanced)
        self.addParameter(p)

        self.addParameter(QgsProcessingParameterRasterDestination(
            self.OUTPUT_REM,
            self.tr("Output REM raster"),
        ))

        p = QgsProcessingParameterVectorDestination(
            self.OUTPUT_RIVER_POINTS,
            self.tr("Output river elevation sample points (optional QC layer)"),
            type=QgsProcessing.TypeVectorPoint,
            optional=True, createByDefault=False,
        )
        self.addParameter(p)

    def processAlgorithm(self, parameters, context, feedback):

        fb = QgsProcessingMultiStepFeedback(7, feedback)

        # ---- Read parameters ------------------------------------------
        dem_layer = self.parameterAsRasterLayer(parameters, self.INPUT_DEM, context)
        if not dem_layer or not dem_layer.isValid():
            raise QgsProcessingException("Invalid DEM layer.")

        river_layer = self.parameterAsVectorLayer(parameters, self.RIVER_LINE, context)
        if river_layer is None:
            raise QgsProcessingException("Invalid river line layer.")

        sample_interval = self.parameterAsDouble(parameters, self.SAMPLE_INTERVAL, context)
        idw_factor      = self.parameterAsInt(parameters, self.IDW_RESOLUTION_FACTOR, context)
        idw_neighbours  = self.parameterAsInt(parameters, self.IDW_NEIGHBOURS, context)
        idw_power       = self.parameterAsDouble(parameters, self.IDW_POWER, context)
        clip_percentile = self.parameterAsDouble(parameters, self.CLIP_PERCENTILE, context)

        if dem_layer.crs().isGeographic():
            fb.reportError(
                "WARNING: DEM appears to use a geographic CRS (degrees). "
                "A projected CRS (metres) is strongly recommended.", False
            )

        # Unique temp file names for this run
        run_id            = str(uuid.uuid4())[:8]
        tmp_points        = os.path.join(tempfile.gettempdir(), f"rem_samples_{run_id}.shp")
        tmp_river_surface = os.path.join(tempfile.gettempdir(), f"rem_surface_{run_id}.tif")
        tmp_rem           = os.path.join(tempfile.gettempdir(), f"rem_raw_{run_id}.tif")

        results = {}

        # ---- Step 1: Points along geometry -----------------------------
        fb.setCurrentStep(1)
        fb.pushInfo("Step 1/7 – Creating points along river …")

        points_along = processing.run(
            "native:pointsalonglines",
            {
                "INPUT":        river_layer,
                "DISTANCE":     sample_interval,
                "START_OFFSET": 0,
                "END_OFFSET":   0,
                "OUTPUT":       "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 2: Sample raster values ------------------------------
        fb.setCurrentStep(2)
        fb.pushInfo("Step 2/7 – Sampling DEM elevation at points …")

        sampled = processing.run(
            "native:rastersampling",
            {
                "INPUT":         points_along["OUTPUT"],
                "RASTERCOPY":    dem_layer,
                "COLUMN_PREFIX": "elev_",
                "OUTPUT":        "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        # Save to .shp — avoids GPKG FID uniqueness crash with many features
        processing.run(
            "native:savefeatures",
            {"INPUT": sampled["OUTPUT"], "OUTPUT": tmp_points},
            context=context, feedback=fb,
        )

        sample_layer = QgsVectorLayer(tmp_points, "river_samples", "ogr")
        n_pts = sample_layer.featureCount()
        fb.pushInfo(f"  → {n_pts} sample points created.")

        if n_pts < 4:
            raise QgsProcessingException(
                f"Only {n_pts} sample points generated. "
                "Check that the river line overlaps the DEM and the "
                "sample interval is not too large."
            )

        # Find elevation field index
        field_idx  = sample_layer.fields().indexFromName("elev_1")
        elev_field = "elev_1"
        if field_idx < 0:
            for f in sample_layer.fields():
                if f.name().startswith("elev_"):
                    field_idx  = sample_layer.fields().indexFromName(f.name())
                    elev_field = f.name()
                    break
        if field_idx < 0:
            raise QgsProcessingException(
                f"Could not find elevation field. "
                f"Available fields: {[f.name() for f in sample_layer.fields()]}"
            )

        fb.pushInfo(f"  Elevation field: '{elev_field}' (index {field_idx})")

        # Save river points QC output if requested
        if parameters.get(self.OUTPUT_RIVER_POINTS):
            processing.run(
                "native:savefeatures",
                {
                    "INPUT":  sample_layer,
                    "OUTPUT": parameters[self.OUTPUT_RIVER_POINTS],
                },
                context=context, feedback=fb,
            )
            results[self.OUTPUT_RIVER_POINTS] = parameters[self.OUTPUT_RIVER_POINTS]

        if fb.isCanceled():
            return {}

        # ---- Step 3: IDW interpolation ---------------------------------
        fb.setCurrentStep(3)
        ext     = dem_layer.extent()
        x_res   = dem_layer.rasterUnitsPerPixelX()
        idw_res = x_res * idw_factor

        fb.pushInfo(
            f"Step 3/7 – IDW interpolation at {idw_res:.1f} m "
            f"({idw_factor}× DEM pixel size) …"
        )

        INTERPOLATION_DATA = f"{sample_layer.source()}::~::0::~::{field_idx}::~::0"

        idw_result = processing.run(
            "qgis:idwinterpolation",
            {
                "INTERPOLATION_DATA":        INTERPOLATION_DATA,
                "DISTANCE_COEFFICIENT":      idw_power,
                "NUMBER_OF_NEIGHBOR_POINTS": idw_neighbours,
                "EXTENT":                    ext,
                "PIXEL_SIZE":                idw_res,
                "OUTPUT":                    "TEMPORARY_OUTPUT",
            },
            context=context, feedback=fb,
        )

        idw_lyr = QgsRasterLayer(idw_result["OUTPUT"], "idw_check")
        fb.pushInfo(
            f"  IDW output: {idw_lyr.width()}×{idw_lyr.height()} px, "
            f"extent {idw_lyr.extent()}"
        )

        if fb.isCanceled():
            return {}

        # ---- Step 4: Warp to DEM resolution/extent ---------------------
        fb.setCurrentStep(4)
        fb.pushInfo(
            f"Step 4/7 – Resampling river surface to {x_res:.1f} m …"
        )

        processing.run(
            "gdal:warpreproject",
            {
                "INPUT":             idw_result["OUTPUT"],
                "SOURCE_CRS":        dem_layer.crs(),
                "TARGET_CRS":        dem_layer.crs(),
                "RESAMPLING":        2,      # bilinear
                "TARGET_RESOLUTION": x_res,
                "TARGET_EXTENT":     ext,
                "TARGET_EXTENT_CRS": dem_layer.crs(),
                "NODATA":            None,
                "OPTIONS":           "COMPRESS=LZW",
                "OUTPUT":            tmp_river_surface,
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 5: REM = DEM − river surface -------------------------
        fb.setCurrentStep(5)
        fb.pushInfo("Step 5/7 – Computing REM (DEM − river surface) …")

        processing.run(
            "gdal:rastercalculator",
            {
                "INPUT_A": dem_layer.source(),
                "BAND_A":  1,
                "INPUT_B": tmp_river_surface,
                "BAND_B":  1,
                "FORMULA": "(A - B) * (A >= B)",
                "NO_DATA": -9999,
                "RTYPE":   5,       # Float32
                "OUTPUT":  tmp_rem,
            },
            context=context, feedback=fb,
        )

        if fb.isCanceled():
            return {}

        # ---- Step 6: Percentile clip (optional) ------------------------
        fb.setCurrentStep(6)
        output_path = self.parameterAsOutputLayer(
            parameters, self.OUTPUT_REM, context
        )

        if clip_percentile > 0:
            fb.pushInfo(
                f"Step 6/7 – Clipping above {clip_percentile}th percentile …"
            )
            final_path = self._clip_by_percentile(
                tmp_rem, clip_percentile, output_path, context, fb
            )
        else:
            import shutil
            shutil.copy(tmp_rem, output_path)
            final_path = output_path

        fb.setCurrentStep(7)
        fb.pushInfo("✓ REM complete.")

        # Apply the bundled REM colour ramp automatically when the layer
        # is loaded into the project.
        if context.willLoadLayerOnCompletion(final_path):
            context.layerToLoadOnCompletionDetails(
                final_path
            ).setPostProcessor(_REMStylePostProcessor.create())

        results[self.OUTPUT_REM] = final_path
        return results

    # ------------------------------------------------------------------ #
    #  Percentile clip                                                     #
    # ------------------------------------------------------------------ #

    def _clip_by_percentile(
        self, raster_path, percentile, output_path, context, feedback
    ):
        try:
            from osgeo import gdal
            import numpy as np

            ds   = gdal.Open(raster_path)
            band = ds.GetRasterBand(1)
            nd   = band.GetNoDataValue()
            arr  = band.ReadAsArray().astype(float)
            ds   = None

            if nd is not None:
                arr = np.where(arr == nd, np.nan, arr)
            valid = arr[~np.isnan(arr)]
            p_val = float(np.percentile(valid, percentile))
            feedback.pushInfo(f"  {percentile}th percentile = {p_val:.3f} m")

            processing.run(
                "gdal:rastercalculator",
                {
                    "INPUT_A": raster_path,
                    "BAND_A":  1,
                    "FORMULA": f"numpy.where(A > {p_val}, {p_val}, A)",
                    "NO_DATA": -9999,
                    "RTYPE":   5,
                    "OUTPUT":  output_path,
                },
                context=context, feedback=feedback,
            )
            return output_path

        except ImportError:
            feedback.pushInfo(
                "  numpy/GDAL not available – skipping percentile clip."
            )
            import shutil
            shutil.copy(raster_path, output_path)
            return output_path
