prop/rust/prop_grid_rs/src/weather_scalar_file.rs

//! Per-cell *derived* weather scalars on disk, the cheap-read sibling
//! of [`pgrid`][crate::pgrid].
//!
//! 1:1 wire-compatible with `Microwaveprop.Weather.ScalarFile` on the
//! Elixir side. Chunked into 5°×5° spatial buckets so a state-sized
//! `/weather` viewport reads only the chunks that overlap.
//!
//! ## Layout
//!
//! ```text
//! <scores_dir>/weather_scalars/
//!   <iso>/                           # e.g. 2026-04-29T12:00:00Z/
//!     <lat_band>_<lon_band>.mp.gz
//! ```
//!
//! `lat_band = floor(lat / 5)`, `lon_band = floor(lon / 5)`. Each chunk
//! is gzipped MessagePack: `[ {row}, {row}, … ]`. Map keys are strings
//! on the wire — the Elixir reader atomizes the whitelist on read, so
//! every key emitted here must match the Elixir whitelist or it gets
//! silently dropped by callers that use atom-style access.
//!
//! Scalar derivation mirrors `Microwaveprop.Weather.WeatherLayers.derive/1`
//! plus the surface fields `Microwaveprop.Weather.build_grid_cache_row/4`
//! adds on top. The Elixir module remains the single source of truth
//! for any cell that hasn't been Rust-materialized yet (the
//! `kickoff_async_scalar_materialize` fallback path).

use std::collections::HashMap;
use std::fs::File;
use std::io::{BufWriter, Write};
use std::path::{Path, PathBuf};

use chrono::{DateTime, Utc};
use flate2::write::GzEncoder;
use flate2::Compression;
use serde::Serialize;

use crate::field_grid::FieldGrid;
use crate::planes::GridPlanes;
use crate::sounding_params::{
    ducting_detected, min_refractivity_gradient, sat_vap_pres, surface_refractivity, Level,
};

const CHUNK_STEP: i32 = 5;
const SUBDIR: &str = "weather_scalars";

/// One row in a chunk file. Field names match the Elixir atom-key
/// whitelist exactly — adding a field here without updating the
/// Elixir `@atom_keys` set means the new field round-trips as a
/// string and is invisible to atom-style callers.
#[derive(Debug, Default, Serialize)]
pub struct ScalarRow {
    pub lat: f64,
    pub lon: f64,
    pub valid_time: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dewpoint_depression: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub surface_rh: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub surface_pressure_mb: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub surface_refractivity: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refractivity_gradient: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bl_height: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pwat: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temp_850mb: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dewpoint_850mb: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temp_700mb: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub dewpoint_700mb: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub lapse_rate: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub mid_lapse_rate: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub inversion_strength: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub inversion_base_m: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub ducting: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duct_base_m: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duct_strength: Option<f64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub duct_cutoff_ghz: Option<f64>,
}

#[derive(Debug, thiserror::Error)]
pub enum WriteError {
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
    #[error("encode: {0}")]
    Encode(#[from] rmp_serde::encode::Error),
}

/// Build a `ScalarRow` from a single grid cell. Returns `None` for
/// cells whose surface temperature is missing or out of physically
/// plausible range — matches Elixir's `build_grid_cache_row` filter
/// (`temp_c < -80 or temp_c > 60` are dropped).
///
/// `levels` is supplied by the caller (via `GridPlanes::levels_at`) so
/// the pressure-level walk is shared with the conditions and profile
/// derivation instead of being redone here.
pub fn derive_row(
    grid: &FieldGrid,
    p: &GridPlanes,
    cell: usize,
    lat: f64,
    lon: f64,
    valid_time: DateTime<Utc>,
    levels: &[Level],
) -> Option<ScalarRow> {
    let temp_c = grid.at_opt(p.tmp_2m, cell).map(|v| v as f64 - 273.15)?;
    if !temp_c.is_finite() || !(-80.0..=60.0).contains(&temp_c) {
        return None;
    }

    let dewpoint_c = grid.at_opt(p.dpt_2m, cell).map(|v| v as f64 - 273.15);
    let dewpoint_depression = dewpoint_c.map(|d| temp_c - d);
    let surface_pressure_mb = grid.at_opt(p.pres_sfc, cell).map(|v| v as f64 / 100.0);
    let bl_height = grid.at_opt(p.hpbl, cell).map(|v| v as f64);
    let pwat = grid.at_opt(p.pwat, cell).map(|v| v as f64);

    let surface_rh = dewpoint_c.map(|d| 100.0 * sat_vap_pres(d) / sat_vap_pres(temp_c));

    let derived_min_grad = if levels.len() >= 3 {
        min_refractivity_gradient(levels.to_vec())
    } else {
        None
    };
    let native_min_grad = grid.at_opt(p.native_min_gradient, cell).map(|v| v as f64);
    let refractivity_gradient = derived_min_grad.or(native_min_grad);
    let ducting = Some(ducting_detected(refractivity_gradient));
    let surface_refractivity_val = surface_refractivity(levels);

    let mut sorted: Vec<&Level> = levels.iter().collect();

    sorted.sort_by(|a, b| {
        b.pres_mb
            .partial_cmp(&a.pres_mb)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    let temp_850mb = level_value(&sorted, 850.0, |l| Some(l.tmpc));
    let dewpoint_850mb = level_value(&sorted, 850.0, |l| l.dwpc);
    let temp_700mb = level_value(&sorted, 700.0, |l| Some(l.tmpc));
    let dewpoint_700mb = level_value(&sorted, 700.0, |l| l.dwpc);

    let lapse_rate = compute_lapse_rate(&sorted);
    let mid_lapse_rate = compute_layer_lapse_rate(&sorted, 850.0, 700.0);
    let inversion_strength = compute_inversion_strength(&sorted);
    let inversion_base_m = compute_inversion_base_m(&sorted);

    // Ducts: the f01..f48 pipeline only stores the duct-summary scalars
    // (`max_duct_thickness_m`, `duct_count`, `best_duct_freq_ghz`) per
    // cell. Duct base requires the full Duct list, which isn't threaded
    // through `CellValues` yet — matches the current Elixir behavior on
    // Rust-produced profiles (their `:duct_characteristics` is also
    // nil). `duct_cutoff_ghz` IS available because
    // `native_duct::reduce_grid_to_ducts` reduces it to a scalar at
    // ingest time.
    let duct_count = grid.at_opt(p.duct_count, cell).unwrap_or(0.0);
    let duct_strength = if duct_count > 0.0 {
        grid.at_opt(p.max_duct_thickness_m, cell).map(|v| v as f64)
    } else {
        None
    };
    let duct_cutoff_ghz = if duct_count > 0.0 {
        grid.at_opt(p.best_duct_freq_ghz, cell).map(|v| v as f64)
    } else {
        None
    };

    Some(ScalarRow {
        lat,
        lon,
        valid_time: valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string(),
        temperature: Some(temp_c),
        dewpoint_depression,
        surface_rh,
        surface_pressure_mb,
        surface_refractivity: surface_refractivity_val,
        refractivity_gradient,
        bl_height,
        pwat,
        temp_850mb,
        dewpoint_850mb,
        temp_700mb,
        dewpoint_700mb,
        lapse_rate,
        mid_lapse_rate,
        inversion_strength,
        inversion_base_m,
        ducting,
        duct_base_m: None,
        duct_strength,
        duct_cutoff_ghz,
    })
}

/// Absolute path the HRRR writer lands at for `valid_time`.
pub fn dir_for(scores_dir: &Path, valid_time: DateTime<Utc>) -> PathBuf {
    let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    scores_dir.join(SUBDIR).join(iso)
}

/// Sibling directory for HRDPS-derived scalar chunks. Coexists with the
/// HRRR `dir_for` directory so the two sources can write 5°×5° chunks
/// independently — `write_atomic`'s "wipe dir first" sweep would
/// otherwise have whichever source ran last clobber the other's chunks.
/// Readers stitch the two sets together at request time.
pub fn dir_for_hrdps(scores_dir: &Path, valid_time: DateTime<Utc>) -> PathBuf {
    let iso = valid_time.format("%Y-%m-%dT%H:%M:%SZ").to_string();
    scores_dir.join(SUBDIR).join(format!("{iso}.hrdps"))
}

/// Persist `rows` for `valid_time` as 5°×5° chunked `.mp.gz` files.
/// Mirrors the Elixir writer: the destination directory is wiped of
/// any prior chunks first so a smaller follow-up write doesn't leave
/// stale files behind. Each chunk is written `tmp + rename` for an
/// atomic NFS-friendly publish.
pub fn write_atomic(
    scores_dir: &Path,
    valid_time: DateTime<Utc>,
    rows: &[ScalarRow],
) -> Result<PathBuf, WriteError> {
    write_atomic_into(dir_for(scores_dir, valid_time), rows)
}

/// HRDPS sibling of `write_atomic` that lands at `dir_for_hrdps`. Wipes
/// only its own dir, leaving the HRRR `<vt>/` untouched.
pub fn write_atomic_hrdps(
    scores_dir: &Path,
    valid_time: DateTime<Utc>,
    rows: &[ScalarRow],
) -> Result<PathBuf, WriteError> {
    write_atomic_into(dir_for_hrdps(scores_dir, valid_time), rows)
}

fn write_atomic_into(dir: PathBuf, rows: &[ScalarRow]) -> Result<PathBuf, WriteError> {
    std::fs::create_dir_all(&dir)?;

    // Only sweep known scalar chunk files (*.mp.gz), not everything in
    // the directory. A crash between this sweep and the chunk writes
    // below would otherwise permanently delete all previously-persisted
    // scalars for this valid_time.
    if let Ok(entries) = std::fs::read_dir(&dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path
                .extension()
                .and_then(|e| e.to_str())
                .is_some_and(|ext| ext == "gz")
                && path
                    .file_name()
                    .and_then(|n| n.to_str())
                    .is_some_and(|name| name.ends_with(".mp.gz"))
            {
                let _ = std::fs::remove_file(path);
            }
        }
    }

    let mut chunks: HashMap<(i32, i32), Vec<&ScalarRow>> = HashMap::new();
    for row in rows {
        let key = (chunk_band(row.lat), chunk_band(row.lon));
        chunks.entry(key).or_default().push(row);
    }

    for ((lat_band, lon_band), chunk_rows) in chunks {
        let path = dir.join(format!("{lat_band}_{lon_band}.mp.gz"));

        let nanos = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .map(|d| d.as_nanos())
            .unwrap_or(0);
        let pid = std::process::id();
        let mut tmp = path.clone().into_os_string();
        tmp.push(format!(".tmp.{nanos}.{pid}"));
        let tmp = PathBuf::from(tmp);

        {
            let buf = rmp_serde::to_vec_named(&chunk_rows)?;
            let file = File::create(&tmp)?;
            let mut gz = GzEncoder::new(BufWriter::new(file), Compression::default());
            gz.write_all(&buf)?;
            gz.finish()?;
        }

        std::fs::rename(&tmp, &path)?;
    }

    Ok(dir)
}

// ── Helpers ──────────────────────────────────────────────────────────

fn chunk_band(value: f64) -> i32 {
    (value / CHUNK_STEP as f64).floor() as i32
}

/// Nearest level to `target_pres` within ±25 mb, mirroring Elixir's
/// `WeatherLayers.level_value/3`.
fn level_value(
    sorted: &[&Level],
    target_pres: f64,
    project: impl Fn(&Level) -> Option<f64>,
) -> Option<f64> {
    let nearest = sorted.iter().min_by(|a, b| {
        (a.pres_mb - target_pres)
            .abs()
            .partial_cmp(&(b.pres_mb - target_pres).abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    })?;
    if (nearest.pres_mb - target_pres).abs() <= 25.0 {
        project(nearest)
    } else {
        None
    }
}

fn compute_lapse_rate(sorted: &[&Level]) -> Option<f64> {
    if sorted.len() < 2 {
        return None;
    }
    let surface = sorted.first()?;
    let top = sorted.last()?;
    let dh_km = (top.hght_m - surface.hght_m) / 1000.0;
    if dh_km > 0.0 {
        Some((surface.tmpc - top.tmpc) / dh_km)
    } else {
        None
    }
}

fn compute_layer_lapse_rate(sorted: &[&Level], lower_pres: f64, upper_pres: f64) -> Option<f64> {
    if lower_pres <= upper_pres {
        return None;
    }
    let lower = sorted
        .iter()
        .find(|l| (l.pres_mb - lower_pres).abs() <= 25.0)?;
    let upper = sorted
        .iter()
        .find(|l| (l.pres_mb - upper_pres).abs() <= 25.0)?;
    let dh_km = (upper.hght_m - lower.hght_m) / 1000.0;
    if dh_km > 0.0 {
        Some((lower.tmpc - upper.tmpc) / dh_km)
    } else {
        None
    }
}

fn compute_inversion_strength(sorted: &[&Level]) -> Option<f64> {
    if sorted.is_empty() {
        return Some(0.0);
    }
    let mut max_str = 0.0_f64;
    for pair in sorted.windows(2) {
        let dt = pair[1].tmpc - pair[0].tmpc;
        if dt > max_str {
            max_str = dt;
        }
    }
    Some(max_str)
}

fn compute_inversion_base_m(sorted: &[&Level]) -> Option<f64> {
    if sorted.is_empty() {
        return None;
    }
    let sfc_hght = sorted.first()?.hght_m;
    let mut max_str = 0.0_f64;
    let mut base: Option<f64> = None;
    for pair in sorted.windows(2) {
        let dt = pair[1].tmpc - pair[0].tmpc;
        if dt > 0.0 && dt > max_str {
            max_str = dt;
            base = Some(pair[0].hght_m - sfc_hght);
        }
    }
    if max_str > 0.0 {
        base
    } else {
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::fetcher;
    use crate::grid::GridSpec;
    use chrono::TimeZone;
    use flate2::read::GzDecoder;
    use std::io::Read;

    fn one_cell_spec() -> GridSpec {
        GridSpec {
            lon_start: -97.0,
            lon_count: 1,
            lon_step: 0.125,
            lat_start: 33.0,
            lat_count: 1,
            lat_step: 0.125,
        }
    }

    fn grid_with(items: &[(&str, f32)]) -> FieldGrid {
        let mut g = FieldGrid::new(one_cell_spec());
        for &(k, v) in items {
            g.push_plane(k, vec![v]);
        }
        g
    }

    fn surface_only_grid() -> FieldGrid {
        grid_with(&[
            ("TMP:2 m above ground", 295.65), // 22.5 °C
            ("DPT:2 m above ground", 285.65), // 12.5 °C → depression 10 °C
            ("PRES:surface", 101_320.0),      // 1013.2 mb
            ("HPBL:surface", 800.0),
            (
                "PWAT:entire atmosphere (considered as a single layer)",
                25.0,
            ),
        ])
    }

    /// Resolve planes, pull the cell's levels, and derive — the same
    /// sequence the pipeline runs, collapsed for test ergonomics.
    fn derive(grid: &FieldGrid, vt: DateTime<Utc>) -> Option<ScalarRow> {
        let p = GridPlanes::resolve(grid);
        let mut levels = Vec::new();
        p.levels_at(grid, 0, &mut levels);
        derive_row(grid, &p, 0, 33.0, -97.0, vt, &levels)
    }

    #[test]
    fn drops_cells_with_unphysical_surface_temp() {
        let mut grid = surface_only_grid();
        grid.push_plane("TMP:2 m above ground", vec![100.0]); // way too cold
        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();
        assert!(derive(&grid, vt).is_none());
    }

    #[test]
    fn surface_only_row_carries_basic_fields() {
        let grid = surface_only_grid();
        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();
        let row = derive(&grid, vt).expect("row");

        assert_eq!(row.lat, 33.0);
        assert_eq!(row.lon, -97.0);
        assert_eq!(row.valid_time, "2026-04-29T12:00:00Z");
        assert!((row.temperature.unwrap() - 22.5).abs() < 1e-3);
        assert!((row.dewpoint_depression.unwrap() - 10.0).abs() < 1e-3);
        assert!((row.surface_pressure_mb.unwrap() - 1013.2).abs() < 1e-3);
        assert!(row.surface_rh.unwrap() > 30.0 && row.surface_rh.unwrap() < 70.0);
        assert_eq!(row.bl_height, Some(800.0));
        assert_eq!(row.pwat, Some(25.0));

        // No pressure-level data → upper-air fields stay nil.
        assert_eq!(row.temp_850mb, None);
        assert_eq!(row.lapse_rate, None);
    }

    #[test]
    fn upper_air_levels_derive_lapse_rate_and_850mb() {
        // Need real fetcher::GRID_PRESSURE_LEVELS keys. Build directly
        // from grid_level_keys() so the synthetic cell mirrors prod.
        let mut grid = surface_only_grid();
        for k in fetcher::grid_level_keys() {
            // Plausible synthetic profile: lapse 6.5 °C/km, scale height
            // ≈ 8 km. Heights derived from std atmosphere approximation.
            let h_m = match k.pres_mb as i32 {
                1000 => 100.0,
                925 => 800.0,
                850 => 1500.0,
                700 => 3000.0,
                500 => 5600.0,
                250 => 10_400.0,
                _ => 1500.0 + (1000.0 - k.pres_mb) * 8.0,
            };
            let t_c = 22.5 - h_m * 6.5e-3; // °C
            let d_c = t_c - 5.0;
            grid.push_plane(&k.tmp, vec![(t_c + 273.15) as f32]);
            grid.push_plane(&k.dpt, vec![(d_c + 273.15) as f32]);
            grid.push_plane(&k.hgt, vec![h_m as f32]);
        }

        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();
        let row = derive(&grid, vt).expect("row");

        let temp_850 = row.temp_850mb.expect("850 mb T");
        // h_m at 850 mb is 1500 → t_c = 22.5 - 1500 * 6.5e-3 = 12.75 °C
        assert!(
            (temp_850 - 12.75).abs() < 0.5,
            "expected ~12.75 °C at 850 mb, got {temp_850}"
        );

        let lapse = row.lapse_rate.expect("lapse_rate");
        assert!(
            (lapse - 6.5).abs() < 0.5,
            "expected ~6.5 °C/km lapse rate, got {lapse}"
        );

        let mid = row.mid_lapse_rate.expect("mid_lapse_rate");
        assert!(
            (mid - 6.5).abs() < 0.5,
            "expected ~6.5 mid lapse, got {mid}"
        );
    }

    #[test]
    fn write_atomic_round_trip_via_msgpack() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();

        let row = ScalarRow {
            lat: 33.0,
            lon: -97.0,
            valid_time: "2026-04-29T12:00:00Z".to_string(),
            temperature: Some(22.5),
            dewpoint_depression: Some(10.0),
            surface_rh: Some(50.0),
            ducting: Some(false),
            ..ScalarRow::default()
        };

        write_atomic(dir.path(), vt, &[row]).unwrap();

        let chunk = dir
            .path()
            .join(SUBDIR)
            .join("2026-04-29T12:00:00Z")
            .join("6_-20.mp.gz");
        assert!(chunk.exists());

        let raw = std::fs::read(&chunk).unwrap();
        let mut gz = GzDecoder::new(&raw[..]);
        let mut buf = Vec::new();
        gz.read_to_end(&mut buf).unwrap();
        let decoded: rmpv::Value = rmp_serde::from_slice(&buf).unwrap();
        let arr = decoded.as_array().expect("top-level array");
        assert_eq!(arr.len(), 1);

        let rmpv::Value::Map(pairs) = &arr[0] else {
            panic!("row should be a map, got {:?}", arr[0]);
        };
        let lat = pairs
            .iter()
            .find(|(k, _)| k.as_str() == Some("lat"))
            .unwrap();
        assert_eq!(lat.1.as_f64(), Some(33.0));

        let temp = pairs
            .iter()
            .find(|(k, _)| k.as_str() == Some("temperature"))
            .unwrap();
        assert_eq!(temp.1.as_f64(), Some(22.5));
    }

    #[test]
    fn duct_cutoff_ghz_reads_best_duct_freq_from_cell() {
        let mut grid = surface_only_grid();
        grid.push_plane("duct_count", vec![2.0]);
        grid.push_plane("max_duct_thickness_m", vec![120.0]);
        grid.push_plane("best_duct_freq_ghz", vec![18.4]);

        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();
        let row = derive(&grid, vt).expect("row");

        assert_eq!(row.duct_strength, Some(120.0));
        let cutoff = row.duct_cutoff_ghz.expect("duct_cutoff_ghz");
        assert!(
            (cutoff - 18.4).abs() < 1e-3,
            "expected ~18.4 GHz, got {cutoff}"
        );
    }

    #[test]
    fn duct_cutoff_ghz_nil_when_no_ducts() {
        let mut grid = surface_only_grid();
        grid.push_plane("duct_count", vec![0.0]);
        grid.push_plane("best_duct_freq_ghz", vec![18.4]);

        let vt = Utc.with_ymd_and_hms(2026, 4, 29, 12, 0, 0).unwrap();
        let row = derive(&grid, vt).expect("row");

        assert_eq!(row.duct_cutoff_ghz, None);
    }

    #[test]
    fn chunk_band_floors_match_elixir() {
        // floor(33.0 / 5) = 6, floor(-97.0 / 5) = -20 (matches Elixir's
        // `chunk_band` Float.floor/trunc).
        assert_eq!(chunk_band(33.0), 6);
        assert_eq!(chunk_band(-97.0), -20);
        assert_eq!(chunk_band(35.0), 7);
        assert_eq!(chunk_band(-95.0), -19);
    }
}