"""
                    GNU GENERAL PUBLIC LICENSE
                      Version 3, 29 June 2007

    Risk Assessment and Adaptation for Critical Infrastructure (RA2CE).
    Copyright (C) 2023-2026 Stichting Deltares

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable

import numpy as np
import rasterio
from affine import Affine
from geopandas import GeoDataFrame
from networkx import Graph, set_edge_attributes
from numpy import nan
from rasterio.features import shapes
from rasterio.mask import mask
from rasterstats import zonal_stats
from shapely.geometry import LineString, box, shape
from shapely.ops import unary_union
from tqdm import tqdm

from ra2ce.network.hazard.hazard_common_functions import (
    get_edges_geoms,
    validate_extent_graph,
)
from ra2ce.network.hazard.hazard_intersect.hazard_intersect_builder_base import (
    HazardIntersectBuilderBase,
)
from ra2ce.network.networks_utils import get_graph_edges_extent, get_valid_mean


@dataclass
class HazardIntersectBuilderForTif(HazardIntersectBuilderBase):
    hazard_aggregate_wl: str = ""
    hazard_names: list[str] = field(default_factory=list)
    ra2ce_names: list[str] = field(default_factory=list)
    hazard_tif_files: list[Path] = field(default_factory=list)

    @property
    def _combined_names(self) -> list[tuple[str, str]]:
        """
        Combines into tuples `hazard_names` and `ra2ce_names`.

        Returns:
            list[tuple[str, str]]: List of `str` tuples (hazard_name, ra2ce_name).
        """
        return list(zip(self.hazard_names, self.ra2ce_names))

    def _fraction_flooded_array(
        self, line: LineString, src: rasterio.io.DatasetReader
    ) -> float:
        """
        Calculates the fraction of a linestring that overlaps with a hazard raster (values > 0).

        Args:
            line (LineString): The linestring to check.
            src (rasterio.io.DatasetReader): An open rasterio dataset.

        Returns:
            float: Fraction of the linestring overlapping raster cells with value > 0.
        """
        bbox_line = box(*line.bounds)
        try:
            out_image, out_transform = mask(
                src, [bbox_line], crop=True, all_touched=True
            )
            flooded_cells = unary_union(
                [
                    shape(x[0])
                    for x in shapes(out_image[0], transform=out_transform)
                    if x[-1] > 0
                ]
            )
            flooded_cells = GeoDataFrame({"geometry": [flooded_cells]})
            line_intersect = flooded_cells.intersection(line)
            return line_intersect.length.sum() / line.length
        except Exception as e:
            logging.info(f"_fraction_flooded_array() {e} \n for line {line}")
            return 0

    def _from_networkx(self, hazard_overlay: Graph) -> Graph:
        """Overlays the hazard raster over the road segments graph.

        Args:
            *hf* (list of Pathlib paths) : #not sure if this is needed as argument if we also read if from the config
            *graph* (NetworkX Graph) : NetworkX graph with geometries that will be intersected with the hazard map raster.

        Returns:
            *graph* (NetworkX Graph) : NetworkX graph with hazard values
        """
        # TODO apply multiprocessing?
        # from tqdm import (
        #     tqdm,  # somehow this only works when importing here and not at the top of the file
        # )

        # Verify the graph type (networkx)
        assert isinstance(hazard_overlay, Graph)
        extent_graph = get_graph_edges_extent(hazard_overlay)

        # Get all edge geometries
        edges_geoms = get_edges_geoms(hazard_overlay)

        def overlay_network_x(hazard_tif_file: Path, hazard_name: str, ra2ce_name: str):
            # Check if the hazard and graph extents overlap
            validate_extent_graph(extent_graph, hazard_tif_file)
            # Add a no-data value for the edges that do not have a geometry
            set_edge_attributes(
                hazard_overlay,
                {
                    (u, v, k): {ra2ce_name + "_" + self.hazard_aggregate_wl[:2]: nan}
                    for u, v, k, edata in hazard_overlay.edges.data(keys=True)
                    if "geometry" not in edata
                },
            )

            # Add the hazard values to the edges that do have a geometry
            gdf = GeoDataFrame(
                {"geometry": [edata["geometry"] for u, v, k, edata in edges_geoms]}
            )

            # Open raster once for both operations
            with rasterio.open(hazard_tif_file) as src:
                raster_array = src.read(1)
                raster_transform = src.transform
                nodata_value = src.nodata

                tqdm.pandas(desc="Graph hazard overlay with " + hazard_name)
                if self.hazard_aggregate_wl == "mean":
                    flood_stats = gdf.geometry.progress_apply(
                        lambda x: zonal_stats(
                            x,
                            raster_array,
                            affine=raster_transform,
                            all_touched=True,
                            add_stats={"mean": get_valid_mean},
                            nodata=nodata_value,
                        )
                    )
                else:
                    flood_stats = gdf.geometry.progress_apply(
                        lambda x: zonal_stats(
                            x,
                            raster_array,
                            affine=raster_transform,
                            all_touched=True,
                            stats=f"{self.hazard_aggregate_wl}",
                            nodata=nodata_value,
                        )
                    )

                try:
                    flood_stats = flood_stats.apply(
                        lambda x: (
                            x[0][self.hazard_aggregate_wl]
                            if x[0][self.hazard_aggregate_wl]
                            else 0
                        )
                    )
                    set_edge_attributes(
                        hazard_overlay,
                        {
                            (edges[0], edges[1], edges[2]): {
                                ra2ce_name + "_" + self.hazard_aggregate_wl[:2]: x
                            }
                            for x, edges in zip(flood_stats, edges_geoms)
                        },
                    )
                except KeyError:
                    logging.warning(
                        "No aggregation method ('aggregate_wl') is chosen - choose from 'max', 'min' or 'mean'."
                    )

                # Get the fraction of the road that is intersecting with the hazard
                tqdm.pandas(
                    desc="Graph fraction with hazard overlay with " + hazard_name
                )

                graph_fraction_flooded = gdf.geometry.progress_apply(
                    lambda x: self._fraction_flooded_array(x, src)
                )
                graph_fraction_flooded = graph_fraction_flooded.fillna(0)
                set_edge_attributes(
                    hazard_overlay,
                    {
                        (edges[0], edges[1], edges[2]): {ra2ce_name + "_fr": x}
                        for x, edges in zip(graph_fraction_flooded, edges_geoms)
                    },
                )

        self._overlay_hazard_files(overlay_network_x)
        return hazard_overlay

    def _from_geodataframe(self, hazard_overlay: GeoDataFrame):
        """Overlays the hazard raster over the road segments GeoDataFrame.

        Args:
            *graph* (GeoDataFrame) : GeoDataFrame that will be intersected with the hazard map raster.

        Returns:

        """
        assert isinstance(hazard_overlay, GeoDataFrame), "Network is not a GeoDataFrame"

        # Make sure none of the geometries is a nonetype object (this will raise an error in zonal_stats)
        empty_entries = hazard_overlay.loc[hazard_overlay.geometry.isnull()]
        if any(empty_entries):
            logging.warning(
                (
                    "Some geometries have NoneType objects (no coordinate information), namely: {}.".format(
                        empty_entries
                    )
                    + "This could be due to segmentation, and might cause an exception in hazard overlay"
                )
            )

        def overlay_geodataframe(
            hazard_tif_file: Path, hazard_name: str, ra2ce_name: str
        ):
            # Validate input
            # Check if network and raster overlap
            extent_graph = (
                hazard_overlay.total_bounds[0],
                hazard_overlay.total_bounds[2],
                hazard_overlay.total_bounds[1],
                hazard_overlay.total_bounds[3],
            )
            validate_extent_graph(extent_graph, hazard_tif_file)

            # Open raster once for both operations
            with rasterio.open(hazard_tif_file) as src:
                raster_array = src.read(1)
                raster_transform = src.transform
                nodata_value = src.nodata

                tqdm.pandas(desc="Network hazard overlay with " + hazard_name)
                flood_stats = hazard_overlay.geometry.progress_apply(
                    lambda _geom_vector: zonal_stats(
                        vectors=_geom_vector,
                        raster=raster_array,
                        affine=raster_transform,
                        all_touched=True,
                        stats="min max",
                        add_stats={"mean": get_valid_mean},
                        nodata=nodata_value,
                    )
                )

                def _get_attributes(gen_flood_stat: list[dict]) -> tuple:
                    # Just get the first element of the generator
                    _flood_stat = gen_flood_stat[0]
                    return _flood_stat["min"], _flood_stat["max"], _flood_stat["mean"]

                (
                    hazard_overlay[ra2ce_name + "_mi"],
                    hazard_overlay[ra2ce_name + "_ma"],
                    hazard_overlay[ra2ce_name + "_me"],
                ) = list(zip(*map(_get_attributes, flood_stats)))

                tqdm.pandas(
                    desc="Network fraction with hazard overlay with " + hazard_name
                )
                hazard_overlay[
                    ra2ce_name + "_fr"
                ] = hazard_overlay.geometry.progress_apply(
                    lambda x: self._fraction_flooded_array(x, src)
                )

        self._overlay_hazard_files(overlay_geodataframe)
        return hazard_overlay

    def _overlay_hazard_files(self, overlay_func: Callable[[str, str, str], None]):
        for i, (hn, rn) in enumerate(self._combined_names):
            overlay_func(self.hazard_tif_files[i], hn, rn)
