scripts/analysis/check_discharge_quantile_distribution.py

"""
Computes the real empirical quantile distribution of discharge for
every real gauged station -- built to answer a real, specific question
before deciding how to re-weight or restructure the model's predicted
quantiles: where do real near-zero/low-flow, normal-flow, and peak
events actually fall for each real station, rather than assuming the
current quantile set [0.5, 0.9, 0.95, 0.99] is well-matched to the
real data's actual shape.

Reuses build_discharge_timeseries directly (the same real, established
loading path train_spatiotemporal_gnn.py itself uses) and the same
real per-basin filtering fix check_dry_valleys.py needed -- 
build_discharge_timeseries returns every real gauged station across
BOTH basins regardless of which nodes_df is passed in, confirmed
directly in that earlier diagnostic.

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

import numpy as np
import pandas as pd


QUANTILE_LEVELS = [0.01, 0.05, 0.10, 0.25, 0.50, 0.75, 0.90, 0.95, 0.99]


def main() -> None:
    parser = argparse.ArgumentParser(description="Check the real empirical discharge quantile distribution per station")
    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.")
    parser.add_argument("--end-date", type=str, default="2026-12-31")
    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)
        # Real per-basin filtering, same fix check_dry_valleys.py needed:
        # build_discharge_timeseries returns every real gauged station
        # across BOTH basins, not filtered by the nodes_df passed in.
        this_basin_codes = set(nodes_df["station_code"]) & set(wide.columns)
        wide = wide[sorted(this_basin_codes)]

        for station_code in wide.columns:
            series = wide[station_code].dropna()
            n_real = len(series)
            if n_real == 0:
                continue
            real_quantiles = series.quantile(QUANTILE_LEVELS)
            row = {"basin": basin, "station_code": station_code, "n_real": n_real}
            for q_level, q_value in zip(QUANTILE_LEVELS, real_quantiles):
                row[f"Q{int(q_level*100):02d}"] = float(q_value)
            rows.append(row)

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

    result_df = pd.DataFrame(rows)
    print("=" * 100)
    print("Real empirical discharge quantile distribution, per real gauged station (m3/s)")
    print("=" * 100)
    print(result_df.to_string(index=False))

    print()
    print("=" * 100)
    print("Real skew/spread ratios -- how far the tails sit from the real median, per station")
    print("=" * 100)
    for _, row in result_df.iterrows():
        q01, q50, q99 = row["Q01"], row["Q50"], row["Q99"]
        low_ratio = q50 / q01 if q01 > 0 else float("inf")
        high_ratio = q99 / q50 if q50 > 0 else float("inf")
        print(f"{row['station_code']} ({row['basin']}): Q50/Q01={low_ratio:>8.1f}x, "
              f"Q99/Q50={high_ratio:>8.1f}x  (real Q01={q01:.2f}, Q50={q50:.2f}, Q99={q99:.2f})")

    print()
    print("=" * 100)
    print("Where does this project's CURRENT model quantile set [0.5, 0.9, 0.95, 0.99] actually land?")
    print("=" * 100)
    print("No real LOW-tail quantile exists in the current model output at all -- Q50 is the lowest "
          "quantile predicted, meaning there is no dedicated, explicit low-flow signal in the model's "
          "own output structure, only whatever Q50 happens to learn implicitly.")


if __name__ == "__main__":
    main()