src/data/river_centerline.py

"""
Real river centerline geometry, digitized from road-map traces (traced
polylines over an OSM-style map, exported as PNG, then extracted via
color-threshold + skeletonization + shortest-path ordering, and
georeferenced using known town locations visible on each map).

This REPLACES the straight-line-between-gauges approximation in
river_line.py with the actual river shape. Gauges are "snapped" onto the
nearest point of the real centerline so we can measure true along-river
distance instead of great-circle distance between two points.

ACCURACY NOTE: georeferencing used ~6-7 manually-read reference points
per map (town label positions cross-referenced against their real
coordinates), fit with a least-squares affine transform. Residuals at
the reference points were on the order of a few km — good enough to
turn "straight line" into "recognizably the right river shape", but
not survey-grade. If precise positions matter (e.g. for the physics-
informed model), replace these CSVs with a real vector source (IGN BD
TOPO hydrography or similar) later; every function below only depends
on having a DataFrame of [longitude, latitude] points in order, so the
swap doesn't require touching calling code.
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

from .river_graph import _haversine_km


@dataclass
class CenterlinePoint:
    """A point on (or projected onto) a river centerline."""
    longitude: float
    latitude: float
    distance_from_mouth_km: float   # cumulative distance along the line from its first vertex
    total_length_km: float


def load_centerline(csv_path: Path) -> pd.DataFrame:
    """Load a digitized centerline CSV: [seq, longitude, latitude], ordered
    from one end of the river to the other (mouth -> source or vice versa,
    whichever the source trace started at)."""
    df = pd.read_csv(csv_path).sort_values("seq").reset_index(drop=True)
    return df[["longitude", "latitude"]]


def cumulative_distance_km(centerline: pd.DataFrame) -> np.ndarray:
    """Cumulative along-line distance (km) at each vertex, starting at 0."""
    lons, lats = centerline["longitude"].values, centerline["latitude"].values
    dists = [0.0]
    for i in range(1, len(centerline)):
        dists.append(dists[-1] + _haversine_km(lats[i-1], lons[i-1], lats[i], lons[i]))
    return np.array(dists)


def interpolate_by_fraction(centerline: pd.DataFrame, fraction: float) -> CenterlinePoint:
    """Interpolate a point at `fraction` (0-1) of the centerline's total length."""
    fraction = float(np.clip(fraction, 0.0, 1.0))
    cum = cumulative_distance_km(centerline)
    total = cum[-1]
    target = fraction * total

    idx = np.searchsorted(cum, target)
    idx = max(1, min(idx, len(cum) - 1))
    seg_start, seg_end = cum[idx-1], cum[idx]
    seg_frac = 0.0 if seg_end == seg_start else (target - seg_start) / (seg_end - seg_start)

    lon = centerline["longitude"].iloc[idx-1] + seg_frac * (
        centerline["longitude"].iloc[idx] - centerline["longitude"].iloc[idx-1])
    lat = centerline["latitude"].iloc[idx-1] + seg_frac * (
        centerline["latitude"].iloc[idx] - centerline["latitude"].iloc[idx-1])

    return CenterlinePoint(longitude=lon, latitude=lat,
                            distance_from_mouth_km=target, total_length_km=total)


def nearest_point_on_centerline(
    centerline: pd.DataFrame, latitude: float, longitude: float
) -> CenterlinePoint:
    """Project an arbitrary (e.g. clicked) lat/lon onto the nearest point
    on the centerline, returning its position along the line."""
    cum = cumulative_distance_km(centerline)
    lons, lats = centerline["longitude"].values, centerline["latitude"].values

    best = None
    for i in range(len(centerline) - 1):
        ax, ay = lons[i], lats[i]
        bx, by = lons[i+1], lats[i+1]
        px, py = longitude, latitude
        abx, aby = bx - ax, by - ay
        denom = abx**2 + aby**2
        t = 0.0 if denom == 0 else np.clip(((px-ax)*abx + (py-ay)*aby) / denom, 0.0, 1.0)
        proj_lon, proj_lat = ax + t*abx, ay + t*aby
        dist_to_line = _haversine_km(latitude, longitude, proj_lat, proj_lon)
        if best is None or dist_to_line < best[0]:
            seg_len = _haversine_km(ay, ax, by, bx)
            along_km = cum[i] + t * seg_len
            best = (dist_to_line, proj_lon, proj_lat, along_km)

    _, proj_lon, proj_lat, along_km = best
    return CenterlinePoint(longitude=proj_lon, latitude=proj_lat,
                            distance_from_mouth_km=along_km, total_length_km=cum[-1])


def snap_gauges_to_centerline(
    centerline: pd.DataFrame, nodes_df: pd.DataFrame
) -> pd.DataFrame:
    """
    For each gauge station in `nodes_df`, find its nearest point on the
    real centerline and record that position's along-line distance.
    Adds columns: centerline_lon, centerline_lat, centerline_km, snap_distance_km
    (snap_distance_km is how far the gauge's actual coordinates are from
    the digitized line — a large value flags a likely-misplaced gauge or
    georeferencing error worth checking).
    """
    out = nodes_df.copy()
    snapped = out.apply(
        lambda r: nearest_point_on_centerline(centerline, r["latitude"], r["longitude"]),
        axis=1,
    )
    out["centerline_lon"] = [p.longitude for p in snapped]
    out["centerline_lat"] = [p.latitude for p in snapped]
    out["centerline_km"] = [p.distance_from_mouth_km for p in snapped]
    out["snap_distance_km"] = [
        _haversine_km(r["latitude"], r["longitude"], p.latitude, p.longitude)
        for (_, r), p in zip(out.iterrows(), snapped)
    ]
    return out.sort_values("centerline_km").reset_index(drop=True)


def resample_centerline_for_clicks(centerline: pd.DataFrame, gauges_snapped: pd.DataFrame, n_points: int = 300) -> pd.DataFrame:
    """
    Densely resample a centerline into `n_points` evenly-spaced-by-distance
    points, each with precomputed elevation — used as click targets for
    the Plotly chart (Streamlit's chart-click selection fires on marker
    points, so we need many closely-spaced markers along the line for
    clicking to feel continuous).

    Returns:
        DataFrame [longitude, latitude, distance_from_mouth_km, elevation_m]
    """
    total_km = cumulative_distance_km(centerline)[-1]
    rows = []
    for i in range(n_points):
        frac = i / (n_points - 1)
        cp = interpolate_by_fraction(centerline, frac)
        elev = elevation_at_km(gauges_snapped, cp.distance_from_mouth_km)
        rows.append({
            "longitude": cp.longitude, "latitude": cp.latitude,
            "distance_from_mouth_km": cp.distance_from_mouth_km, "elevation_m": elev,
        })
    return pd.DataFrame(rows)


def elevation_at_km(gauges_snapped: pd.DataFrame, km: float) -> Optional[float]:
    """
    Piecewise-linear elevation estimate at a given along-centerline
    distance, interpolated between the two nearest snapped gauges (by
    centerline_km). Clamps to the nearest gauge's elevation outside the
    gauged range rather than extrapolating.

    Args:
        gauges_snapped: output of `snap_gauges_to_centerline`, must be
            sorted by centerline_km (it is, by default).
        km: along-centerline distance from the mouth.

    Returns:
        Interpolated elevation in meters, or None if no gauges are available.
    """
    g = gauges_snapped.dropna(subset=["elevation_m"]).sort_values("centerline_km")
    if g.empty:
        return None
    if km <= g["centerline_km"].iloc[0]:
        return float(g["elevation_m"].iloc[0])
    if km >= g["centerline_km"].iloc[-1]:
        return float(g["elevation_m"].iloc[-1])

    for i in range(len(g) - 1):
        a, b = g.iloc[i], g.iloc[i+1]
        if a["centerline_km"] <= km <= b["centerline_km"]:
            span = b["centerline_km"] - a["centerline_km"]
            t = 0.0 if span == 0 else (km - a["centerline_km"]) / span
            return float(a["elevation_m"] + t * (b["elevation_m"] - a["elevation_m"]))
    return None