src/data/loaders/bdtopo_hydro.py

"""
BD TOPO hydrography loader.
Single Responsibility: Load IGN BD TOPO hydrographic data (river
centerline tronçons, hydrographic surfaces, catchment polygons) as
downloaded by scripts/download_bdtopo_hydro.py.

Data source: IGN Geoplateforme WFS (https://data.geopf.fr/wfs), Licence
Ouverte 2.0. See download_bdtopo_hydro.py for the actual download step
and its axis-order/typeName caveats -- this loader assumes that data is
already on disk as GeoJSON.
"""
import json
import math
from pathlib import Path
from typing import Optional, List, Tuple

import pandas as pd
import matplotlib.pyplot as plt
from .base import BaseDataLoader


def _haversine_km(lat1, lon1, lat2, lon2) -> float:
    R = 6371.0
    lat1, lon1, lat2, lon2 = map(math.radians, [lat1, lon1, lat2, lon2])
    dlat, dlon = lat2 - lat1, lon2 - lon1
    a = math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2
    return R * 2 * math.asin(math.sqrt(a))


class BDTopoHydroLoader(BaseDataLoader):
    """
    Loads BD TOPO hydrography for a given area: river centerline
    tronçons (with real elevation at every vertex), hydrographic
    surfaces (includes karst classification), and catchment polygons.

    Unlike the other loaders in this project, BD TOPO data doesn't come
    as flat CSVs -- it's GeoJSON with 3D (lon, lat, altitude) LineString
    geometry. This loader keeps everything as plain dicts/DataFrames
    (no geopandas/shapely dependency) wherever the operation allows it,
    since several of these (elevation profile, karst proximity) only
    need vertex coordinates, not real geometric operations. Catchment
    polygon loading DOES need geopandas (point-in-polygon is a real
    geometric operation), so that one method degrades gracefully with a
    clear message if geopandas isn't installed.
    """

    def __init__(self, data_path: Path):
        """
        Args:
            data_path: directory containing troncon_hydrographique.geojson,
                surface_hydrographique.geojson, and (optionally)
                bassin_versant_topographique.geojson, as produced by
                download_bdtopo_hydro.py.
        """
        super().__init__(data_path)

    def load(self) -> dict:
        """
        Load the raw troncon_hydrographique GeoJSON (the primary
        dataset for this loader -- river centerline tronçons).

        Returns:
            The parsed GeoJSON dict ({"type": "FeatureCollection", "features": [...]})
        """
        path = self.data_path / "troncon_hydrographique.geojson"
        if not path.exists():
            raise FileNotFoundError(f"No troncon_hydrographique.geojson found in {self.data_path}")
        return json.loads(path.read_text())

    def load_surfaces(self) -> dict:
        """Load surface_hydrographique.geojson (includes karst classification)."""
        path = self.data_path / "surface_hydrographique.geojson"
        if not path.exists():
            raise FileNotFoundError(f"No surface_hydrographique.geojson found in {self.data_path}")
        return json.loads(path.read_text())

    def load_catchments(self):
        """
        Load bassin_versant_topographique.geojson as a GeoDataFrame.
        Requires geopandas (real polygon geometry, unlike the other
        methods here). Raises a clear ImportError if unavailable rather
        than a confusing geopandas stack trace.
        """
        try:
            import geopandas as gpd
        except ImportError as e:
            raise ImportError(
                "load_catchments needs geopandas: pip install geopandas"
            ) from e
        path = self.data_path / "bassin_versant_topographique.geojson"
        if not path.exists():
            raise FileNotFoundError(f"No bassin_versant_topographique.geojson found in {self.data_path}")
        return gpd.read_file(path)

    def get_metadata(self) -> dict:
        """Get BD TOPO hydrography metadata."""
        meta = super().get_metadata()
        meta.update({
            "data_type": "bdtopo_hydrography",
            "source_organization": "IGN (Geoplateforme WFS, data.geopf.fr)",
            "license": "Licence Ouverte 2.0",
        })
        return meta

    # ------------------------------------------------------------------
    # River name filtering + topology (shared by several methods below)
    # ------------------------------------------------------------------

    def _tronçons_for_river(
        self, geojson: dict, river_name: str,
        name_fields=("cpx_toponyme_de_cours_d_eau", "toponyme", "nom_cours_d_eau"),
    ) -> List[dict]:
        """Filter troncon_hydrographique features to those matching a river name."""
        features = geojson.get("features", [])
        if not features:
            return []
        name_field = next((f for f in name_fields if f in features[0].get("properties", {})), None)
        if name_field is None:
            raise ValueError(f"None of {name_fields} found in feature properties.")
        river_name_lower = river_name.lower()
        return [
            f for f in features
            if river_name_lower in str(f.get("properties", {}).get(name_field, "")).lower()
        ]

    # ------------------------------------------------------------------
    # Elevation profile from embedded Z-coordinates
    # ------------------------------------------------------------------

    def get_elevation_profile(self, river_name: str, geojson: Optional[dict] = None) -> pd.DataFrame:
        """
        Real elevation profile along a named river, using the Z
        (altitude) coordinate already embedded in every tronçon vertex
        -- far finer resolution than interpolating between the 15-27
        gauge stations we otherwise have elevation for.

        Does NOT solve the "132 disconnected components" ordering
        problem (see analyze_bdtopo_hydro.py's export_real_centerline
        for the graph-based ordering approach) -- this returns raw
        (lon, lat, elevation_m) points, unordered across tronçon
        boundaries, sorted only by elevation descending. Fine for
        elevation range / distribution questions; use
        export_real_centerline's ordered output if you need a proper
        upstream-to-downstream sequence.

        Args:
            river_name: e.g. "Risle" or "Eure".
            geojson: optional pre-loaded troncon_hydrographique dict.

        Returns:
            DataFrame [longitude, latitude, elevation_m], sorted by
            elevation_m descending (upstream-ish to downstream-ish).
        """
        geojson = geojson or self.load()
        matched = self._tronçons_for_river(geojson, river_name)
        if not matched:
            raise ValueError(f"No tronçons matched river name '{river_name}'")

        rows = []
        for f in matched:
            coords = f.get("geometry", {}).get("coordinates", [])
            for c in coords:
                if len(c) >= 3:
                    rows.append({"longitude": c[0], "latitude": c[1], "elevation_m": c[2]})

        if not rows:
            raise ValueError(
                f"Matched {len(matched)} tronçon(s) for '{river_name}' but none had a Z "
                f"(elevation) coordinate -- this BD TOPO export may be 2D-only."
            )

        df = pd.DataFrame(rows).drop_duplicates()
        return df.sort_values("elevation_m", ascending=False).reset_index(drop=True)

    # ------------------------------------------------------------------
    # Karst proximity check
    # ------------------------------------------------------------------

    def check_karst_near_point(
        self,
        latitude: float,
        longitude: float,
        search_radius_km: float = 3.0,
        surfaces_geojson: Optional[dict] = None,
    ) -> List[dict]:
        """
        Find karst-classified surface_hydrographique features within
        `search_radius_km` of a point. General-purpose version of the
        amont/aval-bétoire check built for this project's Risle
        investigation -- works for any point.

        Args:
            latitude, longitude: the point to search around.
            search_radius_km: only features with a vertex within this
                distance are returned.
            surfaces_geojson: optional pre-loaded surface_hydrographique dict.

        Returns:
            List of {"nature": str, "distance_km": float} for every
            matching karst feature, nearest first. Empty list if none found.
        """
        geojson = surfaces_geojson or self.load_surfaces()
        features = geojson.get("features", [])

        results = []
        for f in features:
            props = f.get("properties", {})
            nature = props.get("nature") or props.get("Nature") or ""
            if "karst" not in str(nature).lower():
                continue
            dist = self._feature_min_distance_km(f, latitude, longitude)
            if dist <= search_radius_km:
                results.append({"nature": nature, "distance_km": dist})

        return sorted(results, key=lambda r: r["distance_km"])

    @staticmethod
    def _feature_min_distance_km(feature: dict, lat: float, lon: float) -> float:
        """Minimum distance from (lat, lon) to any vertex of a GeoJSON feature."""
        geom = feature.get("geometry") or {}
        gtype = geom.get("type")
        coords = geom.get("coordinates")
        if coords is None:
            return float("inf")

        def flatten(c, depth):
            if depth == 0:
                yield c
            else:
                for sub in c:
                    yield from flatten(sub, depth - 1)

        depth = {"Point": 0, "LineString": 1, "MultiLineString": 2,
                  "Polygon": 2, "MultiPolygon": 3}.get(gtype)
        if depth is None:
            return float("inf")

        best = float("inf")
        for pt in flatten(coords, depth):
            d = _haversine_km(lat, lon, pt[1], pt[0])
            if d < best:
                best = d
        return best

    # ------------------------------------------------------------------
    # Plots
    # ------------------------------------------------------------------

    def plot_elevation_profile(
        self,
        river_name: str,
        df: Optional[pd.DataFrame] = None,
        figsize: tuple = (10, 5),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Scatter of real BD TOPO vertex elevations along a named river,
        against distance-from-mouth approximated by cumulative point
        index after sorting -- NOT a true ordered profile (see
        get_elevation_profile's docstring), but useful for seeing the
        actual elevation range/density from real data at a glance.

        Args:
            river_name: e.g. "Risle" or "Eure".
            df: optional pre-computed get_elevation_profile() output.
            figsize: figure size in inches.
            save_path: if provided, saves the figure to this path.

        Returns:
            The matplotlib Axes object.
        """
        if df is None:
            df = self.get_elevation_profile(river_name)

        fig, ax = plt.subplots(figsize=figsize)
        ax.scatter(range(len(df)), df["elevation_m"], s=3, alpha=0.4, color="#2E6F95")
        ax.set_title(f"{river_name} — real BD TOPO vertex elevations ({len(df)} points)",
                      fontsize=12, fontweight="bold")
        ax.set_xlabel("Point index (sorted by elevation, not true along-river order)")
        ax.set_ylabel("Elevation (m)")
        ax.grid(True, alpha=0.3)
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches="tight")
            print(f"Plot saved to {save_path}")

        return ax

    def plot_catchments(
        self,
        gdf=None,
        figsize: tuple = (9, 8),
        save_path: Optional[Path] = None,
    ) -> plt.Axes:
        """
        Plot catchment (bassin versant topographique) polygons. Requires
        geopandas.

        Args:
            gdf: optional pre-loaded load_catchments() output.
            figsize: figure size in inches.
            save_path: if provided, saves the figure to this path.

        Returns:
            The matplotlib Axes object.
        """
        if gdf is None:
            gdf = self.load_catchments()

        fig, ax = plt.subplots(figsize=figsize)
        gdf.plot(ax=ax, color="#A9C6D8", edgecolor="#2E6F95", alpha=0.6, linewidth=0.8)
        ax.set_title(f"Catchment Polygons (n={len(gdf)})", fontsize=12, fontweight="bold")
        ax.set_xlabel("Longitude")
        ax.set_ylabel("Latitude")
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches="tight")
            print(f"Plot saved to {save_path}")

        return ax