src/graph/physics_losses.py

"""
Physics-informed loss terms for streamflow prediction on the reach graph.

Four constraints, each tied to real structure the graph now actually
has (see src/graph/build_reach_graph.py):

  1. confluence_mass_balance_loss  -- new mass genuinely enters (is_confluence)
  2. split_rejoin_conservation_loss -- no new mass, paired via braid_id
  3. routing_consistency_loss       -- travel-time lag from distance_km/elevation_drop_m
  4. water_balance_loss             -- P - ET - Q - deltaS ~= 0 per node

All four apply graph-wide, not just at the 27 gauged nodes -- that's the
actual mechanism by which sparse labels generalize to ~4,500 ungauged
nodes, not an incidental detail.

No model exists yet, so these are standalone functions operating on
whatever Q tensor a model eventually produces. They use only operations
that behave identically on a torch.Tensor or a plain numpy array
(indexing, elementwise arithmetic, sum, mean) so the same code path is
testable now with numpy and will work unchanged with real torch tensors
and real gradients once a model exists -- verified by testing this
module against numpy inputs directly.
"""
from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd

try:
    import torch
    _HAS_TORCH = True
except ImportError:
    _HAS_TORCH = False


def _mse(residual):
    """
    Mean squared residual, NaN-masked -- works identically on numpy or
    torch. Real ground-truth Q is naturally, heavily NaN (only real
    gauges with real observations ever have a value; confirmed against
    real data: ~93.5% NaN for the reach graph's discharge tensor) --
    without masking, ANY single NaN anywhere in the residual poisons
    the entire mean to NaN, which isn't a rare edge case for this data,
    it's the normal shape of it. Returns NaN only if truly nothing
    usable exists (every entry NaN), which is a real "no data" signal
    worth surfacing, not silently averaging to 0 and implying perfect
    physics satisfaction when there was actually no evidence either way.
    """
    if _HAS_TORCH and isinstance(residual, torch.Tensor):
        mask = ~torch.isnan(residual)
        if not mask.any():
            return residual.sum() * float("nan")
        return (residual[mask] ** 2).mean()
    mask = ~np.isnan(residual)
    if not mask.any():
        return np.nan
    return (residual[mask] ** 2).mean()


# ---------------------------------------------------------------------------
# 1. Confluence mass balance: Q_confluence ~= sum(Q_upstream_branches)
# ---------------------------------------------------------------------------

def build_confluence_index(nodes_df: pd.DataFrame, edges_df: pd.DataFrame) -> List[Tuple[int, List[int]]]:
    """
    Precompute, once per graph (not per training step), which node
    indices feed into each real confluence. Returns
    [(confluence_idx, [upstream_idx, ...]), ...] using positional
    indices into nodes_df (0..n-1), matching how a model's output
    tensor would be laid out.

    Splitting this out from the loss function itself means the loss can
    just do array indexing every step -- the graph structure doesn't
    change between training steps, so there's no reason to recompute
    which nodes are involved in each confluence on every call.
    """
    code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])}
    confluence_codes = set(nodes_df[nodes_df["is_confluence"]]["station_code"])

    pairs = []
    for conf_code in confluence_codes:
        upstream = edges_df[edges_df["target"] == conf_code]["source"].tolist()
        upstream_idx = [code_to_idx[u] for u in upstream if u in code_to_idx]
        if len(upstream_idx) >= 2:
            pairs.append((code_to_idx[conf_code], upstream_idx))
    return pairs


def confluence_mass_balance_loss(Q, confluence_index: List[Tuple[int, List[int]]]):
    """
    For each real confluence, predicted discharge there should
    approximately equal the sum of its upstream branches' predicted
    discharge -- new mass genuinely enters at a confluence (an
    independent tributary catchment), so this is a straightforward sum,
    unlike the split/rejoin case below.

    Ignores travel time between the branches and the confluence (an
    instantaneous-mass approximation) -- see routing_consistency_loss
    for the piece that accounts for lag separately.

    Args:
        Q: predicted discharge, shape [n_nodes] (one timestep) or
            [n_nodes, T] (multiple timesteps, this loss applies per
            timestep the same way).
        confluence_index: from build_confluence_index.

    Returns:
        Scalar loss (0.0, on the same array type as Q, if no confluences).
    """
    if not confluence_index:
        return Q.sum() * 0.0  # zero, but keeps dtype/type consistent (torch-safe)

    residuals = []
    for conf_idx, upstream_idx in confluence_index:
        upstream_sum = Q[upstream_idx[0]]
        for idx in upstream_idx[1:]:
            upstream_sum = upstream_sum + Q[idx]
        residuals.append(Q[conf_idx] - upstream_sum)

    if _HAS_TORCH and isinstance(Q, torch.Tensor):
        residual_stack = torch.stack(residuals)
    else:
        residual_stack = np.stack(residuals)
    return _mse(residual_stack)


# ---------------------------------------------------------------------------
# 2. Split/rejoin conservation: Q_split ~= Q_rejoin (no new mass between them)
# ---------------------------------------------------------------------------

def build_braid_index(nodes_df: pd.DataFrame) -> List[Tuple[int, int]]:
    """
    Precompute (split_idx, rejoin_idx) pairs from the saved braid_id
    column (see build_reach_graph.py's pair_splits_and_rejoins). Same
    precompute-once rationale as build_confluence_index.
    """
    code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])}
    pairs = []
    for _, row in nodes_df[nodes_df["braid_id"].notna()].iterrows():
        rejoin_code, split_code = row["station_code"], row["braid_id"]
        if rejoin_code in code_to_idx and split_code in code_to_idx:
            pairs.append((code_to_idx[split_code], code_to_idx[rejoin_code]))
    return pairs


def split_rejoin_conservation_loss(Q, braid_index: List[Tuple[int, int]]):
    """
    For each matched split/rejoin pair, predicted discharge should be
    approximately equal at both ends -- the same water dividing into
    parallel channels and recombining adds no new mass, unlike a real
    confluence (see confluence_mass_balance_loss). This is a genuinely
    different physical constraint, not a weaker version of the same one:
    a model that learned "sum inflows" generically would get this wrong,
    since a rejoin's two branches together should equal the SPLIT's
    single value, not add something new on top of it.

    Args:
        Q: predicted discharge, shape [n_nodes] or [n_nodes, T].
        braid_index: from build_braid_index.

    Returns:
        Scalar loss (0.0 if no braids in this graph).
    """
    if not braid_index:
        return Q.sum() * 0.0

    split_idx = [s for s, _ in braid_index]
    rejoin_idx = [r for _, r in braid_index]
    residual = Q[split_idx] - Q[rejoin_idx]
    return _mse(residual)


# ---------------------------------------------------------------------------
# 3. Routing: travel-time lag from real channel distance and slope
# ---------------------------------------------------------------------------

def estimate_travel_time_hours(
    distance_km,
    elevation_drop_m,
    min_velocity_ms: float = 0.1,
    max_velocity_ms: float = 3.0,
    velocity_coefficient: float = 1.0,
) -> float:
    """
    Rough channel-flow velocity from slope, in the spirit of Manning's
    equation's slope dependence (v ~ sqrt(slope)) without the channel
    geometry/roughness terms Manning's actually needs, which we don't
    have real data for -- explicitly an approximation, not a full
    hydraulic solve. Slope = elevation_drop_m / (distance_km * 1000).

    Clamped to [min_velocity_ms, max_velocity_ms] since a near-zero or
    negative slope (a virtually flat reach, or a data artifact) would
    otherwise give a nonsensical near-infinite or negative travel time.

    Returns:
        Travel time in hours for water to traverse this edge.
    """
    distance_m = distance_km * 1000.0
    slope = np.clip(elevation_drop_m / np.maximum(distance_m, 1.0), 1e-6, None)
    velocity = np.clip(velocity_coefficient * np.sqrt(slope) * 10.0, min_velocity_ms, max_velocity_ms)
    return distance_m / velocity / 3600.0


def build_routing_index(
    nodes_df: pd.DataFrame, edges_df: pd.DataFrame, timestep_hours: float = 24.0,
) -> List[Tuple[int, int, int]]:
    """
    Precompute (upstream_idx, downstream_idx, lag_timesteps) for every
    edge, rounding each edge's estimated travel time to the nearest
    whole timestep -- e.g. a 30-hour travel time at a 24-hour (daily)
    timestep rounds to a 1-step lag. An edge whose travel time rounds to
    0 is still included (same-timestep routing, lag=0).
    """
    code_to_idx = {code: i for i, code in enumerate(nodes_df["station_code"])}
    pairs = []
    for _, e in edges_df.iterrows():
        if e["source"] not in code_to_idx or e["target"] not in code_to_idx:
            continue
        drop = e["elevation_drop_m"] if pd.notna(e["elevation_drop_m"]) else 0.1
        hours = estimate_travel_time_hours(e["distance_km"], max(drop, 0.1))
        lag = int(round(hours / timestep_hours))
        pairs.append((code_to_idx[e["source"]], code_to_idx[e["target"]], lag))
    return pairs


def routing_consistency_loss(Q, routing_index: List[Tuple[int, int, int]]):
    """
    Q at a downstream node at time t should approximately equal Q at its
    upstream node at time (t - lag), lag coming from real distance and
    slope (build_routing_index) -- not just "conserve mass at the same
    instant," which routing_consistency_loss's siblings above assume as
    a simplification. This is the piece that makes that simplification
    less necessary over time: a well-trained model satisfying this loss
    is learning the actual travel-time behavior of each reach.

    Args:
        Q: predicted discharge, shape [n_nodes, T] -- REQUIRES a time
            dimension, unlike the other three losses, since travel-time
            lag is meaningless for a single instant.
        routing_index: from build_routing_index.

    Returns:
        Scalar loss (0.0 if no edges have a usable lag within Q's time range).
    """
    T = Q.shape[1]
    residuals = []
    for up_idx, down_idx, lag in routing_index:
        if lag >= T:
            continue  # this edge's travel time exceeds the whole prediction window
        if lag == 0:
            residuals.append(Q[down_idx, :] - Q[up_idx, :])
        else:
            residuals.append(Q[down_idx, lag:] - Q[up_idx, :-lag])

    if not residuals:
        return Q.sum() * 0.0

    if _HAS_TORCH and isinstance(Q, torch.Tensor):
        residual_cat = torch.cat(residuals)
    else:
        residual_cat = np.concatenate(residuals)
    return _mse(residual_cat)


# ---------------------------------------------------------------------------
# 4. Water balance: P - ET - Q - deltaS ~= 0, per node
# ---------------------------------------------------------------------------

def water_balance_loss(
    Q_m3s,
    precip_mm,
    evap_mm,
    catchment_area_km2,
    period_days: float = 365.0,
    delta_storage_m3: Optional[object] = None,
):
    """
    Precipitation minus evapotranspiration minus discharge minus storage
    change should balance to ~0, in volume terms, over the given period.

    UNIT CONVERSION (the easy part to get subtly wrong): 1 mm of depth
    over 1 km^2 is 1000 m^3 (1 km^2 = 1e6 m^2, 1 mm = 1e-3 m,
    1e6 * 1e-3 = 1e3). P and ET (mm, over the period) get converted to
    m^3 via catchment_area_km2 before comparing against Q, which is
    converted from a rate (m^3/s) to a volume by multiplying by the
    period length in seconds.

    delta_storage_m3 defaults to zero (a steady-state approximation) --
    we have no direct storage measurement (soil moisture, groundwater
    volume change) in this project's data, only groundwater LEVEL at
    sparse wells, which isn't the same thing as a basin-wide storage
    volume. Treating deltaS as strictly zero is a real, named
    approximation, not a hidden one -- pass a nonzero delta_storage_m3
    if a proxy for it becomes available later (e.g. derived from
    groundwater level trend where well coverage allows it).

    Args:
        Q_m3s: predicted discharge, shape [n_nodes] (period-average rate).
        precip_mm, evap_mm: node features, already available.
        catchment_area_km2: from catchment.py -- NaN for ungauged/unknown
            catchments, in which case that node is excluded from this
            loss entirely (silently including it with a wrong/zero area
            would corrupt the term, not just add noise).
        period_days: length of the period P/ET/Q are aggregated over.
        delta_storage_m3: optional storage change; zero-array default.

    Returns:
        Scalar loss, computed only over nodes with a real catchment area.
    """
    valid = ~np.isnan(catchment_area_km2) if not _HAS_TORCH or not isinstance(catchment_area_km2, torch.Tensor) \
        else ~torch.isnan(catchment_area_km2)

    period_seconds = period_days * 86400.0
    Q_volume_m3 = Q_m3s * period_seconds
    P_volume_m3 = precip_mm * catchment_area_km2 * 1000.0
    ET_volume_m3 = evap_mm * catchment_area_km2 * 1000.0
    dS = delta_storage_m3 if delta_storage_m3 is not None else (Q_m3s * 0.0)

    residual = P_volume_m3 - ET_volume_m3 - Q_volume_m3 - dS
    residual_valid = residual[valid]
    if (residual_valid.shape[0] if hasattr(residual_valid, "shape") else len(residual_valid)) == 0:
        return Q_m3s.sum() * 0.0
    return _mse(residual_valid)


# ---------------------------------------------------------------------------
# Combined loss
# ---------------------------------------------------------------------------

def physics_informed_loss(
    Q_supervised_pred, Q_supervised_true, gauged_mask,
    Q_full, confluence_index, braid_index,
    weights: Optional[Dict[str, float]] = None,
    Q_timeseries=None, routing_index=None,
    precip_mm=None, evap_mm=None, catchment_area_km2=None,
) -> Dict[str, float]:
    """
    Combines the supervised loss (masked to gauged nodes) with all
    physics terms that have the inputs to compute (routing and water
    balance are optional -- they need a time dimension / climate data
    respectively, which not every training step may have on hand).

    Returns a dict of every individual term plus 'total', rather than
    just the summed scalar -- so it's possible to see which physics
    term is actually driving the loss during training, not just that
    "the loss" went up or down.
    """
    weights = weights or {"confluence": 1.0, "split_rejoin": 1.0, "routing": 1.0, "water_balance": 1.0}

    supervised_residual = (Q_supervised_pred - Q_supervised_true)[gauged_mask]
    losses = {"supervised": _mse(supervised_residual)}

    losses["confluence"] = confluence_mass_balance_loss(Q_full, confluence_index)
    losses["split_rejoin"] = split_rejoin_conservation_loss(Q_full, braid_index)

    if Q_timeseries is not None and routing_index is not None:
        losses["routing"] = routing_consistency_loss(Q_timeseries, routing_index)
    if precip_mm is not None and evap_mm is not None and catchment_area_km2 is not None:
        losses["water_balance"] = water_balance_loss(Q_full, precip_mm, evap_mm, catchment_area_km2)

    total = losses["supervised"]
    for name, w_key in [("confluence", "confluence"), ("split_rejoin", "split_rejoin"),
                          ("routing", "routing"), ("water_balance", "water_balance")]:
        if name in losses:
            total = total + weights.get(w_key, 1.0) * losses[name]
    losses["total"] = total
    return losses