prop/rust/prop_grid_rs/src/pgrid.rs

//! `.pgrid` — dense, cell-major, random-access profile grid.
//!
//! Replaces the gzipped-MessagePack `.mp.gz` profile artifact. Measured
//! on a full CONUS grid (95,073 cells, 13 pressure levels):
//!
//! | | `.mp.gz` | `.pgrid` |
//! |---|---|---|
//! | write | 3.96 s (rmpv tree + gzip -9) | one `write_all` of 22.4 MB |
//! | size | 22.0 MB | 22.4 MB |
//! | read one cell | gunzip + unpack + atomize **the whole file** | one 236-byte `pread` |
//!
//! The write happens 30× an hour (f00 + f01..f48); the single-cell read
//! happens on every `/map` point click, every Skew-T load and every
//! `PathCompute` call. Both were paying for a format designed for
//! neither.
//!
//! ## Layout (little-endian)
//!
//! ```text
//! magic        4   "PGRD"
//! version      1   0x01
//! flags        1   bit0: 0 = hrrr, 1 = hrdps
//! n_fields     2   u16
//! valid_time   8   i64  unix seconds
//! lat_start    8   f64
//! lon_start    8   f64
//! lat_step     8   f64
//! lon_step     8   f64
//! n_rows       2   u16  (latitude)
//! n_cols       2   u16  (longitude)
//! field_table  n_fields × 32   NUL-padded ASCII field names
//! body         n_rows*n_cols*n_fields × 4   f32, CELL-MAJOR
//! ```
//!
//! **Cell-major** (`body[(cell * n_fields + field)]`) rather than
//! plane-major: reading one cell across all fields — the thing the UI
//! actually does — is then a single contiguous `pread`. A viewport read
//! is one contiguous `pread` per grid row.
//!
//! `f32::NAN` is the missing-value sentinel, matching
//! [`FieldGrid`][crate::field_grid::FieldGrid].
//!
//! The field table is written into the header rather than being implied
//! by the version, so the Elixir reader resolves fields by name. Adding a
//! field is backwards compatible: an older reader simply does not look
//! it up.

use std::io::Write;
use std::path::{Path, PathBuf};

use chrono::{DateTime, Utc};

use crate::fetcher;
use crate::grid::GridSpec;

pub const MAGIC: &[u8; 4] = b"PGRD";
pub const VERSION: u8 = 1;
/// Fixed width of one field-table entry, in bytes.
pub const FIELD_NAME_LEN: usize = 32;

/// Per-cell scalar fields, in on-disk order. Names match the atom keys
/// `Microwaveprop.Propagation.ProfilesFile` already whitelists, so the
/// Elixir reader can hand callers the same map shape the `.mp.gz` path
/// produced.
pub const SCALAR_FIELDS: &[&str] = &[
    "surface_temp_c",
    "surface_dewpoint_c",
    "surface_pressure_mb",
    "hpbl_m",
    "pwat_mm",
    "wind_u",
    "wind_v",
    "cloud_cover_pct",
    "precip_mm",
    "native_min_gradient",
    "best_duct_freq_ghz",
    "max_duct_thickness_m",
    "duct_count",
    "nexrad_max_reflectivity_dbz",
    "commercial_degradation_db",
    "commercial_baseline_dbm",
    "commercial_current_dbm",
    "commercial_n_links",
    "surface_refractivity",
    "min_refractivity_gradient",
];

/// Full field list: the scalars above, then `hght_m_<P>mb`, `tmpc_<P>mb`,
/// `dwpc_<P>mb` for each pressure level. `pres_mb` is implied by the
/// field name, so it costs no bytes.
pub fn field_names() -> &'static [String] {
    static NAMES: std::sync::OnceLock<Vec<String>> = std::sync::OnceLock::new();
    NAMES.get_or_init(|| {
        let mut v: Vec<String> = SCALAR_FIELDS.iter().map(|s| (*s).to_string()).collect();
        for &p in fetcher::GRID_PRESSURE_LEVELS {
            v.push(format!("hght_m_{p}mb"));
            v.push(format!("tmpc_{p}mb"));
            v.push(format!("dwpc_{p}mb"));
        }
        for n in &v {
            assert!(
                n.len() < FIELD_NAME_LEN,
                "field name {n} exceeds {FIELD_NAME_LEN} bytes"
            );
        }
        v
    })
}

pub fn n_fields() -> usize {
    field_names().len()
}

/// Byte offset of the body, i.e. the header + field table length.
pub fn header_len(n_fields: usize) -> usize {
    4 + 1 + 1 + 2 + 8 + 8 + 8 + 8 + 8 + 2 + 2 + n_fields * FIELD_NAME_LEN
}

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

/// `<scores_dir>/profiles/<iso>.pgrid`.
pub fn path_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("profiles").join(format!("{iso}.pgrid"))
}

/// One cell's record: `n_fields()` f32 values in `field_names()` order.
/// The pipeline fills these directly, so no intermediate map is built.
pub type Record = Vec<f32>;

pub fn empty_record() -> Record {
    vec![f32::NAN; n_fields()]
}

/// Serialise the header for `spec` / `valid_time`.
pub fn encode_header(spec: &GridSpec, valid_time: DateTime<Utc>, hrdps: bool) -> Vec<u8> {
    let names = field_names();
    let mut out = Vec::with_capacity(header_len(names.len()));
    out.extend_from_slice(MAGIC);
    out.push(VERSION);
    out.push(u8::from(hrdps));
    out.extend_from_slice(&(names.len() as u16).to_le_bytes());
    out.extend_from_slice(&valid_time.timestamp().to_le_bytes());
    out.extend_from_slice(&spec.lat_start.to_le_bytes());
    out.extend_from_slice(&spec.lon_start.to_le_bytes());
    out.extend_from_slice(&spec.lat_step.to_le_bytes());
    out.extend_from_slice(&spec.lon_step.to_le_bytes());
    out.extend_from_slice(&(spec.lat_count as u16).to_le_bytes());
    out.extend_from_slice(&(spec.lon_count as u16).to_le_bytes());
    for name in names {
        let mut buf = [0u8; FIELD_NAME_LEN];
        buf[..name.len()].copy_from_slice(name.as_bytes());
        out.extend_from_slice(&buf);
    }
    debug_assert_eq!(out.len(), header_len(names.len()));
    out
}

/// Write `<scores_dir>/profiles/<iso>.pgrid` atomically (tmp + rename,
/// so an NFS reader sees the old file, the new file, or nothing).
///
/// `body` is the flat cell-major `f32` array, length
/// `n_cells * n_fields()`.
pub fn write_atomic(
    scores_dir: &Path,
    valid_time: DateTime<Utc>,
    spec: &GridSpec,
    body: &[f32],
    hrdps: bool,
) -> Result<PathBuf, WriteError> {
    assert_eq!(
        body.len(),
        spec.lat_count * spec.lon_count * n_fields(),
        "pgrid body does not match grid spec"
    );

    let path = path_for(scores_dir, valid_time);
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }

    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 file = std::fs::File::create(&tmp)?;
        let mut w = std::io::BufWriter::with_capacity(1 << 20, file);
        w.write_all(&encode_header(spec, valid_time, hrdps))?;
        // f32 little-endian. On LE hosts this is the in-memory
        // representation, so it is a bulk copy; the explicit conversion
        // keeps the format host-independent.
        let mut chunk: Vec<u8> = Vec::with_capacity(4 * 8192);
        for values in body.chunks(8192) {
            chunk.clear();
            for v in values {
                chunk.extend_from_slice(&v.to_le_bytes());
            }
            w.write_all(&chunk)?;
        }
        w.flush()?;
    }

    match std::fs::rename(&tmp, &path) {
        Ok(()) => Ok(path),
        Err(e) => {
            let _ = std::fs::remove_file(&tmp);
            Err(WriteError::Io(e))
        }
    }
}

// ── Reader (tests + any Rust-side consumer) ──────────────────────────

#[derive(Debug, Clone)]
pub struct Header {
    pub hrdps: bool,
    pub valid_time: DateTime<Utc>,
    pub spec: GridSpec,
    pub fields: Vec<String>,
}

impl Header {
    pub fn field_index(&self, name: &str) -> Option<usize> {
        self.fields.iter().position(|f| f == name)
    }

    pub fn n_cells(&self) -> usize {
        self.spec.lat_count * self.spec.lon_count
    }
}

pub fn decode_header(bytes: &[u8]) -> Option<Header> {
    if bytes.len() < 52 || &bytes[0..4] != MAGIC || bytes[4] != VERSION {
        return None;
    }
    let hrdps = bytes[5] != 0;
    let n_fields = u16::from_le_bytes(bytes[6..8].try_into().ok()?) as usize;
    let valid_time =
        DateTime::<Utc>::from_timestamp(i64::from_le_bytes(bytes[8..16].try_into().ok()?), 0)?;
    let lat_start = f64::from_le_bytes(bytes[16..24].try_into().ok()?);
    let lon_start = f64::from_le_bytes(bytes[24..32].try_into().ok()?);
    let lat_step = f64::from_le_bytes(bytes[32..40].try_into().ok()?);
    let lon_step = f64::from_le_bytes(bytes[40..48].try_into().ok()?);
    let lat_count = u16::from_le_bytes(bytes[48..50].try_into().ok()?) as usize;
    let lon_count = u16::from_le_bytes(bytes[50..52].try_into().ok()?) as usize;

    if bytes.len() < header_len(n_fields) {
        return None;
    }
    let mut fields = Vec::with_capacity(n_fields);
    for i in 0..n_fields {
        let start = 52 + i * FIELD_NAME_LEN;
        let raw = &bytes[start..start + FIELD_NAME_LEN];
        let end = raw.iter().position(|&b| b == 0).unwrap_or(FIELD_NAME_LEN);
        fields.push(String::from_utf8_lossy(&raw[..end]).into_owned());
    }

    Some(Header {
        hrdps,
        valid_time,
        spec: GridSpec {
            lat_start,
            lon_start,
            lat_step,
            lon_step,
            lat_count,
            lon_count,
        },
        fields,
    })
}

/// Read one cell's record out of a whole-file buffer. Mirrors what the
/// Elixir reader does with `:file.pread/3`.
pub fn read_cell(bytes: &[u8], header: &Header, cell: usize) -> Option<Vec<f32>> {
    let n = header.fields.len();
    let start = header_len(n) + cell * n * 4;
    let end = start + n * 4;
    if end > bytes.len() {
        return None;
    }
    Some(
        bytes[start..end]
            .chunks_exact(4)
            .map(|b| f32::from_le_bytes(b.try_into().unwrap()))
            .collect(),
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::grid;
    use chrono::TimeZone;

    fn spec_3x2() -> GridSpec {
        GridSpec {
            lon_start: -100.0,
            lon_count: 3,
            lon_step: 0.5,
            lat_start: 30.0,
            lat_count: 2,
            lat_step: 0.5,
        }
    }

    #[test]
    fn field_list_shape() {
        let names = field_names();
        assert_eq!(
            names.len(),
            SCALAR_FIELDS.len() + fetcher::GRID_PRESSURE_LEVELS.len() * 3
        );
        assert_eq!(names[0], "surface_temp_c");
        assert!(names.contains(&"tmpc_850mb".to_string()));
        assert!(names.contains(&"dwpc_700mb".to_string()));
        assert!(names.contains(&"hght_m_1000mb".to_string()));
    }

    #[test]
    fn header_round_trips() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = grid::wgrib2_grid_spec();
        let bytes = encode_header(&spec, vt, false);
        let h = decode_header(&bytes).expect("decodes");

        assert!(!h.hrdps);
        assert_eq!(h.valid_time, vt);
        assert_eq!(h.spec, spec);
        assert_eq!(h.fields, field_names());
        assert_eq!(bytes.len(), header_len(n_fields()));
    }

    #[test]
    fn header_carries_hrdps_flag() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let bytes = encode_header(&grid::hrdps_grid_spec(), vt, true);
        assert!(decode_header(&bytes).unwrap().hrdps);
    }

    #[test]
    fn decode_rejects_bad_magic_and_version() {
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let mut bytes = encode_header(&spec_3x2(), vt, false);
        let good = bytes.clone();
        bytes[0..4].copy_from_slice(b"XXXX");
        assert!(decode_header(&bytes).is_none());

        let mut bytes = good;
        bytes[4] = 99;
        assert!(decode_header(&bytes).is_none());
    }

    #[test]
    fn write_read_round_trip_per_cell() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let n = spec.lat_count * spec.lon_count;
        let nf = n_fields();

        // Cell c, field f gets value c*100 + f, except cell 1 which is
        // all-missing so the NaN sentinel is exercised.
        let mut body = vec![f32::NAN; n * nf];
        for c in 0..n {
            if c == 1 {
                continue;
            }
            for f in 0..nf {
                body[c * nf + f] = (c * 100 + f) as f32;
            }
        }

        let path = write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        assert!(path.to_string_lossy().ends_with(".pgrid"));

        let raw = std::fs::read(&path).unwrap();
        let h = decode_header(&raw).unwrap();
        assert_eq!(h.spec, spec);
        assert_eq!(
            raw.len(),
            header_len(nf) + n * nf * 4,
            "file is header + dense body, nothing else"
        );

        let cell0 = read_cell(&raw, &h, 0).unwrap();
        assert_eq!(cell0[0], 0.0);
        assert_eq!(cell0[nf - 1], (nf - 1) as f32);

        let cell2 = read_cell(&raw, &h, 2).unwrap();
        assert_eq!(cell2[0], 200.0);

        let cell1 = read_cell(&raw, &h, 1).unwrap();
        assert!(
            cell1.iter().all(|v| v.is_nan()),
            "missing cell reads as NaN"
        );

        assert!(read_cell(&raw, &h, n).is_none(), "out-of-range cell");
    }

    #[test]
    fn named_field_lookup_resolves_to_the_written_value() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let nf = n_fields();
        let mut body = vec![f32::NAN; spec.lat_count * spec.lon_count * nf];

        let names = field_names();
        let temp_idx = names.iter().position(|n| n == "surface_temp_c").unwrap();
        let t850_idx = names.iter().position(|n| n == "tmpc_850mb").unwrap();
        body[4 * nf + temp_idx] = 22.5;
        body[4 * nf + t850_idx] = 12.75;

        let path = write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        let raw = std::fs::read(&path).unwrap();
        let h = decode_header(&raw).unwrap();
        let rec = read_cell(&raw, &h, 4).unwrap();

        assert_eq!(rec[h.field_index("surface_temp_c").unwrap()], 22.5);
        assert_eq!(rec[h.field_index("tmpc_850mb").unwrap()], 12.75);
        assert!(h.field_index("no_such_field").is_none());
    }

    #[test]
    fn atomic_write_leaves_no_tmp_files() {
        let dir = tempfile::tempdir().unwrap();
        let vt = Utc.with_ymd_and_hms(2026, 4, 19, 14, 0, 0).unwrap();
        let spec = spec_3x2();
        let body = vec![f32::NAN; spec.lat_count * spec.lon_count * n_fields()];

        write_atomic(dir.path(), vt, &spec, &body, false).unwrap();
        write_atomic(dir.path(), vt, &spec, &body, false).unwrap();

        let entries: Vec<_> = std::fs::read_dir(dir.path().join("profiles"))
            .unwrap()
            .map(|e| e.unwrap().file_name().into_string().unwrap())
            .collect();
        assert_eq!(entries.len(), 1, "got {entries:?}");
        assert!(entries[0].ends_with(".pgrid"));
    }
}