scripts/enrich_reach_graph.py

"""
Enriches the reach graph's node tables (from scripts/build_reach_graphs.py)
with every available feature loader, via node_features.py's base_nodes_df
path -- closing the gap where the reach graph and the feature pipeline
existed separately but nothing actually connected them.

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

import pandas as pd

try:
    from src.graph.node_features import build_node_features
except ImportError:
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
    from src.graph.node_features import build_node_features

BASIN_FILE_NAMES = {0: "eure", 1: "risle"}


def enrich_basin(
    data_root: Path, basin_id: int, skip_climate: bool = False,
    date_range: "tuple[str, str] | None" = None,
) -> None:
    file_key = BASIN_FILE_NAMES[basin_id]
    graph_dir = data_root / "reach_graph"
    nodes_path = graph_dir / f"{file_key}_nodes.csv"

    if not nodes_path.exists():
        print(f"{file_key}: no {nodes_path} found -- run scripts/build_reach_graphs.py first. Skipping.")
        return

    base_nodes = pd.read_csv(nodes_path)
    print(f"{file_key}: {len(base_nodes)} nodes loaded from {nodes_path}")

    enriched, report = build_node_features(
        base_nodes_df=base_nodes,
        idpr_path=data_root / "idpr.csv",
        ades_path=data_root / "ades",
        safran_path=None if skip_climate else data_root / "safran",
        # --skip-climate exists specifically because add_safran_features has
        # not been verified at reach-graph scale (~2,900 nodes) and is the
        # prime suspect for a real hang there, unlike everything else in
        # this pipeline which has been checked at this scale already (see
        # add_groundwater_features's rewrite, which hit exactly this class
        # of bug once before). Use this flag to get unblocked on everything
        # else while that gets diagnosed with the real source.
        catchment_path=data_root / "catchment_area.csv",
        hydrometric_path=data_root / "hydrometric",
        date_range=date_range,
        # Safe to include here, unlike an earlier version of this script
        # assumed: target_* attaches via a left-merge on station_code,
        # and real gauge codes ("H...") can never collide with BD TOPO
        # hydrographic node IDs ("NOEUDHYD...") or virtual infill node
        # IDs ("VIRTUAL::..."), confirmed against the real station code
        # namespace -- so only real gauge rows ever get a real target
        # value; every confluence/virtual node correctly gets NaN, not
        # a mislabeled one.
    )
    print(report)

    n_climate = enriched["climate_precip_mm"].notna().sum() if "climate_precip_mm" in enriched.columns else 0
    n_catchment = enriched["catchment_area_km2"].notna().sum() if "catchment_area_km2" in enriched.columns else 0
    n_targets = enriched["target_discharge_m3s_mean"].notna().sum() if "target_discharge_m3s_mean" in enriched.columns else 0
    print(f"  climate coverage: {n_climate}/{len(enriched)} nodes")
    print(f"  target (discharge) coverage: {n_targets}/{len(enriched)} nodes "
          f"(expected: only real gauges -- {int(base_nodes['is_gauged'].sum())} in this basin)")
    print(f"  catchment_area coverage: {n_catchment}/{len(enriched)} nodes "
          f"(expected: only real gauges have one -- see node_features.py's "
          f"add_catchment_features docstring; this is what water_balance_loss "
          f"can actually use right now, everything else needs a cumulative "
          f"catchment estimate this script doesn't build)")

    out_path = graph_dir / f"{file_key}_nodes_enriched.csv"
    enriched.to_csv(out_path, index=False)
    print(f"  saved to {out_path}")
    print()


def main() -> None:
    parser = argparse.ArgumentParser(description="Enrich reach graph nodes with feature loaders")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"))
    parser.add_argument("--skip-climate", action="store_true",
                         help="Skip SAFRAN/ERA5 climate features -- use if that step hangs at reach-graph scale")
    parser.add_argument("--start-date", type=str, default="2013-01-01",
                         help="Start of the training period (inclusive). Default 2013-01-01: the window "
                              "maximizing discharge-gauge coverage (6/8 stations, 13,084 observations) -- "
                              "see the date-range analysis this default came from before changing it.")
    parser.add_argument("--end-date", type=str, default="2026-12-31",
                         help="End of the training period (inclusive).")
    parser.add_argument("--no-date-filter", action="store_true",
                         help="Disable date filtering entirely -- each source uses its own full history")
    args = parser.parse_args()

    date_range = None if args.no_date_filter else (args.start_date, args.end_date)
    if date_range:
        print(f"Using date range: {date_range[0]} to {date_range[1]}")
        print()

    for basin_id in BASIN_FILE_NAMES:
        enrich_basin(args.data_root, basin_id, skip_climate=args.skip_climate, date_range=date_range)


if __name__ == "__main__":
    main()