scripts/analysis/check_dry_valleys.py

"""
Checks whether any real gauged station in this project's actual data
ever shows true dry-valley (zero or near-zero) discharge -- a direct,
factual data check, built to resolve whether zero-inflation-specific
loss functions (Tweedie, dual-head zero-inflated models, etc.) are
actually relevant here, rather than assuming either way from generic
hydrology advice that wasn't written for this specific river system.

Reuses build_discharge_timeseries directly (the same real, established
loading path train_spatiotemporal_gnn.py itself uses), not a
re-derived reading of the raw hydrometric files -- avoids the exact
mistake the first version of check_routing_lag_distribution.py made by
not reusing the real, existing function.

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

import numpy as np
import pandas as pd


def main() -> None:
    parser = argparse.ArgumentParser(description="Check for real dry-valley (zero/near-zero discharge) conditions")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"))
    parser.add_argument("--start-date", type=str, default="2000-01-01",
                         help="As wide a real historical window as available -- a dry-valley check "
                              "benefits from more real history, not just the training window.")
    parser.add_argument("--end-date", type=str, default="2026-12-31")
    parser.add_argument("--near-zero-threshold-m3s", type=float, default=0.01,
                         help="A real, small discharge (L/s equivalent: 10 L/s) below which flow is "
                              "treated as 'near-zero' for reporting purposes -- exact zero is checked "
                              "separately and explicitly, this threshold is only for the 'near' case.")
    args = parser.parse_args()

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

    date_range = (args.start_date, args.end_date)
    rows = []

    for basin in ["eure", "risle"]:
        nodes_path = args.data_root / "reach_graph" / f"{basin}_nodes_enriched.csv"
        if not nodes_path.exists():
            print(f"{basin}: no real nodes file at {nodes_path}, skipping.")
            continue
        nodes_df = pd.read_csv(nodes_path)

        wide = build_discharge_timeseries(nodes_df, args.data_root / "hydrometric", date_range)
        # build_discharge_timeseries returns EVERY real gauged station
        # across BOTH basins (confirmed directly -- it never filters by
        # the nodes_df passed in; that filtering happens later, inside
        # assemble_dynamic_tensor, which this script doesn't call) --
        # restricting to this basin's own real station_codes here is
        # required, or every station gets reported twice, once per
        # basin, with identical values. A real, caught mistake in an
        # earlier version of this exact script.
        this_basin_codes = set(nodes_df["station_code"]) & set(wide.columns)
        wide = wide[sorted(this_basin_codes)]
        print(f"{basin}: real discharge data for {wide.shape[1]} real gauged station(s), "
              f"{wide.shape[0]} real date(s) in range")

        for station_code in wide.columns:
            series = wide[station_code].dropna()
            n_real = len(series)
            if n_real == 0:
                rows.append({"basin": basin, "station_code": station_code, "n_real": 0,
                             "min_m3s": np.nan, "pct_exact_zero": np.nan,
                             "pct_near_zero": np.nan, "p01_m3s": np.nan, "p05_m3s": np.nan})
                continue
            n_zero = int((series == 0).sum())
            n_near_zero = int((series <= args.near_zero_threshold_m3s).sum())
            rows.append({
                "basin": basin, "station_code": station_code, "n_real": n_real,
                "min_m3s": float(series.min()),
                "pct_exact_zero": 100.0 * n_zero / n_real,
                "pct_near_zero": 100.0 * n_near_zero / n_real,
                "p01_m3s": float(series.quantile(0.01)),
                "p05_m3s": float(series.quantile(0.05)),
            })

    if not rows:
        print("No real stations found in either basin -- nothing to report.")
        return

    result_df = pd.DataFrame(rows).sort_values("min_m3s")
    print()
    print("=" * 70)
    print(f"Real per-station discharge floor (near-zero threshold: {args.near_zero_threshold_m3s} m3/s)")
    print("=" * 70)
    print(result_df.to_string(index=False))

    n_ever_exact_zero = int((result_df["pct_exact_zero"] > 0).sum())
    n_ever_near_zero = int((result_df["pct_near_zero"] > 0).sum())
    print()
    if n_ever_exact_zero == 0:
        print("No real station in this dataset ever recorded an exact-zero discharge day -- "
              "zero-inflation-specific losses (Tweedie, dual-head zero-inflated models) do not "
              "appear to match this project's actual real data, based on real evidence.")
    else:
        print(f"{n_ever_exact_zero} real station(s) recorded at least one exact-zero discharge day -- "
              f"real, direct evidence a zero-inflation-aware approach could be worth investigating "
              f"further for those specific stations.")
    print(f"{n_ever_near_zero} real station(s) recorded at least one real day at or below "
          f"{args.near_zero_threshold_m3s} m3/s.")


if __name__ == "__main__":
    main()