Northern-Thailand-Ping-River-Monitor/tests/test_flood_forecast.py

"""Tests for the flood forecast ML package. Synthetic data only -- no DB/network."""

import datetime
from typing import Dict, List, Optional

import joblib
import numpy as np
import pandas as pd
import pytest

from src.ml import features, predict, train


def make_synth(
    n_hours: int,
    stations: List[str],
    seed: int = 0,
    start: str = "2020-01-01",
    pulses: Optional[Dict[str, List[tuple]]] = None,
    missing_patches: Optional[Dict[str, List[tuple]]] = None,
) -> pd.DataFrame:
    """Generate a synthetic long measurement frame with smooth levels, flood pulses,
    and optional missing patches, for `stations` over `n_hours` hourly steps.

    pulses: {station: [(start_hour, width_hours, peak_add), ...]}
    missing_patches: {station: [(start_hour, length_hours), ...]}
    """
    rng = np.random.default_rng(seed)
    idx = pd.date_range(start, periods=n_hours, freq="h")
    rows = []
    for station in stations:
        base = 1.5 + 0.1 * np.sin(np.linspace(0, 6 * np.pi, n_hours))
        noise = rng.normal(0, 0.02, n_hours)
        level = base + noise
        for pulse_start, width, peak_add in (pulses or {}).get(station, []):
            t = np.arange(n_hours)
            bump = peak_add * np.exp(-0.5 * ((t - (pulse_start + width / 2)) / (width / 4)) ** 2)
            level = level + bump
        discharge = 20.0 * level + rng.normal(0, 1.0, n_hours)

        missing = np.zeros(n_hours, dtype=bool)
        for patch_start, length in (missing_patches or {}).get(station, []):
            missing[patch_start : patch_start + length] = True

        for i in range(n_hours):
            if missing[i]:
                continue
            rows.append(
                {
                    "timestamp": idx[i],
                    "station_code": station,
                    "water_level": round(float(level[i]), 3),
                    "discharge": round(float(discharge[i]), 2),
                }
            )
    return pd.DataFrame(rows)


def test_no_future_leakage():
    stations = ["P.1", "P.20"]
    df_a = make_synth(200, stations, seed=1, pulses={"P.1": [(150, 10, 3.0)]})
    grid_a = features.make_hourly_grid(df_a)
    feat_a = features.build_features(grid_a, "P.1")

    t0 = grid_a.observed.index[120]

    df_b = df_a.copy()
    future_mask = df_b["timestamp"] > t0
    df_b.loc[future_mask, "water_level"] = df_b.loc[future_mask, "water_level"] + 50.0
    df_b.loc[future_mask, "discharge"] = df_b.loc[future_mask, "discharge"] + 500.0
    grid_b = features.make_hourly_grid(df_b)
    feat_b = features.build_features(grid_b, "P.1")

    past_a = feat_a.loc[feat_a.index <= t0]
    past_b = feat_b.loc[feat_b.index <= t0]
    pd.testing.assert_frame_equal(past_a, past_b)


def test_label_alignment():
    idx = pd.date_range("2020-01-01", periods=12, freq="h")
    warn_thr, _danger_thr = features.get_thresholds("P.1")
    peak = warn_thr + 0.3
    levels = [1.0, 1.0, 1.0, 1.0, 1.0, peak, peak, 1.0, 1.0, 1.0, 1.0, 1.0]
    df = pd.DataFrame(
        {
            "timestamp": idx,
            "station_code": "P.1",
            "water_level": levels,
            "discharge": [20.0 * lvl for lvl in levels],
        }
    )
    grid = features.make_hourly_grid(df)
    labels = features.build_labels(grid, "P.1", horizons=(6,))

    # Level crosses the warning threshold at t=5. A 6h forward window (t, t+6]
    # first includes t=5 for t=0 .. t=4 (inclusive), so exceed_warn_6 should be
    # 1 for t=0..4 and not (necessarily) for later rows in this hand-built series.
    for t in range(5):
        assert labels["exceed_warn_6"].iloc[t] == 1.0, f"t={t} expected warn exceedance"

    # max_level_6 at t=0 covers hours 1..6 -> includes the peak.
    assert labels["max_level_6"].iloc[0] == pytest.approx(peak)


def test_label_coverage_gate():
    n = 40
    idx = pd.date_range("2020-01-01", periods=n, freq="h")
    levels = [1.0] * n
    df = pd.DataFrame(
        {"timestamp": idx, "station_code": "P.1", "water_level": levels, "discharge": [20.0] * n}
    )
    # Drop 70% of a future window (hours 21..26) for the row at t=20, no exceedance in it.
    df_missing = df[~df["timestamp"].isin(idx[21:26])].copy()
    grid = features.make_hourly_grid(df_missing)
    labels = features.build_labels(grid, "P.1", horizons=(6,))
    t20 = idx[20]
    assert pd.isna(labels.loc[t20, "exceed_warn_6"])

    # Same sparse window, but WITH an observed exceedance inside it -> must be 1, not NaN.
    df_with_peak = df_missing.copy()
    peak_row = pd.DataFrame(
        [{"timestamp": idx[22], "station_code": "P.1", "water_level": 5.0, "discharge": 100.0}]
    )
    df_with_peak = pd.concat([df_with_peak, peak_row], ignore_index=True)
    grid2 = features.make_hourly_grid(df_with_peak)
    labels2 = features.build_labels(grid2, "P.1", horizons=(6,))
    assert labels2.loc[t20, "exceed_warn_6"] == 1.0


def test_ffill_and_staleness():
    n = 20
    idx = pd.date_range("2020-01-01", periods=n, freq="h")
    df = pd.DataFrame(
        {
            "timestamp": idx,
            "station_code": "P.1",
            "water_level": [1.0 + 0.01 * i for i in range(n)],
            "discharge": [20.0] * n,
        }
    )
    # Small gap: drop hours 5,6 (2h gap).
    df_small_gap = df[~df["timestamp"].isin(idx[5:7])].copy()
    grid = features.make_hourly_grid(df_small_gap)
    feat = features.build_features(grid, "P.1")
    assert feat.loc[idx[5], "obs_age_h"] == pytest.approx(1.0)
    assert feat.loc[idx[6], "obs_age_h"] == pytest.approx(2.0)

    # Large gap: drop hours 5..9 (5h gap) -> rows with age>3 dropped (NaN).
    df_big_gap = df[~df["timestamp"].isin(idx[5:10])].copy()
    grid2 = features.make_hourly_grid(df_big_gap)
    feat2 = features.build_features(grid2, "P.1")
    assert feat2.loc[idx[8], "obs_age_h"] != feat2.loc[idx[8], "obs_age_h"]  # NaN
    assert feat2.loc[idx[9], "obs_age_h"] != feat2.loc[idx[9], "obs_age_h"]  # NaN
    assert feat2.loc[idx[7], "obs_age_h"] == pytest.approx(3.0)


_FORECAST_KEYS = {
    "station_code",
    "horizon_hours",
    "p_warning",
    "p_danger",
    "predicted_max_level",
    "current_level",
    "as_of",
    "model_version",
    "trained_at",
    "source",
    "threshold_warning",
    "threshold_danger",
}


def _assert_valid_forecast_row(row: dict) -> None:
    # "stages" is optional: model rows for stations in features.FLOOD_STAGES carry
    # per-inundation-stage exceedance probabilities (currently P.1 only).
    assert _FORECAST_KEYS <= set(row.keys())
    assert set(row.keys()) - _FORECAST_KEYS <= {"stages"}
    assert 0.0 <= row["p_warning"] <= 1.0
    assert 0.0 <= row["p_danger"] <= 1.0
    assert row["p_danger"] <= row["p_warning"]
    assert row["predicted_max_level"] >= row["current_level"]
    assert row["source"] in ("model", "heuristic")
    for stage in row.get("stages", []):
        assert 0.0 <= stage["p_exceed"] <= 1.0
        assert stage["level"] > 0


def test_train_smoke_and_roundtrip(tmp_path):
    # Include every station P.1's feature set actually references (its UPSTREAM_LEADS)
    # so no upstream column is entirely NaN -- HistGradientBoosting's binning step
    # cannot fit a fully-degenerate column (see train._safe_fit).
    upstream = [code for code, _lead in features.UPSTREAM_LEADS["P.1"]]
    data_stations = ["P.1"] + upstream
    target_stations = ["P.1", "P.20"]
    n = 700
    pulses = {station: [(start, 20, 2.0) for start in range(50, n - 50, 110)] for station in data_stations}
    df = make_synth(n, data_stations, seed=7, pulses=pulses)

    metrics = train.train_all(
        df, target_stations, models_dir=tmp_path, skip_eval=True, hgb_overrides={"max_iter": 20}
    )
    assert metrics["stations"]["P.1"]["status"] == "trained"
    assert metrics["stations"]["P.20"]["status"] == "trained"
    assert (tmp_path / "flood_P.1.joblib").exists()
    assert (tmp_path / "metrics.json").exists()

    readings_by_station = {
        code: group[["timestamp", "water_level", "discharge"]].to_dict("records")
        for code, group in df.groupby("station_code")
        if code in target_stations
    }
    now = df["timestamp"].max()
    forecasts = predict.get_forecasts(readings_by_station, models_dir=tmp_path, now=now)

    assert len(forecasts) > 0
    for row in forecasts:
        _assert_valid_forecast_row(row)
    assert any(row["source"] == "model" for row in forecasts)


def test_heuristic_fallback(tmp_path):
    df = make_synth(50, ["P.1"], seed=3)
    readings_by_station = {"P.1": df[["timestamp", "water_level", "discharge"]].to_dict("records")}
    now = df["timestamp"].max()

    forecasts = predict.get_forecasts(readings_by_station, models_dir=tmp_path, now=now)

    assert len(forecasts) == len(predict.DEFAULT_HORIZONS)
    for row in forecasts:
        _assert_valid_forecast_row(row)
        assert row["source"] == "heuristic"
        assert row["model_version"] == "heuristic-v1"
        assert row["trained_at"] is None


def test_feature_name_stability(tmp_path):
    upstream = [code for code, _lead in features.UPSTREAM_LEADS["P.1"]]
    data_stations = ["P.1"] + upstream
    df = make_synth(300, data_stations, seed=11, pulses={"P.1": [(100, 20, 2.0)]})
    train.train_all(df, ["P.1"], models_dir=tmp_path, skip_eval=True, hgb_overrides={"max_iter": 10})
    # Safe: loading the bundle this same test just wrote to tmp_path, not an external file.
    bundle = joblib.load(tmp_path / "flood_P.1.joblib")

    grid = features.make_hourly_grid(df)
    fresh_columns = list(features.build_features(grid, "P.1").columns)
    assert fresh_columns == bundle["feature_names"]