Northern-Thailand-Ping-River-Monitor/rust/forecast/src/model.rs

//! Trained bundles exported by src/ml/export.py (`flood_<station>.json.gz`),
//! reloaded when a retrain replaces a file.

use crate::hgb::Head;
use anyhow::{Context, Result, bail};
use serde::Deserialize;
use std::collections::BTreeMap;
use std::io::Read;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::SystemTime;

pub const FORMAT_VERSION: u32 = 1;

#[derive(Debug, Deserialize)]
pub struct BundleThresholds {
    pub warning: f64,
    pub danger: f64,
}

#[derive(Debug, Deserialize)]
pub struct Bundle {
    pub format_version: u32,
    pub station_code: String,
    pub model_version: String,
    pub trained_at: Option<String>,
    pub regression_target: Option<String>,
    pub feature_names: Vec<String>,
    pub horizons: Vec<u32>,
    /// Validation-residual sigma per horizon (keys are the hours as text).
    pub sigma: BTreeMap<String, f64>,
    pub thresholds: BundleThresholds,
    /// `max_<h>` regressors and `warn_<h>` / `danger_<h>` classifiers; a head
    /// skipped at training time is null.
    pub heads: BTreeMap<String, Option<Head>>,
}

impl Bundle {
    pub fn head(&self, name: &str) -> Option<&Head> {
        self.heads.get(name).and_then(Option::as_ref)
    }

    pub fn sigma(&self, horizon: u32, default: f64) -> f64 {
        self.sigma
            .get(&horizon.to_string())
            .copied()
            .unwrap_or(default)
    }

    pub fn from_json(text: &str) -> Result<Self> {
        let bundle: Bundle = serde_json::from_str(text)?;
        if bundle.format_version != FORMAT_VERSION {
            bail!(
                "bundle format {} (this build reads {FORMAT_VERSION}): rebuild ping-collector",
                bundle.format_version
            );
        }
        for (name, head) in &bundle.heads {
            if let Some(h) = head {
                h.validate(bundle.feature_names.len())
                    .with_context(|| format!("head {name}"))?;
            }
        }
        Ok(bundle)
    }

    pub fn load(path: &Path) -> Result<Self> {
        let raw = std::fs::read(path).with_context(|| format!("reading {}", path.display()))?;
        let text = if path.extension().is_some_and(|e| e == "gz") {
            let mut s = String::new();
            flate2::read::GzDecoder::new(&raw[..])
                .read_to_string(&mut s)
                .with_context(|| format!("decompressing {}", path.display()))?;
            s
        } else {
            String::from_utf8(raw)?
        };
        Self::from_json(&text).with_context(|| format!("parsing {}", path.display()))
    }
}

/// The bundles available for one forecast run, by station.
pub type ModelSet = BTreeMap<String, Arc<Bundle>>;

/// `models/flood_<station>.json.gz`, cached by modification time.
pub struct ModelStore {
    dir: PathBuf,
    loaded: BTreeMap<String, (SystemTime, Arc<Bundle>)>,
}

fn station_of(path: &Path) -> Option<String> {
    let name = path.file_name()?.to_str()?;
    let code = name.strip_prefix("flood_")?.strip_suffix(".json.gz")?;
    (!code.is_empty()).then(|| code.to_string())
}

impl ModelStore {
    pub fn new(dir: impl Into<PathBuf>) -> Self {
        Self {
            dir: dir.into(),
            loaded: BTreeMap::new(),
        }
    }

    pub fn dir(&self) -> &Path {
        &self.dir
    }

    /// Rescan the directory: load new or replaced bundles, drop removed ones.
    /// A bundle that fails to load is skipped (that station falls back to the
    /// heuristic, as in Python) and the previous copy, if any, is kept.
    pub fn refresh(&mut self) -> Result<ModelSet> {
        let entries = std::fs::read_dir(&self.dir)
            .with_context(|| format!("reading models directory {}", self.dir.display()))?;
        let mut seen = BTreeMap::new();
        for entry in entries.flatten() {
            let path = entry.path();
            let Some(code) = station_of(&path) else {
                continue;
            };
            let mtime = entry.metadata().and_then(|m| m.modified()).ok();
            seen.insert(code, (path, mtime));
        }
        self.loaded.retain(|code, _| seen.contains_key(code));
        for (code, (path, mtime)) in seen {
            let current = self.loaded.get(&code).map(|(t, _)| *t);
            if mtime.is_some() && current == mtime {
                continue;
            }
            self.warn_if_export_is_stale(&code, &path);
            match Bundle::load(&path) {
                Ok(b) if b.station_code == code => {
                    tracing::info!("model {code}: {} loaded", b.model_version);
                    self.loaded
                        .insert(code, (mtime.unwrap_or(SystemTime::UNIX_EPOCH), Arc::new(b)));
                }
                Ok(b) => tracing::error!(
                    "{}: bundle is for {}, not {code}; ignored",
                    path.display(),
                    b.station_code
                ),
                Err(e) => tracing::error!("model {code}: {e:#}"),
            }
        }
        Ok(self
            .loaded
            .iter()
            .map(|(k, (_, b))| (k.clone(), b.clone()))
            .collect())
    }

    fn warn_if_export_is_stale(&self, code: &str, json: &Path) {
        let joblib = self.dir.join(format!("flood_{code}.joblib"));
        let modified = |p: &Path| std::fs::metadata(p).and_then(|m| m.modified()).ok();
        if let (Some(j), Some(e)) = (modified(&joblib), modified(json))
            && j > e
        {
            tracing::warn!(
                "{} is newer than its export {}: run `python -m src.ml.export`",
                joblib.display(),
                json.display()
            );
        }
    }
}