src/graph/dynamic_features.py

"""
Builds genuine [n_nodes, T] time series for the features that
node_features.py's add_safran_features/add_groundwater_features/
add_hydrometric_target_stats only ever hand over as a single period-
aggregated number. That aggregation was always a real, named
limitation (see build_graph.py's DYNAMIC_PREFIXES docstring and the
README's §2.3) -- this module is what actually closes it.

physics_losses.py's routing_consistency_loss specifically needs a real
time dimension and had nothing to consume before this existed.

Three sources, three different real challenges:
  - discharge: already daily, already per-station -- just needs
    pivoting to wide form, restricted to real gauges (no fabrication
    for non-gauge nodes).
  - groundwater: wells report on wildly irregular schedules (confirmed
    against real data: 13 different "latest dates" among 18 wells
    within 20km of one station). Needs resampling to a common daily
    grid via forward-fill (a well's level changes slowly -- carrying
    the last known reading forward is standard practice for sparse
    level data, not an invented shortcut) BEFORE the spatial average,
    not after -- averaging raw irregular readings per exact calendar
    date is exactly the bug that made the old "aggregate_to_stations"
    path undercount real well coverage by 5-10x (see node_features.py's
    add_groundwater_features docstring).
  - climate: SAFRANLoader.load() already returns per-date rows before
    node_features.py's add_safran_features aggregates them -- this
    just skips that aggregation step and pivots instead. UNTESTED in
    this environment: no real ERA5/safran files were available here to
    validate against (only real ADES and hydrometric data were).
"""
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd


def build_discharge_timeseries(
    nodes_df: pd.DataFrame, hydrometric_path: Path, date_range: Tuple[str, str],
    freq: str = "D",
) -> pd.DataFrame:
    """
    Real daily QmnJ discharge, pivoted to [date x station_code] wide
    form, restricted to the date range and reindexed onto a regular
    daily grid (missing days stay NaN -- no fabrication).

    Deliberately does NOT forward-fill discharge the way groundwater
    gets resampled below: discharge genuinely changes day to day and a
    missing daily reading should stay missing, not be papered over with
    the previous day's value the way a slowly-changing water table can
    reasonably be.

    Only real gauge station_codes appear as columns -- confluences and
    virtual nodes never had discharge observations to begin with, same
    principle as the period-aggregate target_* columns.
    """
    from ..data.loaders.hydrometric import HydrometricLoader

    loader = HydrometricLoader(data_path=hydrometric_path)
    df = loader.load()
    if "discharge_m3s" not in df.columns:
        raise ValueError("Loaded hydrometric data has no discharge_m3s column")

    start, end = date_range
    df = df[(df["date"] >= start) & (df["date"] <= end)]
    wide = df.pivot_table(index="date", columns="station_code", values="discharge_m3s", aggfunc="mean")

    full_index = pd.date_range(start, end, freq=freq)
    wide = wide.reindex(full_index)
    wide.index.name = "date"
    return wide


def map_stations_to_precipitation_gauges(nodes_df: pd.DataFrame, precip_stations: pd.DataFrame) -> pd.Series:
    """
    Which real Météo-France precipitation station each graph node reads
    from -- nearest by real haversine distance, same reasoning and same
    underlying helper as map_stations_to_grid_cells (the ERA5 grid-cell
    assignment): computed once, since the spatial relationship doesn't
    change over time, and reused for every date.

    Unlike the ERA5 grid (a uniform ~20-30 cell mesh with predictable
    spacing), real precipitation gauges are irregularly placed --
    genuinely possible for the nearest one to be tens of km away in a
    sparsely-instrumented area. No distance cap is applied here: an
    imperfect nearest real gauge is still real, directly-measured data,
    which was the whole reason to prefer this over ERA5 reanalysis in
    the first place -- capping it back down to "use reanalysis instead
    beyond some radius" would undo that.

    Args:
        nodes_df: real graph nodes with latitude/longitude columns.
        precip_stations: real Météo-France stations with NUM_POSTE,
            LAT, LON columns (as returned by
            fetch_meteofrance_precipitation.py's combined output).

    Returns:
        Series indexed like nodes_df, values are real NUM_POSTE station
        codes.
    """
    unique_stations = precip_stations.drop_duplicates(subset="NUM_POSTE")
    station_lat = unique_stations["LAT"].values
    station_lon = unique_stations["LON"].values
    station_ids = unique_stations["NUM_POSTE"].values

    assignments = []
    for _, row in nodes_df.iterrows():
        dist = _haversine_km_vec(row["latitude"], row["longitude"], station_lat, station_lon)
        assignments.append(station_ids[dist.argmin()])
    return pd.Series(assignments, index=nodes_df.index, name="precip_station_id")


def compute_rolling_precip_sum(precip_tensor: np.ndarray, window_days: int = 30, min_coverage: float = 0.5) -> np.ndarray:
    """
    Real rolling sum of precipitation over the trailing window_days,
    per node per date -- the "recent rainfall" side of the specific-
    discharge relationship below. A window built from mostly-missing
    days is treated as unreliable (NaN), not silently understated,
    same completeness discipline already used for the per-example
    lookback-window precip sum in train_spatiotemporal_gnn.py and for
    aggregate_anomalies in the climate work.
    """
    n_nodes, n_dates = precip_tensor.shape
    rolling_sum = np.full((n_nodes, n_dates), np.nan, dtype=np.float32)
    for t in range(window_days, n_dates):
        window = precip_tensor[:, t - window_days:t]
        real_fraction = (~np.isnan(window)).mean(axis=1)
        window_sum = np.nansum(window, axis=1)
        rolling_sum[:, t] = np.where(real_fraction >= min_coverage, window_sum, np.nan)
    return rolling_sum


def fit_specific_discharge_coefficient(
    discharge_tensor: np.ndarray, rolling_precip_sum: np.ndarray, catchment_area_km2: np.ndarray,
    dates, train_end: str,
) -> float:
    """
    Fits k in Q ~= k * (recent precipitation sum) * (catchment area) --
    a standard, well-established hydrological technique for estimating
    discharge at ungauged points ("regionalization" via specific
    discharge scaling), not a novel or speculative one. Fit ONLY on
    real (Q, precip, area) triples at GAUGED stations during the real
    TRAINING period specifically -- same "fit on train only" discipline
    as standardization, avoiding any leakage from val/test into a
    coefficient that then gets applied everywhere, including at
    ungauged nodes evaluated during val/test.

    Closed-form least squares through the origin (no intercept -- zero
    rainfall over the window should imply ~zero attributable runoff,
    a physically reasonable constraint, not an arbitrary modeling
    choice): k = sum(x*y) / sum(x^2), x = precip_sum * area, y = Q.

    Returns:
        The fitted real scalar k. Raises if there's no real data to
        fit against, rather than returning a meaningless default.
    """
    train_col_mask = pd.DatetimeIndex(dates) <= pd.Timestamp(train_end)
    Q_train = discharge_tensor[:, train_col_mask]
    P_train = rolling_precip_sum[:, train_col_mask]
    A = catchment_area_km2[:, None] * np.ones_like(Q_train)  # broadcast area across all real dates

    x = P_train * A
    y = Q_train
    valid = ~np.isnan(x) & ~np.isnan(y) & (x != 0)
    n_real = int(valid.sum())
    if n_real == 0:
        raise ValueError("No real (discharge, precipitation, catchment area) triples available to fit "
                          "the specific-discharge coefficient -- check that all three are real for at "
                          "least some gauged station during the training period.")

    x_valid, y_valid = x[valid], y[valid]
    k = float(np.sum(x_valid * y_valid) / np.sum(x_valid ** 2))
    print(f"Fitted specific-discharge coefficient k={k:.6g} on {n_real} real (Q, precip, area) "
          f"triple(s) from gauged stations, training period only")
    return k


def estimate_discharge_from_precip(
    rolling_precip_sum: np.ndarray, catchment_area_km2: np.ndarray, k: float,
) -> np.ndarray:
    """
    Applies the fitted specific-discharge coefficient to EVERY node
    (gauged or not), wherever real precipitation and catchment area
    both exist -- this is the actual point: an ungauged node with real
    precipitation and real catchment area gets a real, physically-
    motivated discharge ESTIMATE, computed the same way for every node
    rather than only where discharge happens to already be known.

    This is deliberately an INPUT feature, never a target -- it must
    never be treated as if it were a real observation. It's fit from
    real data and grounded in a real, standard hydrological
    relationship, but it is still an estimate, not a measurement, and
    should reach the model exactly the way every other real-but-
    imperfect input does: as a value the model can learn to weigh,
    never as something the supervised loss is held accountable to
    match exactly.
    """
    n_nodes = catchment_area_km2.shape[0]
    area_broadcast = catchment_area_km2[:, None] * np.ones((n_nodes, rolling_precip_sum.shape[1]))
    estimated = k * rolling_precip_sum * area_broadcast
    # NaN propagates automatically wherever rolling_precip_sum or
    # catchment_area_km2 is real NaN -- no explicit masking needed here,
    # multiplication by NaN is already NaN.
    return estimated.astype(np.float32)


FORECAST_LEAD_TIMES = [1, 2, 3, 5, 7]  # real confirmed coverage in previous_runs_forecast.csv -- NOT every horizon this project predicts at


def build_forecast_tensor(
    nodes_df: pd.DataFrame, forecast_csv_path: Path, anchor_dates, horizons: List[int],
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Real forecasted precipitation, per (anchor_date, node, horizon) --
    built from previous_runs_forecast.csv, whose real confirmed schema
    is genuinely different from every other timeseries builder in this
    file: it's not a simple [date x station] wide table, it's one row
    per real OBSERVATION date with several "precipitation_previous_
    dayN" columns, each holding what was FORECASTED N days before that
    date, for that date.

    The real mapping this requires, worked through explicitly: for an
    anchor date t and horizon h, the forecast we want is "what was
    predicted h days in advance, for date t+h" -- which was ISSUED on
    date (t+h) - h = t, exactly our anchor. In this file's real column
    convention, that value lives in the ROW for date=t+h, in column
    precipitation_previous_day{h}. Getting this backwards (e.g. using
    the row for date=t instead) would silently use a forecast for the
    wrong target date entirely.

    station_code in this real file already matches this project's real
    graph node station codes directly (confirmed against real sample
    data) -- no nearest-neighbor mapping needed here, unlike the
    Meteo-France precipitation gauges, which come from a different
    network entirely.

    Real, confirmed constraints, not assumptions: only horizons in
    FORECAST_LEAD_TIMES ([1,2,3,5,7]) have any real forecast data at
    all -- every other horizon this project predicts at (4, 10, 30, 60,
    90, 180) has NONE, and gets missing=True unconditionally, not
    silently filled with anything. Real coverage also only starts
    2024-03-01 -- anchor dates before that (the large majority of this
    project's 2013+ training period) get missing=True for every
    horizon, real lead time or not.

    Returns:
        forecast_tensor, missing_tensor: both [n_examples, n_nodes,
        n_horizons], real values (or 0.0 placeholder) and an explicit
        real/missing flag respectively -- same value+missingness
        pattern used for every other real-but-incomplete input in this
        project.
    """
    df = pd.read_csv(forecast_csv_path)
    df["date"] = pd.to_datetime(df["date"])

    # Pre-indexed for fast repeated lookup -- {(date, station_code): row}
    # rather than re-filtering the real dataframe for every single
    # (anchor_date, node, horizon) combination, which would be
    # prohibitively slow at this project's real scale (thousands of
    # examples x 100 nodes x 10 horizons).
    df_indexed = df.set_index(["date", "station_code"])

    station_codes = nodes_df["station_code"].tolist()
    n_nodes = len(station_codes)
    n_examples = len(anchor_dates)
    n_horizons = len(horizons)

    forecast_tensor = np.zeros((n_examples, n_nodes, n_horizons), dtype=np.float32)
    missing_tensor = np.ones((n_examples, n_nodes, n_horizons), dtype=np.float32)  # starts fully missing

    for h_idx, h in enumerate(horizons):
        if h not in FORECAST_LEAD_TIMES:
            continue  # no real forecast data for this horizon at all -- stays fully missing
        col = f"precipitation_previous_day{h}"
        for ex_idx, anchor in enumerate(anchor_dates):
            target_date = pd.Timestamp(anchor) + pd.Timedelta(days=h)
            for node_idx, station_code in enumerate(station_codes):
                try:
                    value = df_indexed.loc[(target_date, station_code), col]
                except KeyError:
                    continue  # genuinely no real row for this date/station -- stays missing
                if pd.notna(value):
                    forecast_tensor[ex_idx, node_idx, h_idx] = value
                    missing_tensor[ex_idx, node_idx, h_idx] = 0.0

    return forecast_tensor, missing_tensor


def compute_days_since_last_real(tensor, sentinel_days=365.0):
    """
    For each [node, date] position, how many real days have passed
    since the most recent real (non-NaN) observation at or before that
    date -- 0 if the current day itself is real. Built specifically
    because a binary missing-flag alone can't distinguish a single-day
    sensor glitch from a station offline for months; both look
    identical to the model at any given missing timestep without this.

    Computed over the FULL real date range, not per-window -- correctly
    looks backward past wherever a later 30-day lookback window happens
    to start, rather than resetting to "unknown" the moment a window
    boundary is crossed. This is exactly why it has to be built once
    here, before windowing, and then sliced the same way value/missing
    already are, rather than computed fresh inside each window.

    Nodes with NO real observation at all up through a given date get
    the real sentinel value (a fixed ceiling, not literal infinity) --
    "genuinely never observed" is a real, distinct case from "observed
    365+ days ago", but both should read as "don't trust this" to the
    model without a numerically unbounded value reaching it.
    """
    import numpy as np
    n_nodes, n_dates = tensor.shape
    is_real = ~np.isnan(tensor)
    date_indices = np.arange(n_dates, dtype=np.float64)

    idx_where_real = np.where(is_real, date_indices[None, :], -1.0)
    last_real_idx = np.maximum.accumulate(idx_where_real, axis=1)

    dt = np.where(last_real_idx < 0, sentinel_days, date_indices[None, :] - last_real_idx)
    dt = np.minimum(dt, sentinel_days)
    return dt.astype(np.float32)


def compress_days_since_last_real(dt):
    """
    log(1 + dt) -- compresses the real scale so a 3-vs-10-day
    distinction isn't drowned out by a 3-vs-3000-day one the way a raw
    linear day count would. Same unstandardized-large-number
    instability this project already hit once with raw discharge in
    L/s, applied preemptively here rather than found the hard way again.
    """
    import numpy as np
    return np.log1p(dt).astype(np.float32)


def compute_forward_filled_tensor(tensor):
    """
    Real forward-filled value per [node, date] position: the most
    recent real (non-NaN) observation at or before that date, NaN if
    no real observation has occurred yet anywhere before it. Causally
    safe by construction -- forward-fill only ever looks backward in
    time, so unlike a fitted statistic (a per-node mean/median, which
    needs a "fit on train only" discipline to avoid leaking future
    information into the fill), this needs no train/val/test split
    guard at all: a forward-filled value at any date, in any split,
    only ever depends on dates strictly before it.

    Built specifically as a fill-value SOURCE for real gaps during a
    real low- (or high-) flow episode, where a station-wide central-
    tendency fill (compute_per_node_historical_median, in
    physics_losses.py) is itself the wrong reference point --
    confirmed directly via a real false-spike investigation: even a
    real per-node MEDIAN fill (1142.0, for station H605641401) was
    still ~17x that station's real low-flow values (57-70 L/s) during
    a real gap that fell IN THE MIDDLE of a real low-flow episode,
    because a median describes a station's typical day, not
    necessarily the day immediately preceding a specific gap. A real,
    direct ablation (zeroing the missingness flag and days-since
    channel for this exact example) confirmed the flag mechanism alone
    cannot fix this: even with the model given no signal at all that
    the value was filled, the median-based fill still produced a
    predicted 0.95 quantile ~28x the real observed value, because the
    filled VALUE itself (1142.0) was implausible for the real
    conditions at that moment, not because the model failed to
    discount a flagged fill. Forward-fill instead carries the real,
    most-recently-known value forward, directly informative about
    real LOCAL, CURRENT conditions in a way a station-wide statistic
    cannot be.

    Same accumulate-forward pattern already validated in
    compute_days_since_last_real -- computed over the FULL real date
    range, not per-window, for the same reason: correctly looks
    backward past wherever a later lookback window happens to start.
    """
    n_nodes, n_dates = tensor.shape
    is_real = ~np.isnan(tensor)
    date_indices = np.arange(n_dates)

    idx_where_real = np.where(is_real, date_indices[None, :], -1)
    last_real_idx = np.maximum.accumulate(idx_where_real, axis=1)

    filled = np.take_along_axis(tensor, np.clip(last_real_idx, 0, None), axis=1)
    filled = np.where(last_real_idx < 0, np.nan, filled)
    return filled.astype(np.float32)


def build_precipitation_timeseries(
    nodes_df: pd.DataFrame, precip_csv_path: Path, date_range: Tuple[str, str], freq: str = "D",
) -> pd.DataFrame:
    """
    Real daily precipitation (mm), pivoted to [date x station_code] wide
    form -- same output convention as build_discharge_timeseries and
    build_waterlevel_timeseries, so it drops into the same downstream
    pipeline (assemble_dynamic_tensor etc.) without special-casing.

    Real Météo-France gauge data (fetch_meteofrance_precipitation.py),
    not ERA5 reanalysis -- deliberately: precipitation varies sharply
    over short distances in ways a reanalysis grid cell can miss, so a
    real nearby gauge is more locally authoritative where one exists.
    Each of OUR graph nodes reads from its own nearest real gauge (see
    map_stations_to_precipitation_gauges), which can mean several of
    our nodes share the same real gauge if they're close together --
    expected and correct, not a bug, the same way several ERA5-grid
    climate stations already share one grid cell elsewhere in this
    project.
    """
    precip_df = pd.read_csv(precip_csv_path)
    precip_df["date"] = pd.to_datetime(precip_df["date"])

    # CRITICAL, CONFIRMED FIX -- exclude real "ghost" stations (present
    # in the real Météo-France station list, with real LAT/LON, but
    # ZERO real precip_mm_daily readings across every real row they
    # have) from the candidate pool BEFORE nearest-neighbor matching.
    # Confirmed directly against a real daily_precipitation.csv: 56 of
    # 110 real candidate stations (51%) have precip_mm_daily entirely
    # NaN -- a real fetch/ingestion gap in fetch_meteofrance_
    # precipitation.py, not a downstream bug -- yet
    # map_stations_to_precipitation_gauges has no liveness check, only
    # a distance check, so any real graph node whose real geographically
    # -nearest station happened to be one of these 56 real ghosts got
    # permanently, silently matched to a station that can never
    # contribute a single real reading. Confirmed as the real, direct
    # cause of 22 of 27 real gauge stations in one real evaluation run
    # having ZERO real precipitation anywhere in the full 2013-2026
    # record -- every one of forward-fill, the per-node median fallback,
    # AND compress_days_since_last_real's real 365-day sentinel were all
    # working exactly as designed; there was simply never a real
    # observation for any of them to find. Filtering here, once, before
    # the nearest-neighbor search, routes every real node to the
    # nearest real LIVE station instead -- a real, if sometimes farther,
    # station that can actually contribute real data, which is strictly
    # better than a real ghost at zero distance.
    stations_with_real_data = precip_df.loc[precip_df["precip_mm_daily"].notna(), "NUM_POSTE"].unique()
    n_candidates_before = precip_df["NUM_POSTE"].nunique()
    precip_stations_live = precip_df[precip_df["NUM_POSTE"].isin(stations_with_real_data)]
    n_candidates_after = precip_stations_live["NUM_POSTE"].nunique()
    if n_candidates_after < n_candidates_before:
        print(f"build_precipitation_timeseries: excluding {n_candidates_before - n_candidates_after} of "
              f"{n_candidates_before} real candidate precipitation station(s) with ZERO real "
              f"precip_mm_daily readings from nearest-neighbor matching, leaving "
              f"{n_candidates_after} real live station(s).")

    assignments = map_stations_to_precipitation_gauges(nodes_df, precip_stations_live[["NUM_POSTE", "LAT", "LON"]])

    start, end = date_range
    full_index = pd.date_range(start, end, freq=freq)
    wide = pd.DataFrame(index=full_index)
    wide.index.name = "date"

    for node_idx, station_id in assignments.items():
        station_code = nodes_df.loc[node_idx, "station_code"]
        station_series = precip_df[precip_df["NUM_POSTE"] == station_id].set_index("date")["precip_mm_daily"]
        wide[station_code] = station_series.reindex(full_index)

    return wide


def build_waterlevel_timeseries(
    nodes_df: pd.DataFrame, hydrometric_path: Path, date_range: Tuple[str, str],
    freq: str = "D",
) -> pd.DataFrame:
    """
    Real daily water level (HIXnJ), pivoted to [date x station_code]
    wide form -- mirrors build_discharge_timeseries exactly, using
    waterlevel_mm instead of discharge_m3s (both already present in
    HydrometricLoader's real output, confirmed directly earlier this
    project: columns ['date', 'station_code', 'discharge_m3s',
    'waterlevel_mm']).

    Confirmed real coverage: 13 of 27 real stations have water-level
    data at all -- a genuinely different, larger set than discharge's
    6-8 in-window stations, though not identical (some stations have
    one variable but not the other). Built specifically to give the
    model's water-level output channel real supervision for the first
    time -- previously it only had one loose consistency constraint
    (manning_consistency_loss) and no direct target at all.
    """
    from ..data.loaders.hydrometric import HydrometricLoader

    loader = HydrometricLoader(data_path=hydrometric_path)
    df = loader.load()
    if "waterlevel_mm" not in df.columns:
        raise ValueError("Loaded hydrometric data has no waterlevel_mm column")

    start, end = date_range
    df = df[(df["date"] >= start) & (df["date"] <= end)]
    wide = df.pivot_table(index="date", columns="station_code", values="waterlevel_mm", aggfunc="mean")

    full_index = pd.date_range(start, end, freq=freq)
    wide = wide.reindex(full_index)
    wide.index.name = "date"
    return wide


def build_groundwater_timeseries(
    nodes_df: pd.DataFrame, ades_path: Path, date_range: Tuple[str, str],
    max_distance_km: float = 20.0, freq: str = "D",
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """
    Real groundwater level, resampled to a daily grid per well
    (forward-fill) BEFORE spatial averaging, then averaged per node
    using the SAME nearby-well set every day -- the spatial join
    (which wells are near which node) is time-invariant, so it's
    computed once, not repeated per date. That's what keeps this
    tractable at real reach-graph scale: O(n_nodes) spatial lookups
    total, not O(n_nodes x n_dates).

    Returns:
        (level_wide, depth_wide) -- both [date x station_code], node
        ordering matching nodes_df. A node with zero wells within
        max_distance_km gets NaN for every date, not zero -- "no
        nearby monitoring" is meaningfully different from "measured
        zero," same principle as node_features.py's static version.
    """
    from ..data.loaders.ades import ADESLoader
    from scipy.spatial import cKDTree

    loader = ADESLoader(data_path=ades_path)
    gw_df = loader.load()
    start, end = date_range
    full_index = pd.date_range(start, end, freq=freq)

    # Resample each well independently to the common daily grid, forward-fill.
    wells_wide = gw_df.pivot_table(index="date", columns="code_bss", values="groundwater_level_m", aggfunc="mean")
    wells_wide = wells_wide.reindex(wells_wide.index.union(full_index)).sort_index().ffill()
    wells_wide = wells_wide.reindex(full_index)

    depth_wide_raw = gw_df.pivot_table(index="date", columns="code_bss", values="groundwater_depth_m", aggfunc="mean") \
        if "groundwater_depth_m" in gw_df.columns else None
    if depth_wide_raw is not None:
        depth_wide_raw = depth_wide_raw.reindex(depth_wide_raw.index.union(full_index)).sort_index().ffill()
        depth_wide_raw = depth_wide_raw.reindex(full_index)

    well_coords = gw_df.drop_duplicates("code_bss").set_index("code_bss")[["lat", "lon"]]
    well_coords = well_coords.reindex(wells_wide.columns)  # align to the pivoted columns' order
    tree = cKDTree(well_coords[["lon", "lat"]].values)
    coarse_radius_deg = (max_distance_km / 111.0) * 1.5

    def _haversine_km(lat1, lon1, lat2, lon2):
        R = 6371.0
        lat1r, lon1r, lat2r, lon2r = map(np.radians, [lat1, lon1, lat2, lon2])
        dlat, dlon = lat2r - lat1r, lon2r - lon1r
        a = np.sin(dlat / 2) ** 2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlon / 2) ** 2
        return R * 2 * np.arcsin(np.sqrt(a))

    station_lon = nodes_df["longitude"].values
    station_lat = nodes_df["latitude"].values
    candidate_lists = tree.query_ball_point(np.column_stack([station_lon, station_lat]), r=coarse_radius_deg)

    level_cols, depth_cols = {}, {}
    for i, station_code in enumerate(nodes_df["station_code"]):
        candidates = candidate_lists[i]
        if not candidates:
            level_cols[station_code] = pd.Series(np.nan, index=full_index)
            depth_cols[station_code] = pd.Series(np.nan, index=full_index)
            continue
        cand_idx = np.array(candidates)
        cand_lat = well_coords["lat"].values[cand_idx]
        cand_lon = well_coords["lon"].values[cand_idx]
        dist = _haversine_km(station_lat[i], station_lon[i], cand_lat, cand_lon)
        within = cand_idx[dist <= max_distance_km]
        if len(within) == 0:
            level_cols[station_code] = pd.Series(np.nan, index=full_index)
            depth_cols[station_code] = pd.Series(np.nan, index=full_index)
            continue
        well_names = wells_wide.columns[within]
        level_cols[station_code] = wells_wide[well_names].mean(axis=1)
        if depth_wide_raw is not None:
            depth_cols[station_code] = depth_wide_raw[well_names].mean(axis=1)
        else:
            depth_cols[station_code] = pd.Series(np.nan, index=full_index)

    level_wide = pd.DataFrame(level_cols)
    depth_wide = pd.DataFrame(depth_cols)
    return level_wide, depth_wide


def build_climate_grid_timeseries(
    safran_path: Path, bbox: Tuple[float, float, float, float], date_range: Tuple[str, str],
) -> Tuple[Dict[str, pd.DataFrame], pd.DataFrame]:
    """
    Real ERA5 climate at native 0.25 deg grid resolution -- not
    interpolated per station. Confirmed against real data that
    station-level interpolation was redundant: several real stations
    landed on bit-identical values for the same variable, since they
    share the same underlying grid cell. This is the actual fix, not
    a bigger version of the same approach.

    Returns:
        (climate_dict, grid_coords) --
        climate_dict: {variable_name: wide DataFrame [date x grid_cell_id]},
          same shape/contract as build_climate_timeseries's return, so
          every downstream consumer (climatology, later the anomaly
          model) works completely unchanged against either.
        grid_coords: DataFrame [grid_cell_id, grid_lat, grid_lon] --
          needed by map_stations_to_grid_cells (and later, by any river
          node) to know which real coordinate each cell ID refers to.
    """
    from ..data.loaders.safran import SAFRANLoader

    loader = SAFRANLoader(data_path=safran_path)
    df = loader.load_grid(bbox=bbox)
    df = loader.convert_units(df)

    start, end = date_range
    df = df[(df["date"] >= start) & (df["date"] <= end)]

    df["grid_cell_id"] = df["grid_lat"].round(3).astype(str) + "_" + df["grid_lon"].round(3).astype(str)
    grid_coords = df.drop_duplicates("grid_cell_id")[["grid_cell_id", "grid_lat", "grid_lon"]].reset_index(drop=True)

    # convert_units() ADDS converted columns rather than replacing the raw
    # ones -- confirmed against real data that both temp_2m_K and temp_C
    # (and every other raw/converted pair) were being fit separately,
    # doubling work for identical underlying signal (a linear unit shift
    # produces the exact same climatology skill profile either way).
    # wind_u_ms/wind_v_ms are NOT dropped alongside wind_speed_ms -- unlike
    # the other pairs, speed is a genuinely different derived quantity
    # (loses direction information u/v carry), not a unit conversion of
    # either component.
    RAW_COLUMNS_SUPERSEDED_BY_CONVERTED = {
        "temp_2m_K", "precip_m", "evap_m", "solar_Jm2", "snow_m", "runoff_m",
    }
    var_cols = [c for c in df.columns
                if c not in ("date", "grid_lat", "grid_lon", "grid_cell_id")
                and c not in RAW_COLUMNS_SUPERSEDED_BY_CONVERTED]
    climate_dict = {}
    for var in var_cols:
        wide = df.pivot_table(index="date", columns="grid_cell_id", values=var, aggfunc="mean")
        climate_dict[var] = wide
    return climate_dict, grid_coords


def map_stations_to_grid_cells(nodes_df: pd.DataFrame, grid_coords: pd.DataFrame) -> pd.Series:
    """
    Which grid cell each real node reads its climate forecast from --
    computed once (the spatial join is time-invariant, same principle
    already used for groundwater's nearby-well lookup), not repeated
    per timestep or per query. Nearest-cell by exact haversine -- at
    only ~20-30 candidate cells for this study area, no KD-tree
    prefilter is needed the way it was for ADES's ~100 wells against
    ~4,500 nodes; a direct distance computation per node against every
    cell is cheap at this scale.

    Returns:
        Series indexed like nodes_df, values are grid_cell_id strings
        matching build_climate_grid_timeseries's wide DataFrame columns.
    """
    cell_lat = grid_coords["grid_lat"].values
    cell_lon = grid_coords["grid_lon"].values
    cell_ids = grid_coords["grid_cell_id"].values

    assignments = []
    for _, row in nodes_df.iterrows():
        dist = _haversine_km_vec(row["latitude"], row["longitude"], cell_lat, cell_lon)
        assignments.append(cell_ids[dist.argmin()])
    return pd.Series(assignments, index=nodes_df.index, name="grid_cell_id")


def _haversine_km_vec(lat1, lon1, lat2, lon2):
    R = 6371.0
    lat1r, lon1r = np.radians(lat1), np.radians(lon1)
    lat2r, lon2r = np.radians(lat2), np.radians(lon2)
    dlat, dlon = lat2r - lat1r, lon2r - lon1r
    a = np.sin(dlat / 2) ** 2 + np.cos(lat1r) * np.cos(lat2r) * np.sin(dlon / 2) ** 2
    return R * 2 * np.arcsin(np.sqrt(a))


def build_climate_timeseries(
    nodes_df: pd.DataFrame, safran_path: Path, date_range: Tuple[str, str],
) -> Dict[str, pd.DataFrame]:
    """
    Real per-date ERA5 variables, pivoted per variable to [date x
    station_code] wide form -- skips node_features.py's add_safran_
    features aggregation step entirely rather than reversing it.

    SUPERSEDED for new work by build_climate_grid_timeseries +
    map_stations_to_grid_cells: confirmed against real ERA5 data (the
    real climatology test run) that several stations land on
    bit-identical values here, since ERA5's 0.25 deg cells are coarser
    than the spacing between some real gauges -- this function still
    works and is now real-data-tested, but it's doing per-station
    interpolation onto what's provably a shared, coarser grid. Kept for
    any caller that specifically wants a station-indexed table.
    """
    from ..data.loaders.safran import SAFRANLoader

    station_coords = nodes_df.rename(columns={"latitude": "lat", "longitude": "lon"})[
        ["station_code", "lat", "lon"]
    ]
    loader = SAFRANLoader(data_path=safran_path, station_coords=station_coords)
    df = loader.load()
    df = loader.convert_units(df)

    start, end = date_range
    df = df[(df["date"] >= start) & (df["date"] <= end)]

    RAW_COLUMNS_SUPERSEDED_BY_CONVERTED = {
        "temp_2m_K", "precip_m", "evap_m", "solar_Jm2", "snow_m", "runoff_m",
    }
    var_cols = [c for c in df.columns
                if c not in ("date", "station_code") and c not in RAW_COLUMNS_SUPERSEDED_BY_CONVERTED]
    result = {}
    for var in var_cols:
        wide = df.pivot_table(index="date", columns="station_code", values=var, aggfunc="mean")
        result[var] = wide
    return result


def assemble_dynamic_tensor(nodes_df: pd.DataFrame, wide_df: pd.DataFrame) -> Tuple[np.ndarray, List[pd.Timestamp]]:
    """
    Aligns a [date x station_code] wide DataFrame onto nodes_df's own
    row order, so the result matches data.x's node indexing exactly --
    this is the piece that makes a build_*_timeseries output directly
    usable as physics_losses.py's routing_consistency_loss's `Q`
    argument (shape [n_nodes, T]).

    A node whose station_code never appears as a column (every
    confluence/virtual node, for discharge) gets an all-NaN row, not a
    dropped row -- shape stays [n_nodes, T] regardless of coverage.
    """
    aligned = wide_df.reindex(columns=nodes_df["station_code"])
    return aligned.values.T, list(wide_df.index)