src/data/river_line.py

"""
Continuous-position interpolation along the river graph.

APPROXIMATION NOTICE: without a real river centerline shapefile, each
graph edge (station -> station) is treated as a straight line between
the two gauges. This is a placeholder for the true meandering river
path. Once a real centerline (e.g. IGN BD TOPO hydrography linestrings)
is available, replace `interpolate_point` with a lookup against the
actual geometry, snapped to the nearest point on the line and measured
along its length instead of great-circle distance between endpoints.

Everything downstream of this module (the Gradio app) is written
against the `RiverPoint` interface below, so swapping the underlying
geometry source later shouldn't require changing the UI.
"""
from dataclasses import dataclass
from typing import Optional, List
import numpy as np
import pandas as pd

from .river_graph import _haversine_km


@dataclass
class RiverPoint:
    """A queried location along the (approximated) river line."""
    basin_id: int
    upstream_station: str
    downstream_station: str
    fraction: float          # 0 = at upstream station, 1 = at downstream station
    latitude: float
    longitude: float
    elevation_m: float
    distance_from_upstream_km: float
    edge_length_km: float


def edges_as_segments(nodes_df: pd.DataFrame, edges_df: pd.DataFrame) -> pd.DataFrame:
    """Attach endpoint coordinates/elevation to each edge, for interpolation and plotting."""
    n = nodes_df.set_index("station_code")
    seg = edges_df.copy()
    seg["lat1"] = seg["source"].map(n["latitude"])
    seg["lon1"] = seg["source"].map(n["longitude"])
    seg["elev1"] = seg["source"].map(n["elevation_m"])
    seg["lat2"] = seg["target"].map(n["latitude"])
    seg["lon2"] = seg["target"].map(n["longitude"])
    seg["elev2"] = seg["target"].map(n["elevation_m"])
    return seg


def interpolate_point(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    source: str,
    target: str,
    fraction: float,
) -> RiverPoint:
    """
    Interpolate a point a given fraction of the way along one edge
    (straight line between the two station endpoints).

    Args:
        nodes_df, edges_df: from `build_basin_graph`.
        source: upstream station code of the edge.
        target: downstream station code of the edge.
        fraction: 0.0 (at `source`) to 1.0 (at `target`).

    Returns:
        A RiverPoint with interpolated lat/lon/elevation.
    """
    fraction = float(np.clip(fraction, 0.0, 1.0))
    edge = edges_df[(edges_df["source"] == source) & (edges_df["target"] == target)]
    if edge.empty:
        raise ValueError(f"No edge {source} -> {target} in the graph")
    edge = edge.iloc[0]

    n = nodes_df.set_index("station_code")
    up, down = n.loc[source], n.loc[target]

    lat = up["latitude"] + fraction * (down["latitude"] - up["latitude"])
    lon = up["longitude"] + fraction * (down["longitude"] - up["longitude"])
    elev = up["elevation_m"] + fraction * (down["elevation_m"] - up["elevation_m"])
    edge_len = edge["distance_km"]

    return RiverPoint(
        basin_id=int(edge["basin_id"]),
        upstream_station=source,
        downstream_station=target,
        fraction=fraction,
        latitude=lat,
        longitude=lon,
        elevation_m=elev,
        distance_from_upstream_km=fraction * edge_len,
        edge_length_km=edge_len,
    )


def nearest_point_on_graph(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    latitude: float,
    longitude: float,
) -> RiverPoint:
    """
    Given an arbitrary (e.g. map-clicked) lat/lon, find the closest point
    on any edge's straight-line segment and return it as a RiverPoint.
    Used to translate a free-form click into a position on the graph.
    """
    seg = edges_as_segments(nodes_df, edges_df)
    if seg.empty:
        raise ValueError("Graph has no edges to project onto")

    best = None
    for _, row in seg.iterrows():
        # Project (latitude, longitude) onto the segment in a simple
        # equirectangular local approximation (fine at this spatial scale).
        ax, ay = row["lon1"], row["lat1"]
        bx, by = row["lon2"], row["lat2"]
        px, py = longitude, latitude

        abx, aby = bx - ax, by - ay
        denom = abx ** 2 + aby ** 2
        t = 0.0 if denom == 0 else ((px - ax) * abx + (py - ay) * aby) / denom
        t = float(np.clip(t, 0.0, 1.0))

        proj_lon = ax + t * abx
        proj_lat = ay + t * aby
        dist_km = _haversine_km(latitude, longitude, proj_lat, proj_lon)

        if best is None or dist_km < best[0]:
            best = (dist_km, row["source"], row["target"], t)

    _, source, target, t = best
    return interpolate_point(nodes_df, edges_df, source, target, t)


def interpolate_along_chain(
    nodes_df: pd.DataFrame,
    edges_df: pd.DataFrame,
    basin_id: int,
    global_fraction: float,
) -> RiverPoint:
    """
    Interpolate a point by a single 0-1 fraction of the ENTIRE basin chain's
    length (as opposed to `interpolate_point`, which takes a fraction of one
    edge). Convenient for a single slider spanning source -> outlet.

    Args:
        nodes_df, edges_df: from `build_basin_graph`.
        basin_id: which basin's chain to walk.
        global_fraction: 0.0 (source) to 1.0 (outlet).

    Returns:
        A RiverPoint at that position along the chain.
    """
    global_fraction = float(np.clip(global_fraction, 0.0, 1.0))
    b_edges = edges_df[edges_df["basin_id"] == basin_id].reset_index(drop=True)
    if b_edges.empty:
        raise ValueError(f"No edges found for basin_id={basin_id}")

    total_km = b_edges["distance_km"].sum()
    target_km = global_fraction * total_km

    cumulative = 0.0
    for _, edge in b_edges.iterrows():
        if cumulative + edge["distance_km"] >= target_km or edge is b_edges.iloc[-1]:
            local_fraction = 0.0 if edge["distance_km"] == 0 else \
                (target_km - cumulative) / edge["distance_km"]
            return interpolate_point(
                nodes_df, edges_df, edge["source"], edge["target"],
                float(np.clip(local_fraction, 0.0, 1.0)),
            )
        cumulative += edge["distance_km"]

    # Fallback: end of chain.
    last = b_edges.iloc[-1]
    return interpolate_point(nodes_df, edges_df, last["source"], last["target"], 1.0)


def chain_total_length_km(edges_df: pd.DataFrame, basin_id: int) -> float:
    """Total length (sum of edge distances) of a basin's chain, in km."""
    return float(edges_df[edges_df["basin_id"] == basin_id]["distance_km"].sum())


def groundwater_at_point(
    point: RiverPoint,
    groundwater_df: pd.DataFrame,
    max_distance_km: float = 20.0,
) -> tuple[Optional[float], float]:
    """
    Estimate groundwater level at a RiverPoint by averaging nearby wells
    (raw ADES well data with [lat, lon, groundwater_level_m]), inverse-
    distance weighted.

    Args:
        point: a RiverPoint (e.g. from `interpolate_point`).
        groundwater_df: raw ADES well data with columns
            [code_bss, lat, lon, groundwater_level_m] (most recent
            reading per well, or already time-filtered by the caller).
        max_distance_km: only wells within this radius are considered.

    Returns:
        (estimate, nearest_well_km): the inverse-distance-weighted
        groundwater level in meters (None if no wells are within range),
        and the distance to the single nearest well regardless of range
        — so callers can explain *why* an estimate is missing (e.g.
        "nearest well is 34 km away") instead of just showing a blank.
    """
    required = {"lat", "lon", "groundwater_level_m"}
    if not required.issubset(groundwater_df.columns):
        return None, float("inf")

    df = groundwater_df.dropna(subset=list(required)).copy()
    if df.empty:
        return None, float("inf")

    df["distance_km"] = df.apply(
        lambda r: _haversine_km(point.latitude, point.longitude, r["lat"], r["lon"]), axis=1
    )
    nearest_km = float(df["distance_km"].min())
    nearby = df[df["distance_km"] <= max_distance_km]
    if nearby.empty:
        return None, nearest_km

    weights = 1.0 / np.maximum(nearby["distance_km"], 0.1)
    estimate = float((nearby["groundwater_level_m"] * weights).sum() / weights.sum())
    return estimate, nearest_km