src/data/river_graph.py

"""
River basin graph construction for La Eure and La Risle.

Ordering is elevation-first (falling back to latitude per-basin when any
station in that basin is missing elevation — see `order_stations_by_elevation`),
and edges are built by simply connecting consecutive stations in that order
within each basin. This guarantees every edge in the graph is monotonically
downhill by construction: there is no way for a "downstream" edge to point
to a higher point, because the ordering itself defines what "downstream"
means (previous versions of this module inferred edges from a nearest-lower-
elevation-within-radius heuristic instead, which could leave stations
disconnected — this version favors always producing a connected path per
basin as the default).

VISUALIZATION NOTE: a lat/lon map plot with north up will show "downstream"
pointing UP the page whenever the river flows north (true for both La Eure
and La Risle in this dataset). That's geographically correct, not a bug —
but it reads as backwards at a glance. `plot_elevation_profile` avoids this
ambiguity entirely by plotting elevation directly: downstream always goes
down and to the right, regardless of the river's compass direction.

IMPORTANT — still a heuristic, not ground truth:
  - `basin_id` is inferred from the station code prefix (H4xx.. vs
    H6xx..). Confirm against real basin delineation (e.g. the watershed
    shapefile) before trusting it.
  - Connecting stations strictly by sorted order assumes each basin is a
    single unbranched chain. Real confluences (a basin with more than one
    headwater branch merging into a shared downstream reach) will be
    flattened into one path. If you have known branch structure, build
    edges per-branch and merge instead of relying on the whole-basin order.
"""
from dataclasses import dataclass
from typing import Optional, Dict, List, Tuple

import pandas as pd
import matplotlib.pyplot as plt

_PALETTE = ["#2E6F95", "#D1495B", "#3E8914", "#E1A730", "#7C6BAF", "#3D9992"]


def _style_axes(ax: plt.Axes, title: str, xlabel: str, ylabel: str) -> None:
    ax.set_title(title, fontsize=13, fontweight="bold", pad=12)
    ax.set_xlabel(xlabel, fontsize=10)
    ax.set_ylabel(ylabel, fontsize=10)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.grid(True, alpha=0.25, linewidth=0.6)
    ax.tick_params(labelsize=9)


def _haversine_km(lat1, lon1, lat2, lon2) -> float:
    import numpy as np
    R = 6371.0
    lat1, lon1, lat2, lon2 = map(np.radians, [lat1, lon1, lat2, lon2])
    dlat, dlon = lat2 - lat1, lon2 - lon1
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * np.sin(dlon / 2) ** 2
    return R * 2 * np.arcsin(np.sqrt(a))


def assign_basin_id(
    station_code: str,
    prefix_map: Optional[Dict[str, int]] = None,
) -> Optional[int]:
    """
    Infer basin_id from a Hub'Eau station code prefix.

    Default mapping (heuristic, based on observed codes in this project):
        'H4' -> 0  (La Eure)
        'H6' -> 1  (La Risle)
    """
    prefix_map = prefix_map or {"H4": 0, "H6": 1}
    for prefix, basin_id in prefix_map.items():
        if station_code.startswith(prefix):
            return basin_id
    return None


@dataclass
class StationNode:
    station_id: str
    basin_id: int
    latitude: float
    longitude: float
    elevation: Optional[float] = None


class NodeBuilder:
    @staticmethod
    def build_from_dataframe(
        elevations_df: pd.DataFrame,
        prefix_map: Optional[Dict[str, int]] = None,
    ) -> List[StationNode]:
        """
        Build StationNode objects from a station elevations DataFrame.
        `elevation_m` may contain NaN — those become `elevation=None`
        rather than being dropped, so `order_stations_by_elevation` can
        apply its per-basin latitude fallback instead of silently losing
        stations (the earlier version of this module dropped them).
        """
        nodes = []
        for _, row in elevations_df.iterrows():
            basin_id = assign_basin_id(row["station_code"], prefix_map)
            if basin_id is None:
                print(f"NodeBuilder: '{row['station_code']}' matched no basin prefix, skipping")
                continue
            elev = row.get("elevation_m")
            nodes.append(StationNode(
                station_id=row["station_code"],
                basin_id=basin_id,
                latitude=row["latitude"],
                longitude=row["longitude"],
                elevation=None if pd.isna(elev) else float(elev),
            ))
        return nodes


def order_stations_by_elevation(
    nodes: List[StationNode],
    descending: bool = True
) -> List[str]:
    """
    Order station IDs by elevation — a physically-grounded proxy for a station's position
    along a river's course (higher elevation = further upstream, lower = further
    downstream) — for use as `station_order` in sequential edge building.

    This matters because a station list's raw row order commonly does NOT follow the
    river's actual course (verified for this dataset: station_list.csv's row order jumps
    around in both latitude and elevation, not monotonic either way).

    Stations are grouped by basin_id first (so ordering is computed independently per
    river). Within each basin:
    - If every station in the basin has a known elevation, sort by elevation (the
      physically-grounded signal — holds regardless of which way a river flows on a map).
    - If any station in the basin is missing elevation (e.g. a topology-only station with
      no time-series data, using virtual/placeholder data), the WHOLE basin falls back to
      latitude instead. Mixing elevation and latitude values as sort keys within one basin
      would be comparing incompatible scales (elevation in meters vs. latitude in degrees)
      and silently produce garbage ordering — falling back per-basin, not per-node, avoids
      that. Note this latitude convention (ascending latitude = downstream) was verified
      against this dataset's two rivers, both of which flow south -> north; it is not a
      universal rule for all river networks.

    Args:
        nodes: StationNode objects (e.g. from NodeBuilder.build_from_dataframe / .nodes)
        descending: True for upstream -> downstream order (highest elevation / lowest
            latitude first, the convention `EdgeBuilder.build_sequential` expects).

    Returns:
        Station IDs ordered upstream -> downstream within each basin
    """
    basins: Dict[int, List[StationNode]] = {}
    for n in nodes:
        basins.setdefault(n.basin_id, []).append(n)

    ordered_ids: List[str] = []
    for basin_id in sorted(basins.keys()):
        group = basins[basin_id]
        has_full_elevation = all(n.elevation is not None for n in group)

        if has_full_elevation:
            group_sorted = sorted(group, key=lambda n: -n.elevation if descending else n.elevation)
        else:
            missing = sum(1 for n in group if n.elevation is None)
            print(
                f"order_stations_by_elevation: basin {basin_id} has {missing}/{len(group)} "
                f"station(s) without elevation — ordering this basin by latitude instead "
                f"(mixing elevation and latitude as one sort key would be inconsistent)."
            )
            group_sorted = sorted(group, key=lambda n: n.latitude if descending else -n.latitude)

        ordered_ids.extend(n.station_id for n in group_sorted)

    return ordered_ids


class EdgeBuilder:
    @staticmethod
    def build_sequential(nodes: List[StationNode], station_order: List[str]) -> pd.DataFrame:
        """
        Connect consecutive stations within `station_order`, one edge per
        adjacent pair that share a basin_id (a basin boundary in the order
        list breaks the chain rather than creating a cross-basin edge).

        Returns:
            edges_df: [source, target, basin_id, distance_km, elevation_drop_m]
        """
        by_id = {n.station_id: n for n in nodes}
        edges = []
        for a_id, b_id in zip(station_order, station_order[1:]):
            a, b = by_id[a_id], by_id[b_id]
            if a.basin_id != b.basin_id:
                continue  # basin boundary: don't chain across rivers
            elev_drop = None
            if a.elevation is not None and b.elevation is not None:
                elev_drop = round(a.elevation - b.elevation, 2)
            edges.append({
                "source": a.station_id,
                "target": b.station_id,
                "basin_id": a.basin_id,
                "distance_km": round(_haversine_km(a.latitude, a.longitude, b.latitude, b.longitude), 3),
                "elevation_drop_m": elev_drop,
            })
        return pd.DataFrame(edges)


def build_basin_graph(
    elevations_df: pd.DataFrame,
    prefix_map: Optional[Dict[str, int]] = None,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Build station_nodes and station_edges tables for both basins using
    elevation-ordered (latitude-fallback) sequential chaining.

    Args:
        elevations_df: [station_code, latitude, longitude, elevation_m]
            (as returned by StationElevationsLoader.load()). elevation_m
            may be NaN for some rows — those stations still get placed via
            the per-basin latitude fallback, not dropped.
        prefix_map: optional override for basin assignment, see `assign_basin_id`.

    Returns:
        (nodes_df, edges_df):
            nodes_df: [station_code, basin_id, latitude, longitude, elevation_m]
            edges_df: [source, target, basin_id, distance_km, elevation_drop_m]
    """
    nodes = NodeBuilder.build_from_dataframe(elevations_df, prefix_map)
    station_order = order_stations_by_elevation(nodes, descending=True)
    edges_df = EdgeBuilder.build_sequential(nodes, station_order)

    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])

    return nodes_df, edges_df


def save_graph(nodes_df: pd.DataFrame, edges_df: pd.DataFrame, output_dir) -> None:
    """Write station_nodes.csv and station_edges.csv."""
    from pathlib import Path
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    nodes_df.to_csv(output_dir / "station_nodes.csv", index=False)
    edges_df.to_csv(output_dir / "station_edges.csv", index=False)
    print(f"Saved {len(nodes_df)} nodes, {len(edges_df)} edges to {output_dir}")


def plot_basin_graphs(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    basin_names: Optional[Dict[int, str]] = None,
    figsize: tuple = (14, 7),
    save_path=None,
) -> plt.Figure:
    """
    Geographic map of both basin graphs: stations positioned by lon/lat,
    colored by elevation, with arrows showing downstream flow.

    Reminder: "downstream" will point UP the page for a north-flowing
    river — that's correct geography, not a bug. Use
    `plot_elevation_profile` if you want flow direction to always read
    top-to-bottom / left-to-right regardless of compass direction.
    """
    basin_names = basin_names or {b: f"Basin {b}" for b in nodes_df["basin_id"].unique()}
    basin_ids = sorted(nodes_df["basin_id"].unique())

    fig, axes = plt.subplots(1, len(basin_ids), figsize=figsize)
    if len(basin_ids) == 1:
        axes = [axes]

    for ax, basin_id in zip(axes, basin_ids):
        b_nodes = nodes_df[nodes_df["basin_id"] == basin_id]
        b_edges = edges_df[edges_df["basin_id"] == basin_id] if not edges_df.empty else edges_df

        for _, edge in b_edges.iterrows():
            src = b_nodes[b_nodes["station_code"] == edge["source"]].iloc[0]
            tgt = b_nodes[b_nodes["station_code"] == edge["target"]].iloc[0]
            ax.annotate(
                "", xy=(tgt["longitude"], tgt["latitude"]),
                xytext=(src["longitude"], src["latitude"]),
                arrowprops=dict(arrowstyle="-|>", color="#888888", lw=1.4,
                                 shrinkA=8, shrinkB=8),
            )

        scatter = ax.scatter(
            b_nodes["longitude"], b_nodes["latitude"], c=b_nodes["elevation_m"],
            cmap="terrain", s=140, edgecolor="black", linewidth=0.6, zorder=3,
        )
        for _, row in b_nodes.iterrows():
            ax.annotate(row["station_code"], (row["longitude"], row["latitude"]),
                        fontsize=6.5, xytext=(4, 4), textcoords="offset points")

        cbar = plt.colorbar(scatter, ax=ax, fraction=0.046, pad=0.04)
        cbar.set_label("Elevation (m)", fontsize=9)
        cbar.ax.tick_params(labelsize=8)

        _style_axes(ax, basin_names.get(basin_id, f"Basin {basin_id}"), "Longitude", "Latitude")
        ax.set_aspect("equal")

    plt.tight_layout()
    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Plot saved to {save_path}")

    return fig


def plot_elevation_profile(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    basin_names: Optional[Dict[int, str]] = None,
    figsize: tuple = (11, 6),
    save_path=None,
) -> plt.Axes:
    """
    Schematic river profile: x = downstream sequence position, y = elevation.
    Unlike the geographic map, this is unambiguous regardless of which
    compass direction the river actually flows — downstream always reads
    left-to-right and (by construction, since edges are built from sorted
    order) downhill.
    """
    basin_names = basin_names or {b: f"Basin {b}" for b in nodes_df["basin_id"].unique()}
    n = nodes_df.set_index("station_code")

    fig, ax = plt.subplots(figsize=figsize)
    for i, basin_id in enumerate(sorted(nodes_df["basin_id"].unique())):
        b_edges = edges_df[edges_df["basin_id"] == basin_id]
        if b_edges.empty:
            continue
        # Walk the chain in edge order to get sequence position.
        chain = [b_edges.iloc[0]["source"]] + b_edges["target"].tolist()
        elevations = [n.loc[s, "elevation_m"] for s in chain]
        color = _PALETTE[i % len(_PALETTE)]

        ax.plot(range(len(chain)), elevations, marker="o", markersize=7,
                 linewidth=2, color=color, label=basin_names.get(basin_id, f"Basin {basin_id}"))
        for x, (station, elev) in enumerate(zip(chain, elevations)):
            ax.annotate(station, (x, elev), fontsize=6.5, rotation=45,
                        xytext=(0, 8), textcoords="offset points", ha="left")

    _style_axes(ax, "River Elevation Profile (upstream → downstream)",
                "Position along river (upstream → downstream)", "Elevation (m)")
    ax.set_xticks([])
    ax.legend(fontsize=9, frameon=False)
    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Plot saved to {save_path}")

    return ax