"""
Computes the real travel-time lag for every real edge in the ACTUAL
collapsed training subgraph (the ~150-edge, 100-node graph
train_spatiotemporal_gnn.py really trains on), not the full,
fine-grained reach graph.

CORRECTED VERSION: an earlier version of this script read the raw
reach_graph/{basin}_edges.csv files directly -- real, but the WRONG
edges. Those are the full, uncollapsed tronçon-level graph (~3300-5200
edges per basin), where most real distances between adjacent nodes are
tiny (a few hundred meters), giving travel times of minutes to hours.
The graph actually used for training is the COLLAPSED subgraph
(build_collapsed_subgraph), where chains through dropped virtual nodes
are merged into single, real, much longer edges (summed distance_km/
elevation_drop_m) -- confirmed as the real, meaningful distinction: a
first run against the wrong edges showed 100% of edges at zero real
lag days, which would have been a real, actionable, and WRONG
conclusion about this project's actual routing physics.

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

import numpy as np
import pandas as pd


def estimate_travel_time_hours(distance_km, elevation_drop_m, min_velocity_ms=0.1, max_velocity_ms=3.0):
    """Exact real formula from src/graph/physics_losses.py's own
    estimate_travel_time_hours -- kept in sync manually since this is a
    read-only diagnostic script, not meant to import the training
    pipeline's own torch-dependent modules."""
    distance_m = distance_km * 1000.0
    slope = np.clip(elevation_drop_m / np.maximum(distance_m, 1.0), 1e-6, None)
    velocity = np.clip(np.sqrt(slope) * 10.0, min_velocity_ms, max_velocity_ms)
    return distance_m / velocity / 3600.0


def main() -> None:
    parser = argparse.ArgumentParser(description="Check the real routing lag distribution on the actual collapsed training subgraph")
    parser.add_argument("--data-root", type=Path, default=Path("datasets"))
    parser.add_argument("--max-nodes", type=int, default=100,
                         help="Must match train_spatiotemporal_gnn.py's own --max-nodes for this to "
                              "reflect the real, actual training graph.")
    parser.add_argument("--timestep-hours", type=float, default=24.0)
    args = parser.parse_args()

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

    max_nodes_per_basin = args.max_nodes // 2
    all_lags_days = []

    for basin in ["eure", "risle"]:
        nodes_path = args.data_root / "reach_graph" / f"{basin}_nodes_enriched.csv"
        edges_path = args.data_root / "reach_graph" / f"{basin}_edges.csv"
        if not nodes_path.exists() or not edges_path.exists():
            print(f"{basin}: missing {nodes_path} or {edges_path}, skipping.")
            continue

        full_nodes_df = pd.read_csv(nodes_path)
        full_edges_df = pd.read_csv(edges_path)
        print(f"{basin}: {len(full_nodes_df)} real node(s), {len(full_edges_df)} real edge(s) "
              f"(full reach graph)")

        small_nodes_df, small_edges_df = build_collapsed_subgraph(
            full_nodes_df, full_edges_df, max_nodes=max_nodes_per_basin,
        )
        n_edges = len(small_edges_df)
        print(f"{basin}: collapsed to {len(small_nodes_df)} real node(s), {n_edges} real "
              f"collapsed edge(s) -- THIS is the real graph training actually uses")

        if n_edges == 0:
            print(f"{basin}: zero real collapsed edges -- nothing to report.")
            continue

        drop = small_edges_df["elevation_drop_m"].fillna(0.1).clip(lower=0.1)
        hours = estimate_travel_time_hours(small_edges_df["distance_km"].values, drop.values)
        days = hours / 24.0
        lag_steps = np.round(hours / args.timestep_hours).astype(int)

        print(f"  Real collapsed-edge distance (km): min={small_edges_df['distance_km'].min():.2f}, "
              f"median={small_edges_df['distance_km'].median():.2f}, max={small_edges_df['distance_km'].max():.2f}")
        print(f"  Real travel time (days): min={days.min():.3f}, median={np.median(days):.3f}, "
              f"mean={days.mean():.3f}, max={days.max():.3f}, p90={np.percentile(days, 90):.3f}, "
              f"p99={np.percentile(days, 99):.3f}")
        print(f"  Real lag (whole {args.timestep_hours:.0f}h timesteps): "
              f"min={lag_steps.min()}, median={int(np.median(lag_steps))}, max={lag_steps.max()}")
        all_lags_days.extend(days.tolist())

    if not all_lags_days:
        print("No real collapsed edges found in either basin -- nothing to report.")
        return

    all_lags_days = np.array(all_lags_days)
    print()
    print("=" * 70)
    print("Real combined lag distribution, both basins (real collapsed training graph)")
    print("=" * 70)
    for count in [1, 2, 3, 5, 7, 10, 14, 21]:
        coverage = float((all_lags_days <= count).mean() * 100)
        print(f"ROUTING_HORIZON_COUNT={count:>2}: covers {coverage:.1f}% of real collapsed edges' own real travel time")


if __name__ == "__main__":
    main()