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 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_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.

    UNTESTED in this environment: no real era5_*.nc files were
    available here (only ADES and hydrometric real data were). The
    logic mirrors build_discharge_timeseries's pivot pattern directly,
    but please verify the real output shape/values before trusting it
    for training.
    """
    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)]

    var_cols = [c for c in df.columns if c not in ("date", "station_code")]
    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)