scripts/compute_cumulative_catchment.py

"""
Fills catchment_area_km2 for every reach graph node (confluences,
splits, rejoins, virtual infill -- not just the ~8 real gauges Hub'Eau
publishes it for), using BD TOPO's incremental catchment polygons
(bassin_versant_topographique.geojson).

ALGORITHM:
  1. Spatially assign each incremental polygon to the edges it contains
     (an edge's midpoint falling inside a polygon means that edge's
     reach is the one that polygon's incremental drainage feeds).
  2. For any node, cumulative catchment area = sum of the AREAS OF
     DISTINCT POLYGONS among every edge upstream of it (walked via the
     real directed graph, unbounded -- unlike the split/rejoin check's
     15km search radius, catchment area needs the true full upstream
     drainage, however far that extends). "Distinct" matters: many
     edges (and the virtual nodes on them) share the same polygon, since
     BD TOPO's confluence-to-confluence stretches are far coarser than
     our node resolution -- summing per-edge would massively overcount.

CAVEAT: the polygon assignment step needs geopandas, which isn't
available in the environment this was developed in (no network access
to install it). The graph-traversal/area-summation logic (the actual
novel part) is tested with synthetic polygon assignments standing in
for a real spatial join -- but the spatial join itself has not been run
against real data. Please verify the polygon assignment step
specifically (e.g. spot-check a few edges against the map) before
trusting the final areas.

Usage:
    python -m scripts.compute_cumulative_catchment --data-root datasets
"""
import argparse
import sys
from pathlib import Path

import networkx as nx
import pandas as pd


def assign_edges_to_polygons(edges_df: pd.DataFrame, nodes_df: pd.DataFrame, polygons_path: Path) -> pd.Series:
    """
    Spatial join: for each edge, which incremental catchment polygon
    contains its midpoint. Returns a Series indexed like edges_df, with
    the polygon's unique identifier (or NaN if no polygon contains it --
    can happen near basin edges or bbox boundaries).

    NEEDS geopandas -- not verified in the environment this was written
    in, see module docstring.
    """
    import geopandas as gpd
    from shapely.geometry import Point

    polygons = gpd.read_file(polygons_path)
    if "cleabs" in polygons.columns:
        poly_id_col = "cleabs"
    else:
        polygons = polygons.reset_index().rename(columns={"index": "_poly_idx"})
        poly_id_col = "_poly_idx"

    coord = nodes_df.set_index("station_code")[["latitude", "longitude"]]
    midpoints = []
    for _, e in edges_df.iterrows():
        try:
            src, tgt = coord.loc[e["source"]], coord.loc[e["target"]]
        except KeyError:
            midpoints.append(None)
            continue
        midpoints.append(Point((src["longitude"] + tgt["longitude"]) / 2,
                                (src["latitude"] + tgt["latitude"]) / 2))

    edge_points = gpd.GeoDataFrame({"edge_idx": range(len(edges_df))}, geometry=midpoints, crs=polygons.crs)
    joined = gpd.sjoin(edge_points.dropna(subset=["geometry"]), polygons[[poly_id_col, "geometry"]],
                        how="left", predicate="within")
    joined = joined.drop_duplicates("edge_idx")  # a midpoint exactly on a shared boundary could match >1 polygon

    result = pd.Series(index=range(len(edges_df)), dtype=object)
    result.loc[joined["edge_idx"]] = joined[poly_id_col].values
    return result


def compute_cumulative_catchment(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    edge_polygon_id: pd.Series,
    polygon_areas: dict,
) -> pd.Series:
    """
    For every node, sum the areas of distinct polygons among all edges
    upstream of it (inclusive of the edge(s) arriving at the node
    itself). Pure graph-traversal + arithmetic -- no geospatial
    dependency, fully testable without geopandas (see the test in this
    script's accompanying verification, or run this module's functions
    directly against synthetic data).

    Args:
        nodes_df: reach graph nodes [station_code, ...]
        edges_df: reach graph edges [source, target, ...], same row
            order/index as edge_polygon_id.
        edge_polygon_id: from assign_edges_to_polygons (or an equivalent
            manually-built mapping for testing) -- polygon ID per edge,
            aligned to edges_df's index.
        polygon_areas: {polygon_id: area_km2}

    Returns:
        Series indexed like nodes_df, cumulative catchment area in km2
        (NaN for a node with no upstream edges assigned to any polygon
        at all -- a true headwater with no BD TOPO catchment coverage,
        distinct from zero).
    """
    G = nx.DiGraph()
    G.add_nodes_from(nodes_df["station_code"])
    edge_to_poly = {}
    for i, e in edges_df.iterrows():
        G.add_edge(e["source"], e["target"])
        poly = edge_polygon_id.get(i)
        if pd.notna(poly):
            edge_to_poly[(e["source"], e["target"])] = poly

    results = {}
    for node in nodes_df["station_code"]:
        ancestors = nx.ancestors(G, node) | {node}
        # every edge whose TARGET is in this node's ancestor-or-self set
        # is an edge that feeds into this node's upstream drainage
        polys_here = set()
        for (u, v), poly in edge_to_poly.items():
            if v in ancestors:
                polys_here.add(poly)
        if not polys_here:
            results[node] = float("nan")
        else:
            results[node] = sum(polygon_areas.get(p, 0.0) for p in polys_here)

    return pd.Series(results).reindex(nodes_df["station_code"])


def run_for_basin(data_root: Path, basin_id: int, file_key: str) -> None:
    graph_dir = data_root / "reach_graph"
    enriched_path = graph_dir / f"{file_key}_nodes_enriched.csv"
    edges_path = graph_dir / f"{file_key}_edges.csv"
    polygons_path = data_root / "bdtopo_hydro" / "bassin_versant_topographique.geojson"

    if not enriched_path.exists() or not edges_path.exists():
        print(f"{file_key}: missing enriched nodes/edges, run build_reach_graphs.py + "
              f"enrich_reach_graph.py first. Skipping.")
        return
    if not polygons_path.exists():
        print(f"{file_key}: no {polygons_path} found -- run download_bdtopo_hydro.py first. Skipping.")
        return

    nodes_df = pd.read_csv(enriched_path)
    edges_df = pd.read_csv(edges_path)
    print(f"{file_key}: {len(nodes_df)} nodes, {len(edges_df)} edges")

    try:
        import geopandas as gpd
    except ImportError:
        print(f"{file_key}: geopandas not installed -- cannot do the spatial join. Skipping.")
        return

    polygons = gpd.read_file(polygons_path)
    poly_id_col = "cleabs" if "cleabs" in polygons.columns else "_poly_idx"
    if poly_id_col == "_poly_idx":
        polygons = polygons.reset_index().rename(columns={"index": poly_id_col})
    polygons_metric = polygons.to_crs(epsg=2154)
    polygon_areas = dict(zip(polygons[poly_id_col], polygons_metric.geometry.area / 1e6))

    edge_polygon_id = assign_edges_to_polygons(edges_df, nodes_df, polygons_path)
    n_assigned = edge_polygon_id.notna().sum()
    print(f"  {n_assigned}/{len(edges_df)} edges assigned to a polygon "
          f"({len(edges_df) - n_assigned} unassigned -- near bbox/basin edges, expected)")

    cumulative = compute_cumulative_catchment(nodes_df, edges_df, edge_polygon_id, polygon_areas)
    nodes_df["cumulative_catchment_area_km2"] = cumulative.values

    n_covered = nodes_df["cumulative_catchment_area_km2"].notna().sum()
    print(f"  cumulative_catchment_area_km2 coverage: {n_covered}/{len(nodes_df)} nodes "
          f"(vs. {nodes_df['catchment_area_km2'].notna().sum()} from Hub'Eau's own reported values)")

    # Cross-check against Hub'Eau's real values where both exist, as a
    # sanity check on the whole pipeline -- they measure different things
    # (Hub'Eau: official gauge catchment; this: BD TOPO polygon sum) so
    # exact agreement isn't expected, but they should be in the same
    # ballpark for a real gauge, not off by an order of magnitude.
    both = nodes_df.dropna(subset=["catchment_area_km2", "cumulative_catchment_area_km2"])
    if not both.empty:
        both = both.copy()
        both["ratio"] = both["cumulative_catchment_area_km2"] / both["catchment_area_km2"]
        print(f"  Cross-check vs Hub'Eau (real gauges only): ratio mean={both['ratio'].mean():.2f}, "
              f"range=[{both['ratio'].min():.2f}, {both['ratio'].max():.2f}]")
        print(both[["station_code", "catchment_area_km2", "cumulative_catchment_area_km2", "ratio"]]
              .to_string(index=False))

    nodes_df.to_csv(enriched_path, index=False)
    print(f"  updated {enriched_path}")
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description="Compute cumulative catchment area from BD TOPO polygons")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"))
    args = parser.parse_args()

    for basin_id, file_key in [(0, "eure"), (1, "risle")]:
        run_for_basin(args.data_root, basin_id, file_key)


if __name__ == "__main__":
    main()