src/graph/build_graph.py

"""
PyTorch Geometric graph construction for the hydrometric station network.

SCOPE — surface connectivity only. Groundwater (ADES well) data is
deliberately NOT used here. We don't yet have strong enough evidence
(correlated well hydrographs, shared BDLISA aquifer units) to justify
drawing subsurface/karst edges between stations or basins, and a wrong
edge would quietly corrupt a physics-informed model rather than help it.
Revisit if that evidence gets built later (see the wells discussion
around river_line.py / river_centerline.py).

KARST CAVEAT — disappearing rivers ("pertes"): La Eure and La Risle run
over Normandy chalk plateau terrain, a geological setting where a river
can lose flow underground for a stretch and resurface downstream. An
edge in this graph means "these two gauges are sequential along the
river's real course" — the course itself (from the digitized centerline)
is reliable, but that does NOT by itself verify surface flow is
continuous end to end: a correctly-mapped river can still have a losing
reach along it. If/when specific losing reaches are known, pass them via
`known_losing_reaches` to flag (default) or exclude the affected edges,
rather than silently assuming continuity everywhere.

USAGE — combining with node_features.py's enriched table: `build_surface_edges`
rebuilds its own minimal nodes_df (station_code, basin_id, latitude,
longitude, elevation_m) from scratch, since edge construction only needs
those columns. To get the enriched feature set (IDPR, groundwater,
climate, ...) into the graph, merge it back on top, keyed by
station_code, so node ordering stays consistent with edge_index. Note
La Eure and La Risle are built as TWO SEPARATE Data objects, not one
combined graph — see `build_pyg_graphs_per_basin`:

    from src.graph.node_features import build_node_features
    from src.graph.build_graph import build_surface_edges, build_pyg_graphs_per_basin

    enriched, feat_report = build_node_features(
        station_elevations_path=..., idpr_path=..., ades_path=..., hydrometric_path=...,
    )
    base_nodes, edges_df, graph_report = build_surface_edges(
        enriched, centerline_dir=..., basin_file_names={0: "eure", 1: "risle"},
    )
    full_nodes = base_nodes.merge(
        enriched.drop(columns=["basin_id", "latitude", "longitude", "elevation_m"]),
        on="station_code", how="left",
    )
    graphs = build_pyg_graphs_per_basin(full_nodes, edges_df)
    eure_graph, risle_graph = graphs[0], graphs[1]
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

try:
    import torch
    from torch_geometric.data import Data
except ImportError as e:
    raise ImportError(
        "This module requires torch and torch_geometric: "
        "pip install torch torch_geometric"
    ) from e

from ..data.river_graph import (
    NodeBuilder, StationNode, order_stations_by_elevation, EdgeBuilder, assign_basin_id,
)
from ..data.river_centerline import (
    load_centerline, snap_gauges_to_centerline, cumulative_distance_km,
)


@dataclass
class GraphBuildReport:
    """Summary of what happened during construction, for a sanity check
    before trusting the graph — print/log this, don't just discard it."""
    basin_ids: List[int]
    n_nodes: int
    n_edges: int
    edges_from_centerline: int   # built from real digitized geometry
    edges_from_fallback: int     # built from elevation/latitude ordering only (weaker evidence)
    flagged_losing_reaches: int  # edges marked verified_continuous=False
    excluded_losing_reaches: int # edges dropped entirely due to known_losing_reaches

    def __str__(self) -> str:
        return (
            f"Graph: {self.n_nodes} nodes, {self.n_edges} edges across basins {self.basin_ids}\n"
            f"  from real centerline geometry: {self.edges_from_centerline}\n"
            f"  from elevation/latitude fallback (no centerline available): {self.edges_from_fallback}\n"
            f"  flagged as possibly-discontinuous (known losing reach): {self.flagged_losing_reaches}\n"
            f"  excluded entirely (known losing reach): {self.excluded_losing_reaches}"
        )


def _edges_from_real_centerline(
    basin_id: int,
    b_nodes_df: pd.DataFrame,
    centerline: pd.DataFrame,
) -> pd.DataFrame:
    """
    Build sequential surface edges using real digitized centerline order
    (mouth -> source measured as centerline_km, so upstream = larger km).
    This is preferred over elevation/latitude ordering when available,
    since it follows the river's actual traced course rather than
    assuming station elevation alone determines sequence.
    """
    gauges = snap_gauges_to_centerline(centerline, b_nodes_df)
    gauges = gauges.sort_values("centerline_km", ascending=False).reset_index(drop=True)

    edges = []
    for i in range(len(gauges) - 1):
        a, b = gauges.iloc[i], gauges.iloc[i + 1]
        edges.append({
            "source": a["station_code"], "target": b["station_code"], "basin_id": basin_id,
            "distance_km": round(float(a["centerline_km"] - b["centerline_km"]), 3),
            "elevation_drop_m": (
                round(float(a["elevation_m"] - b["elevation_m"]), 2)
                if pd.notna(a["elevation_m"]) and pd.notna(b["elevation_m"]) else None
            ),
            "source_of_evidence": "real_centerline",
        })
    return pd.DataFrame(edges)


def _edges_from_fallback_ordering(basin_id: int, b_nodes: List[StationNode]) -> pd.DataFrame:
    """Fallback when no digitized centerline exists for this basin: order by
    elevation (or latitude, per-basin, if elevation is incomplete — see
    order_stations_by_elevation), then chain sequentially. Weaker evidence
    than a real traced course."""
    order = order_stations_by_elevation(b_nodes, descending=True)
    edges_df = EdgeBuilder.build_sequential(b_nodes, order)
    if not edges_df.empty:
        edges_df["source_of_evidence"] = "elevation_or_latitude_fallback"
    return edges_df


def build_surface_edges(
    elevations_df: pd.DataFrame,
    centerline_dir: Optional[Path] = None,
    basin_file_names: Optional[Dict[int, str]] = None,
    known_losing_reaches: Optional[List[Tuple[str, str]]] = None,
    exclude_losing_reaches: bool = False,
    prefix_map: Optional[Dict[str, int]] = None,
) -> Tuple[pd.DataFrame, pd.DataFrame, GraphBuildReport]:
    """
    Build station nodes and surface-only edges, preferring real digitized
    centerline order per basin when available, falling back to elevation/
    latitude ordering otherwise.

    Args:
        elevations_df: [station_code, latitude, longitude, elevation_m].
        centerline_dir: directory containing {name}_centerline.csv files
            (see river_centerline.py). If None, every basin uses the
            fallback ordering.
        basin_file_names: {basin_id: file_key} mapping to find each
            basin's centerline CSV, e.g. {0: "eure", 1: "risle"}.
        known_losing_reaches: list of (source_station_code, target_station_code)
            pairs known or suspected to have discontinuous surface flow
            (karst losing reaches). These edges are flagged
            (verified_continuous=False) by default, or dropped entirely
            if exclude_losing_reaches=True.
        exclude_losing_reaches: if True, remove known losing-reach edges
            from the graph instead of just flagging them.
        prefix_map: optional override for basin assignment from station
            code prefix, see river_graph.assign_basin_id.

    Returns:
        (nodes_df, edges_df, report)
    """
    basin_file_names = basin_file_names or {0: "eure", 1: "risle"}
    known_losing_reaches = set(known_losing_reaches or [])

    nodes = NodeBuilder.build_from_dataframe(elevations_df, prefix_map)
    if not nodes:
        raise ValueError("No stations matched a basin prefix — check the data or prefix_map.")

    basins: Dict[int, List[StationNode]] = {}
    for n in nodes:
        basins.setdefault(n.basin_id, []).append(n)

    all_edges = []
    n_from_centerline, n_from_fallback = 0, 0

    for basin_id, basin_nodes in basins.items():
        b_nodes_df = pd.DataFrame([{
            "station_code": n.station_id, "latitude": n.latitude,
            "longitude": n.longitude, "elevation_m": n.elevation,
        } for n in basin_nodes])

        csv_path = (centerline_dir / f"{basin_file_names.get(basin_id, basin_id)}_centerline.csv"
                    if centerline_dir else None)
        if csv_path and csv_path.exists():
            centerline = load_centerline(csv_path)
            edges_df = _edges_from_real_centerline(basin_id, b_nodes_df, centerline)
            n_from_centerline += len(edges_df)
        else:
            edges_df = _edges_from_fallback_ordering(basin_id, basin_nodes)
            n_from_fallback += len(edges_df)

        all_edges.append(edges_df)

    edges_df = pd.concat(all_edges, ignore_index=True) if all_edges else pd.DataFrame()

    n_flagged, n_excluded = 0, 0
    if not edges_df.empty:
        edges_df["verified_continuous"] = ~edges_df.apply(
            lambda r: (r["source"], r["target"]) in known_losing_reaches, axis=1
        )
        n_flagged = int((~edges_df["verified_continuous"]).sum())
        if exclude_losing_reaches and n_flagged:
            n_excluded = n_flagged
            n_flagged = 0
            edges_df = edges_df[edges_df["verified_continuous"]].reset_index(drop=True)

    nodes_df = pd.DataFrame([{
        "station_code": n.station_id, "basin_id": n.basin_id,
        "latitude": n.latitude, "longitude": n.longitude, "elevation_m": n.elevation,
    } for n in nodes])

    report = GraphBuildReport(
        basin_ids=sorted(basins.keys()), n_nodes=len(nodes_df), n_edges=len(edges_df),
        edges_from_centerline=n_from_centerline, edges_from_fallback=n_from_fallback,
        flagged_losing_reaches=n_flagged, excluded_losing_reaches=n_excluded,
    )
    return nodes_df, edges_df, report


def build_pyg_graph(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    feature_columns: Optional[List[str]] = None,
    add_missingness_flags: bool = True,
    bidirectional: bool = False,
    standardize_features: bool = True,
) -> Data:
    """
    Convert the surface-connectivity tables into a torch_geometric Data
    object. Basins with no cross-basin edges naturally form disconnected
    components in the graph (as intended — surface connectivity only).

    Node features (x): auto-detected numeric columns from `nodes_df` if
    `feature_columns` is None — this means it works directly with the
    plain 4-column table from `build_surface_edges` OR the enriched
    table from `node_features.build_node_features` (elevation, IDPR,
    groundwater, climate, ...). Columns named `target_*` are always
    excluded from `x` regardless of `feature_columns`, since those are
    the model's targets, not inputs — see node_features.py's docstring
    on why interpolating/using them as features would be label leakage.
    If present, target columns are attached separately as `data.y`
    (raw, NaN preserved — the training loop should mask NaN targets,
    not have them silently imputed).

    STRUCTURAL_COLUMNS (is_gauged, is_confluence, is_split_point,
    is_rejoin_point, snap_distance_km, braid_id) are ALSO always
    excluded from `x`, even though several are boolean — and pandas
    treats bool as a numeric dtype, so without this explicit exclusion
    they'd silently get z-scored and fed to the model as if they were
    physical covariates like elevation, which they are not (confirmed:
    this actually happened before this exclusion list existed). They're
    still attached to the returned Data object as their own named
    attributes (data.is_gauged, etc.) rather than being fully discarded
    — needed downstream for masking supervised loss to gauged nodes and
    for physics_losses.py's confluence/braid index builders, just not
    as model input.

    Edge features (edge_attr): [distance_km, elevation_drop_m, verified_continuous]

    Args:
        nodes_df, edges_df: from `build_surface_edges` (optionally with
            nodes_df enriched via node_features.build_node_features,
            re-merged onto the same station ordering — see module usage
            example).
        feature_columns: explicit list of columns to use as `x`. If
            None, auto-detects all numeric, non-target, non-structural
            columns.
        add_missingness_flags: if True, adds a `{col}__was_missing`
            binary column for any feature column that had NaNs, before
            mean-filling those NaNs — so the model can distinguish "no
            data nearby" from "value happens to be near the mean".
        bidirectional: if True, add a reverse edge for every upstream->
            downstream edge (common in GNN practice to let information
            flow both ways even though physical flow is directional).
            The reverse edges get elevation_drop_m negated.
        standardize_features: if True, z-score each feature column
            (after mean-filling). Set False to keep raw units.

    Returns:
        A torch_geometric.data.Data with `.station_codes` (index ->
        station_code), `.node_id_map` (station_code -> index),
        `.feature_names` (x column order), `.basin_id`, and — if any
        `target_*` columns were present — `.y` (raw values, NaN
        preserved) and `.target_names`. Any STRUCTURAL_COLUMNS present
        on nodes_df are attached as their own same-named attributes
        (booleans as a bool tensor/array; snap_distance_km as float,
        NaN for non-gauged nodes; braid_id kept as a plain Python list,
        since it's station-code strings or None, not something to
        tensor-ify).
    """
    station_codes = nodes_df["station_code"].tolist()
    node_id_map = {code: i for i, code in enumerate(station_codes)}

    STRUCTURAL_COLUMNS = {
        "is_gauged", "is_confluence", "is_split_point", "is_rejoin_point",
        "snap_distance_km", "braid_id",
    }
    # Feature columns that are genuinely time-varying in reality, even
    # though today's pipeline only ever produces a period-aggregate
    # value for them (see node_features.py's add_safran_features/
    # add_groundwater_features/add_ndvi_features) -- distinct from
    # STATIC_COLUMNS below, which are physically time-invariant (a
    # location's elevation or IDPR index doesn't change on any
    # timescale relevant here). n_nearby_wells is classified static
    # despite living under the groundwater loader: it describes monitor
    # coverage (which wells exist nearby), not a water-table quantity
    # that itself varies day to day.
    DYNAMIC_PREFIXES = ("climate_", "avg_groundwater_level", "avg_groundwater_depth", "ndvi_")

    def _is_dynamic(col: str) -> bool:
        base = col[:-len("__was_missing")] if col.endswith("__was_missing") else col
        return any(base.startswith(p) for p in DYNAMIC_PREFIXES)

    target_cols = [c for c in nodes_df.columns if c.startswith("target_")]
    if feature_columns is None:
        feature_columns = [
            c for c in nodes_df.columns
            if c not in ("station_code",) and c not in target_cols
            and c not in STRUCTURAL_COLUMNS
            and pd.api.types.is_numeric_dtype(nodes_df[c])
        ]

    feat = nodes_df[feature_columns].copy()

    missingness_cols = []
    if add_missingness_flags:
        for col in feature_columns:
            if feat[col].isna().any():
                flag_col = f"{col}__was_missing"
                feat[flag_col] = feat[col].isna().astype(float)
                missingness_cols.append(flag_col)

    for col in feature_columns:
        if feat[col].isna().any():
            fill_value = feat[col].mean()
            feat[col] = feat[col].fillna(0.0 if pd.isna(fill_value) else fill_value)

    if standardize_features:
        for col in feature_columns:
            std = feat[col].std()
            feat[col] = (feat[col] - feat[col].mean()) / std if std and std > 0 else 0.0

    all_feature_cols = feature_columns + missingness_cols
    x = torch.tensor(feat[all_feature_cols].values, dtype=torch.float)

    if edges_df.empty:
        edge_index = torch.zeros((2, 0), dtype=torch.long)
        edge_attr = torch.zeros((0, 3), dtype=torch.float)
    else:
        src_idx = edges_df["source"].map(node_id_map).values
        tgt_idx = edges_df["target"].map(node_id_map).values
        if pd.isna(src_idx).any() or pd.isna(tgt_idx).any():
            raise ValueError("Some edge endpoints are not present in nodes_df — check station codes.")

        elev_drop = edges_df["elevation_drop_m"].fillna(0.0).values
        dist_km = edges_df["distance_km"].values
        verified = edges_df.get("verified_continuous", pd.Series([True] * len(edges_df))).astype(float).values

        if bidirectional:
            index_pairs = np.concatenate([
                np.stack([src_idx, tgt_idx]), np.stack([tgt_idx, src_idx]),
            ], axis=1)
            attr = np.concatenate([
                np.stack([dist_km, elev_drop, verified], axis=1),
                np.stack([dist_km, -elev_drop, verified], axis=1),
            ], axis=0)
        else:
            index_pairs = np.stack([src_idx, tgt_idx])
            attr = np.stack([dist_km, elev_drop, verified], axis=1)

        edge_index = torch.tensor(index_pairs, dtype=torch.long)
        edge_attr = torch.tensor(attr, dtype=torch.float)

    data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
    data.station_codes = station_codes
    data.node_id_map = node_id_map
    data.feature_names = all_feature_cols
    data.basin_id = torch.tensor(nodes_df["basin_id"].values, dtype=torch.long)

    # Static/dynamic split: `x` stays the full combined tensor (nothing
    # existing that reads data.x breaks), but every feature is also
    # tagged and split out separately -- x_static for physically
    # time-invariant quantities (elevation, IDPR, catchment area,
    # landcover), x_dynamic for quantities that vary in reality even
    # though the current pipeline only ever hands them over as a single
    # period-aggregate (climate, groundwater level/depth, NDVI). See
    # DYNAMIC_PREFIXES above for exactly which columns land where, and
    # node_features.py's add_safran_features/add_groundwater_features/
    # add_ndvi_features docstrings for why "dynamic" here still means
    # "one aggregated number," not a real time series yet -- a genuine
    # multi-timestep pipeline is a separate, larger piece of future work
    # this split does not attempt to solve on its own.
    dynamic_mask = [_is_dynamic(c) for c in all_feature_cols]
    static_idx = [i for i, d in enumerate(dynamic_mask) if not d]
    dynamic_idx = [i for i, d in enumerate(dynamic_mask) if d]
    data.static_feature_names = [all_feature_cols[i] for i in static_idx]
    data.dynamic_feature_names = [all_feature_cols[i] for i in dynamic_idx]
    x_np_for_split = feat[all_feature_cols].values
    data.x_static = torch.tensor(x_np_for_split[:, static_idx], dtype=torch.float) if static_idx else None
    data.x_dynamic = torch.tensor(x_np_for_split[:, dynamic_idx], dtype=torch.float) if dynamic_idx else None

    for bool_col in ("is_gauged", "is_confluence", "is_split_point", "is_rejoin_point"):
        if bool_col in nodes_df.columns:
            setattr(data, bool_col, torch.tensor(nodes_df[bool_col].values, dtype=torch.bool))
    if "snap_distance_km" in nodes_df.columns:
        data.snap_distance_km = torch.tensor(
            nodes_df["snap_distance_km"].astype(float).values, dtype=torch.float
        )
    if "braid_id" in nodes_df.columns:
        # station-code strings or None -- not tensor-able, kept as a plain
        # list so physics_losses.py's build_braid_index can still use it
        # (that function already works on nodes_df directly, so this is
        # for convenience when only the Data object is at hand, not a
        # hard requirement).
        data.braid_id = nodes_df["braid_id"].tolist()

    if target_cols:
        data.y = torch.tensor(nodes_df[target_cols].values, dtype=torch.float)
        data.target_names = target_cols

    return data


def build_pyg_graphs_per_basin(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    feature_columns: Optional[List[str]] = None,
    add_missingness_flags: bool = True,
    bidirectional: bool = False,
    standardize_features: bool = True,
) -> Dict[int, Data]:
    """
    Build a SEPARATE torch_geometric Data object per basin, rather than
    one combined graph with disconnected components. La Eure and La
    Risle are distinct hydrographic systems with no surface connection
    between them (see the KARST CAVEAT above and river_graph.py's basin
    assignment) — one merged Data object would blur that distinction,
    and PyG's own batching (`torch_geometric.data.Batch.from_data_list`)
    already expects a list of separate small graphs, not one pre-merged
    blob, so this is also the more natural shape for training.

    `basin_id` is dropped from the per-graph feature set automatically
    (it's constant within a single basin's graph and would carry zero
    information there — standardizing a zero-variance column is
    meaningless). It's still available as `data.basin_id` metadata on
    each graph if you need it for bookkeeping.

    Args:
        nodes_df, edges_df: from `build_surface_edges` (optionally
            enriched, see the module docstring's usage example).
        feature_columns: explicit column list PER BASIN. If None,
            auto-detected the same way as `build_pyg_graph`, minus
            `basin_id`.
        add_missingness_flags, bidirectional, standardize_features:
            passed straight through to `build_pyg_graph` for each basin.

    Returns:
        {basin_id: Data}, each with its own LOCAL node indexing (0..n-1
        within that basin) and its own `.station_codes` / `.node_id_map`
        / `.feature_names`.
    """
    graphs = {}
    for basin_id in sorted(nodes_df["basin_id"].unique()):
        b_nodes = nodes_df[nodes_df["basin_id"] == basin_id].reset_index(drop=True)
        b_edges = edges_df[edges_df["basin_id"] == basin_id].reset_index(drop=True)

        cols = feature_columns
        if cols is None:
            target_cols = [c for c in b_nodes.columns if c.startswith("target_")]
            cols = [
                c for c in b_nodes.columns
                if c not in ("station_code", "basin_id") and c not in target_cols
                and pd.api.types.is_numeric_dtype(b_nodes[c])
            ]

        graphs[basin_id] = build_pyg_graph(
            b_nodes, b_edges, feature_columns=cols,
            add_missingness_flags=add_missingness_flags,
            bidirectional=bidirectional, standardize_features=standardize_features,
        )

    return graphs