Northern-Thailand-Ping-River-Monitor/rust/forecast/tests/parity.rs

//! Replays the fixtures written by testdata/gen_forecast_fixtures.py (the
//! Python production path on the same inputs) through this crate.
//!
//! Exact equality for the merged readings, every feature and every level;
//! probabilities may differ in the last bits of `exp`, and `doy_sin`/`doy_cos`
//! in the last bits of the platform's `sin`/`cos` (numpy vs libm).

use chrono::{NaiveDate, NaiveDateTime};
use ping_forecast::dam::{DamDay, DamFrame, hourly_frame};
use ping_forecast::features::{build_features, row_at};
use ping_forecast::grid::{Grid, Row};
use ping_forecast::hii::{self, HiiRow};
use ping_forecast::model::{Bundle, ModelSet, ModelStore};
use ping_forecast::rain::{RainSeries, catchment_mean};
use ping_forecast::{Forecast, forecasts};
use serde_json::Value;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::Arc;

const SCENARIOS: [&str; 4] = ["flood2024", "peak2024", "recent2026", "stale2026"];
/// Max distance in units in the last place for transcendental results.
const ULPS: u64 = 4;

fn testdata() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata")
}

fn load(name: &str) -> Value {
    let raw = std::fs::read(testdata().join(format!("fixture_{name}.json.gz"))).unwrap();
    let mut s = String::new();
    flate2::read::GzDecoder::new(&raw[..])
        .read_to_string(&mut s)
        .unwrap();
    serde_json::from_str(&s).unwrap()
}

fn ts(v: &Value) -> NaiveDateTime {
    let s = v.as_str().unwrap();
    NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M:%S%.f").unwrap()
}

fn f(v: &Value) -> f64 {
    v.as_f64().unwrap_or(f64::NAN)
}

fn rows(v: &Value) -> Vec<Row> {
    v.as_array()
        .unwrap()
        .iter()
        .map(|r| Row {
            timestamp: ts(&r[0]),
            station_code: r[1].as_str().unwrap().into(),
            water_level: f(&r[2]),
            discharge: f(&r[3]),
        })
        .collect()
}

struct Inputs {
    now: NaiveDateTime,
    rid: Vec<Row>,
    hii: Vec<HiiRow>,
    rain: RainSeries,
    dam: DamFrame,
}

fn inputs(fx: &Value) -> Inputs {
    let hii = fx["hii_rows"]
        .as_array()
        .unwrap()
        .iter()
        .map(|r| HiiRow {
            timestamp: ts(&r[0]),
            station_code: r[1].as_str().unwrap().into(),
            wl_msl: f(&r[2]),
            discharge: f(&r[3]),
        })
        .collect();
    let rain = fx["rain_points"]
        .as_array()
        .unwrap()
        .iter()
        .map(|r| {
            let points: Vec<Option<f64>> =
                r[1].as_array().unwrap().iter().map(Value::as_f64).collect();
            (ts(&r[0]), catchment_mean(&points).unwrap_or(f64::NAN))
        })
        .collect();
    let days: Vec<DamDay> = fx["dam_days"]
        .as_array()
        .unwrap()
        .iter()
        .map(|r| DamDay {
            date: NaiveDate::parse_from_str(r[0].as_str().unwrap(), "%Y-%m-%d").unwrap(),
            storage_pct: f(&r[1]),
            inflow_mcm: f(&r[2]),
            outflow_mcm: f(&r[3]),
        })
        .collect();
    Inputs {
        now: ts(&fx["now"]),
        rid: rows(&fx["rid_rows"]),
        hii,
        rain,
        dam: hourly_frame(&days).unwrap_or_default(),
    }
}

fn same_bits(a: f64, b: f64) -> bool {
    a.to_bits() == b.to_bits() || (a.is_nan() && b.is_nan())
}

fn ulps(a: f64, b: f64) -> u64 {
    if same_bits(a, b) || a == b {
        return 0;
    }
    if a.is_sign_negative() != b.is_sign_negative() || a.is_nan() || b.is_nan() {
        return u64::MAX;
    }
    a.to_bits().abs_diff(b.to_bits())
}

fn merged(inp: &Inputs) -> Vec<Row> {
    let (mut m, _) = hii::fill(&inp.rid, &hii::prepare(&inp.hii));
    m.sort_by(|a, b| (&a.station_code, a.timestamp).cmp(&(&b.station_code, b.timestamp)));
    m
}

#[test]
fn hii_fill_matches_python() {
    for name in SCENARIOS {
        let fx = load(name);
        let got = merged(&inputs(&fx));
        let want = rows(&fx["expected"]["merged_rows"]);
        assert_eq!(got.len(), want.len(), "{name}: merged row count");
        for (g, w) in got.iter().zip(&want) {
            assert!(
                g.timestamp == w.timestamp
                    && g.station_code == w.station_code
                    && same_bits(g.water_level, w.water_level)
                    && same_bits(g.discharge, w.discharge),
                "{name}: {g:?} != {w:?}"
            );
        }
    }
}

#[test]
fn features_match_python_bit_for_bit() {
    let mut checked = 0;
    for name in SCENARIOS {
        let fx = load(name);
        let inp = inputs(&fx);
        let grid = Grid::from_rows(&merged(&inp));
        for (code, exp) in fx["expected"]["features"].as_object().unwrap() {
            let cols = build_features(&grid, code, Some(&inp.rain), Some(&inp.dam));
            let i = grid.index_of(ts(&exp["as_of"])).unwrap();
            assert_eq!(grid.last_observed(code), Some(i), "{name} {code}: as_of");
            let row = row_at(&cols, i);
            let want = exp["features"].as_object().unwrap();
            let mut names: Vec<&str> = row.keys().copied().collect();
            names.sort_unstable();
            let mut want_names: Vec<&str> = want.keys().map(String::as_str).collect();
            want_names.sort_unstable();
            assert_eq!(names, want_names, "{name} {code}: feature set");
            for (k, v) in want {
                let (g, w) = (row[k.as_str()], f(v));
                let ok = if k.starts_with("doy_") {
                    ulps(g, w) <= ULPS
                } else {
                    same_bits(g, w)
                };
                assert!(ok, "{name} {code} {k}: rust {g:?} python {w:?}");
                checked += 1;
            }
        }
    }
    assert!(checked > 3000, "only {checked} feature values compared");
}

fn compare(name: &str, got: &[Forecast], want: &Value) {
    let want = want.as_array().unwrap();
    assert_eq!(got.len(), want.len(), "{name}: row count");
    for (g, w) in got.iter().zip(want) {
        let ctx = format!("{name} {} {}h", g.station_code, g.horizon_hours);
        assert_eq!(g.station_code, w["station_code"].as_str().unwrap(), "{ctx}");
        assert_eq!(
            g.horizon_hours as u64,
            w["horizon_hours"].as_u64().unwrap(),
            "{ctx}"
        );
        assert_eq!(g.as_of, ts(&w["as_of"]), "{ctx}");
        assert_eq!(g.source, w["source"].as_str().unwrap(), "{ctx}");
        assert_eq!(
            g.model_version,
            w["model_version"].as_str().unwrap(),
            "{ctx}"
        );
        assert_eq!(g.trained_at.as_deref(), w["trained_at"].as_str(), "{ctx}");
        for (field, gv) in [
            ("current_level", g.current_level),
            ("predicted_max_level", g.predicted_max_level),
            ("threshold_warning", g.threshold_warning),
            ("threshold_danger", g.threshold_danger),
        ] {
            assert!(
                same_bits(gv, f(&w[field])),
                "{ctx} {field}: {gv} vs {}",
                w[field]
            );
        }
        for (field, gv) in [("p_warning", g.p_warning), ("p_danger", g.p_danger)] {
            let d = ulps(gv, f(&w[field]));
            assert!(d <= ULPS, "{ctx} {field}: {gv} vs {} ({d} ulp)", w[field]);
        }
        match (&g.stages, w.get("stages")) {
            (None, None) => {}
            (Some(gs), Some(ws)) => {
                let ws = ws.as_array().unwrap();
                assert_eq!(gs.len(), ws.len(), "{ctx}: stages");
                for (s, w) in gs.iter().zip(ws) {
                    assert_eq!(s.stage as u64, w["stage"].as_u64().unwrap());
                    assert!(same_bits(s.level, f(&w["level"])));
                    assert!(
                        ulps(s.p_exceed, f(&w["p_exceed"])) <= ULPS,
                        "{ctx} stage {}",
                        s.stage
                    );
                }
            }
            _ => panic!("{ctx}: stages present on one side only"),
        }
    }
}

fn p1_models() -> ModelSet {
    ModelStore::new(testdata().join("models"))
        .refresh()
        .unwrap()
}

#[test]
fn forecasts_match_python_with_the_p1_model() {
    let models = p1_models();
    assert_eq!(models.keys().collect::<Vec<_>>(), ["P.1"]);
    let mut model_rows = 0;
    for name in SCENARIOS {
        let fx = load(name);
        let inp = inputs(&fx);
        let got = forecasts(
            &merged(&inp),
            &models,
            inp.now,
            Some(&inp.rain),
            Some(&inp.dam),
        );
        model_rows += got.iter().filter(|r| r.source == "model").count();
        compare(name, &got, &fx["expected"]["forecasts"]["p1_only"]);
    }
    assert_eq!(model_rows, 9, "P.1 model path in the three fresh scenarios");
}

#[test]
fn moved_thresholds_use_the_regression_path_like_python() {
    let path = testdata().join("models/flood_P.1.json.gz");
    let mut bundle = Bundle::load(&path).unwrap();
    bundle.thresholds.warning = 3.5;
    bundle.thresholds.danger = 4.0;
    let models: ModelSet = [("P.1".to_string(), Arc::new(bundle))].into();
    let fx = load("recent2026");
    let inp = inputs(&fx);
    let got = forecasts(
        &merged(&inp),
        &models,
        inp.now,
        Some(&inp.rain),
        Some(&inp.dam),
    );
    compare(
        "recent2026 moved",
        &got,
        &fx["expected"]["forecasts"]["p1_moved_thresholds"],
    );
}

/// Every local model (needs `python -m src.ml.export` run in ../../models on
/// a machine with the same trained bundles as the fixture generator).
#[test]
fn forecasts_match_python_with_all_local_models() {
    let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models");
    let Ok(models) = ModelStore::new(&dir).refresh() else {
        eprintln!("skipped: no {}", dir.display());
        return;
    };
    let fx = load("recent2026");
    let want: Vec<&str> = fx["expected"]["all_models_versions"]
        .as_array()
        .unwrap()
        .iter()
        .map(|v| v.as_str().unwrap())
        .collect();
    let mut have: Vec<&str> = models.values().map(|b| b.model_version.as_str()).collect();
    have.sort_unstable();
    have.dedup();
    if models.len() < 15 || have != want {
        eprintln!("skipped: local models {have:?} are not the fixture's {want:?}");
        return;
    }
    for name in SCENARIOS {
        let fx = load(name);
        let inp = inputs(&fx);
        let got = forecasts(
            &merged(&inp),
            &models,
            inp.now,
            Some(&inp.rain),
            Some(&inp.dam),
        );
        compare(name, &got, &fx["expected"]["forecasts"]["all_models"]);
    }
}