# -*- coding: utf-8 -*-
"""📦 Wee By Population Score Processor module.

This module contains functionality for geoe3 by population score processor.
"""

import os
import shutil
import traceback
from typing import List, Optional

from qgis import processing
from qgis.core import (
    QgsCoordinateReferenceSystem,
    QgsRasterLayer,
    QgsTask,
    QgsVectorLayer,
)

from geest.core.algorithms import AreaIterator
from geest.utilities import log_message, resources_path


class WEEByPopulationScoreProcessingTask(QgsTask):
    """
    A QgsTask subclass for calculating GeoE3 x Population SCORE using raster algebra.

    It iterates over study areas, calculates the GeoE3 Score using aligned input rasters
    (GeoE3 and POP), combines the resulting rasters into a VRT, and applies a QML style.

    It takes as input a GeoE3 Score layer (output as the result of Analysis level aggregation
    of GEOE3 workflow). This GeoE3 Score layer has the following classes:

    | Range  | Description               | Color      |
    |--------|---------------------------|------------|
    | 0 - 1  | Very Low Enablement       | ![#FF0000](#) `#FF0000` |
    | 1 - 2  | Low Enablement            | ![#FFA500](#) `#FFA500` |
    | 2 - 3  | Moderately Enabling       | ![#FFFF00](#) `#FFFF00` |
    | 3 - 4  | Enabling                  | ![#90EE90](#) `#90EE90` |
    | 4 - 5  | Highly Enabling           | ![#0000FF](#) `#0000FF` |

    Additionally a population layer containing counts per cell is required. The population
    data is masked to only include pixels contained within study areas and then reclassified into
    three classes:

    | Color      | Description         |
    |------------|---------------------|
    | ![#FFFF00](#) `#FFFF00` | Low Population      |
    | ![#FFA500](#) `#FFA500` | Medium Population   |
    | ![#800000](#) `#800000` | High Population     |


    The output GeoE3 x Population Score can be one of 15 classes calculated as

    ((A - 1) * 3) + B

    Where A is GeoE3 Score and B is Population class. The output classes are as follows:

    | Color      | Description                                 |
    |------------|---------------------------------------------|
    | ![#FF0000](#) `#FF0000` | Very low enablement, low population       |
    | ![#FF0000](#) `#FF0000` | Very low enablement, medium population    |
    | ![#FF0000](#) `#FF0000` | Very low enablement, high population      |
    | ![#FFA500](#) `#FFA500` | Low enablement, low population            |
    | ![#FFA500](#) `#FFA500` | Low enablement, medium population         |
    | ![#FFA500](#) `#FFA500` | Low enablement, high population           |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, low population       |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, medium population    |
    | ![#FFFF00](#) `#FFFF00` | Moderately enabling, high population      |
    | ![#90EE90](#) `#90EE90` | Enabling, low population                  |
    | ![#90EE90](#) `#90EE90` | Enabling, medium population               |
    | ![#90EE90](#) `#90EE90` | Enabling, high population                 |
    | ![#0000FF](#) `#0000FF` | Highly enabling, low population           |
    | ![#0000FF](#) `#0000FF` | Highly enabling, medium population        |
    | ![#0000FF](#) `#0000FF` | Highly enabling, high population          |

    One sub product is made per study area and then all of the study area outputs are
    combined as a VRT and assigned a QML with the correct legend colours.

    📒 The population processing phases resamples the population into a grid aligned
        raster with the same output pixel dimensions as the study area. The resampling
        process uses the gdalwarp -r sum flag which is not currently available in QGIS
        processing framework, so we manually call gdalwarp to perform this operation.

        I tried to implement a generic way to discover the location of gdalwarp but this
        may not work well on all systems.

    Args:
        study_area_gpkg_path (str): Path to the GeoPackage containing study area masks.
        working_directory (str): Directory to save the output rasters.
        target_crs (Optional[QgsCoordinateReferenceSystem]): CRS for the output rasters.
        force_clear (bool): Flag to force clearing of all outputs before processing.
    """

    def __init__(
        self,
        study_area_gpkg_path: str,
        working_directory: str,
        target_crs: Optional[QgsCoordinateReferenceSystem] = None,
        force_clear: bool = False,
    ):
        super().__init__("GeoE3 Score Processor", QgsTask.CanCancel)
        self.study_area_gpkg_path = study_area_gpkg_path

        self.output_dir = os.path.join(working_directory, "geoe3_by_population_score")
        os.makedirs(self.output_dir, exist_ok=True)

        # These folders should already exist from the aggregation analysis and population raster processing
        self.population_folder = os.path.join(working_directory, "population")
        self.geoe3_folder = os.path.join(working_directory, "geoe3_score")

        self.force_clear = force_clear
        if self.force_clear and os.path.exists(self.output_dir):
            for file in os.listdir(self.output_dir):
                os.remove(os.path.join(self.output_dir, file))

        self.target_crs = target_crs
        if not self.target_crs:
            layer: QgsVectorLayer = QgsVectorLayer(
                f"{self.study_area_gpkg_path}|layername=study_area_clip_polygons",
                "study_area_clip_polygons",
                "ogr",
            )
            self.target_crs = layer.crs()
            del layer
        self.output_rasters: List[str] = []

        log_message("Initialized GeoE3 Score Processing Task")

    def run(self) -> bool:
        """
        Executes the GeoE3 Score calculation task.

        Returns:
            bool: True if the task completed successfully, False otherwise.
        """
        try:
            self.calculate_score()
            self.generate_vrt()
            return True
        except Exception as e:
            log_message(f"Task failed: {e}")
            log_message(traceback.format_exc())
            return False

    def validate_rasters(
        self,
        geoe3_raster: QgsRasterLayer,
        pop_raster: QgsRasterLayer,
        dimension_check=False,
    ) -> None:
        """
        Checks if GEOE3 and POP rasters have the same origin, dimensions, and pixel sizes.

        Args:
            geoe3_raster (QgsRasterLayer): The GEOE3 raster layer.
            pop_raster (QgsRasterLayer): The population raster layer.
            dimension_check (bool): Flag to check if the rasters have the same dimensions. Defaults to False.

        Raises:
            ValueError: If one or both input rasters are invalid, or if rasters don't share the same extent or dimensions when dimension_check is True.
        """
        log_message("Validating input rasters")
        log_message(f"GEOE3 Raster: {geoe3_raster.source()}")
        log_message(f"POP Raster  : {pop_raster.source()}")  # noqa E203

        if not geoe3_raster.isValid() or not pop_raster.isValid():
            raise ValueError("One or both input rasters are invalid.")

        if not dimension_check:
            return

        geoe3_provider = geoe3_raster.dataProvider()
        pop_provider = pop_raster.dataProvider()

        geoe3_extent = geoe3_provider.extent()
        pop_extent = pop_provider.extent()
        if geoe3_extent != pop_extent:
            raise ValueError("Input rasters do not share the same extent.")

        geoe3_size = geoe3_provider.xSize(), geoe3_provider.ySize()
        pop_size = pop_provider.xSize(), pop_provider.ySize()
        if geoe3_size != pop_size:
            raise ValueError("Input rasters do not share the same dimensions.")

        log_message("Validation successful: rasters are aligned.")

    def calculate_score(self) -> None:
        """
        Calculates GeoE3 by POP SCORE using raster algebra and saves the result for each area.
        """
        area_iterator = AreaIterator(self.study_area_gpkg_path)
        for index, (_, _, _, _) in enumerate(area_iterator):
            if self.isCanceled():
                return

            geoe3_path = os.path.join(self.geoe3_folder, f"geoe3_masked_{index}.tif")
            population_path = os.path.join(self.population_folder, f"reclassified_{index}.tif")
            geoe3_layer = QgsRasterLayer(geoe3_path, "GeoE3")
            pop_layer = QgsRasterLayer(population_path, "POP")
            self.validate_rasters(geoe3_layer, pop_layer, dimension_check=False)

            output_path = os.path.join(self.output_dir, f"geoe3_by_population_score_{index}.tif")
            if not self.force_clear and os.path.exists(output_path):
                log_message(f"Reusing existing raster: {output_path}")
                self.output_rasters.append(output_path)
                continue

            log_message(f"Calculating GeoE3 by POP SCORE for area {index}")

            params = {
                "INPUT_A": geoe3_layer,
                "BAND_A": 1,
                "INPUT_B": pop_layer,
                "BAND_B": 1,
                "FORMULA": "((A - 1) * 3) + B",
                "NO_DATA": None,
                "EXTENT_OPT": 3,
                "PROJWIN": None,
                "RTYPE": 0,
                "OPTIONS": "",
                "EXTRA": "",
                "OUTPUT": output_path,
            }

            processing.run("gdal:rastercalculator", params)
            self.output_rasters.append(output_path)

            log_message(f"GeoE3 Score raster saved to {output_path}")

    def generate_vrt(self) -> None:
        """
        Combines all GeoE3 Score rasters into a single VRT and applies a QML style.
        """
        vrt_path = os.path.join(self.output_dir, "geoe3_by_population_score.vrt")
        qml_path = os.path.join(self.output_dir, "geoe3_by_population_score.qml")
        source_qml = resources_path("resources", "qml", "geoe3_by_population_score.qml")

        params = {
            "INPUT": self.output_rasters,
            "RESOLUTION": 0,  # Use highest resolution
            "SEPARATE": False,  # Combine into a single band
            "OUTPUT": vrt_path,
        }

        processing.run("gdal:buildvirtualraster", params)
        log_message(f"Generated VRT at {vrt_path}")

        # Apply QML Style
        if os.path.exists(source_qml):
            shutil.copy(source_qml, qml_path)
            log_message(f"Copied QML style from {source_qml} to {qml_path}")
        else:
            log_message("QML style file not found. Skipping QML copy.")

    def finished(self, result: bool) -> None:
        """
        Called when the task completes.

        Args:
            result (bool): The result of the task execution.
        """
        if result:
            log_message("GeoE3 Score calculation completed successfully.")
        else:
            log_message("GeoE3 Score calculation failed.")
