prop/rust/prop_grid_rs/src/decoder.rs

//! wgrib2 subprocess wrapper. 1:1 port of the hot paths in
//! `lib/microwaveprop/weather/grib2/wgrib2.ex` used by f01..f18:
//!   * `extract_grid` (write GRIB2 to tmp → `wgrib2 -lola … bin` → parse)
//!   * inventory parsing from stdout
//!   * `parse_lola_binary` (Fortran-unformatted IEEE 754 little-endian f32)
//!
//! wgrib2 is called as a child process exactly as the Elixir version
//! does — same CLI flags, same temp-file lifecycle. Output binary is
//! stride-addressed: each message = 4-byte LE record length + N*4 bytes
//! of data + 4-byte duplicate record length.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::process::{Command, Output};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

use byteorder::{ByteOrder, LittleEndian};

use crate::grid::{round3, GridSpec};

const UNDEFINED_VALUE: f32 = 9.999e20;
static UNIQUE: AtomicU64 = AtomicU64::new(1);

#[derive(Debug, thiserror::Error)]
pub enum DecodeError {
    #[error("wgrib2 not available on PATH")]
    Wgrib2NotAvailable,
    #[error("wgrib2 failed (exit {code}): {stderr}")]
    Wgrib2Failed { code: i32, stderr: String },
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
}

#[derive(Debug, Clone)]
pub struct Message {
    pub var: String,
    pub level: String,
}

/// Per-cell: `"VAR:LEVEL" -> f32`. Key format matches Elixir exactly so
/// scorer input can be wired unchanged.
pub type CellValues = HashMap<String, f32>;

/// Lat/lon rounded to 3 decimals. f64 keys can't be hashed directly, so
/// we key by `(lat_mdeg, lon_mdeg)` where `*_mdeg = round(x * 1000) as i32`.
pub type PointGrid = HashMap<(i32, i32), CellValues>;

pub fn wgrib2_available() -> bool {
    which_wgrib2().is_some()
}

fn which_wgrib2() -> Option<PathBuf> {
    if let Ok(val) = std::env::var("WGRIB2") {
        return Some(PathBuf::from(val));
    }
    let path_var = std::env::var_os("PATH")?;
    for dir in std::env::split_paths(&path_var) {
        let candidate = dir.join("wgrib2");
        if candidate.is_file() {
            return Some(candidate);
        }
    }
    None
}

pub fn normalize_lon(lon: f64) -> f64 {
    if lon < 0.0 {
        lon + 360.0
    } else {
        lon
    }
}

pub fn denormalize_lon(lon: f64) -> f64 {
    if lon > 180.0 {
        lon - 360.0
    } else {
        lon
    }
}

/// Extract values at the given grid spec from an in-memory GRIB2 blob.
/// Writes the blob to a temp file, runs wgrib2, parses the output.
pub fn extract_grid(
    grib: &[u8],
    match_pattern: &str,
    grid_spec: GridSpec,
) -> Result<PointGrid, DecodeError> {
    let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?;
    let tmp_dir = std::env::temp_dir();
    let uniq = UNIQUE.fetch_add(1, Ordering::Relaxed);
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let grib_path = tmp_dir.join(format!("hrrr_{nanos}_{uniq}.grib2"));
    let bin_path = tmp_dir.join(format!("hrrr_{nanos}_{uniq}.grib2.lola.bin"));

    let result: Result<PointGrid, DecodeError> = (|| {
        std::fs::write(&grib_path, grib)?;
        run_wgrib2_lola(&wgrib2, &grib_path, match_pattern, &grid_spec, &bin_path)
    })();

    let _ = std::fs::remove_file(&grib_path);
    let _ = std::fs::remove_file(&bin_path);
    result
}

/// File-path variant — avoids the intermediate `write(tmp, binary)` copy.
pub fn extract_grid_from_file(
    grib_path: &Path,
    match_pattern: &str,
    grid_spec: GridSpec,
) -> Result<PointGrid, DecodeError> {
    let wgrib2 = which_wgrib2().ok_or(DecodeError::Wgrib2NotAvailable)?;
    let bin_path = sibling_tmp_path(grib_path, "lola");
    let result = run_wgrib2_lola(&wgrib2, grib_path, match_pattern, &grid_spec, &bin_path);
    let _ = std::fs::remove_file(&bin_path);
    result
}

fn sibling_tmp_path(path: &Path, suffix: &str) -> PathBuf {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_nanos())
        .unwrap_or(0);
    let uniq = UNIQUE.fetch_add(1, Ordering::Relaxed);
    let mut s = path.as_os_str().to_owned();
    s.push(format!(".{suffix}.{nanos}.{uniq}.bin"));
    PathBuf::from(s)
}

fn run_wgrib2_lola(
    wgrib2: &Path,
    grib_path: &Path,
    match_pattern: &str,
    grid_spec: &GridSpec,
    bin_path: &Path,
) -> Result<PointGrid, DecodeError> {
    let lon_spec = format!(
        "{}:{}:{}",
        normalize_lon(grid_spec.lon_start),
        grid_spec.lon_count,
        grid_spec.lon_step
    );
    let lat_spec = format!(
        "{}:{}:{}",
        grid_spec.lat_start, grid_spec.lat_count, grid_spec.lat_step
    );

    let Output { status, stdout, stderr } = Command::new(wgrib2)
        .arg(grib_path)
        .arg("-match")
        .arg(match_pattern)
        .arg("-lola")
        .arg(&lon_spec)
        .arg(&lat_spec)
        .arg(bin_path)
        .arg("bin")
        .output()?;

    if !status.success() {
        let mut combined = stdout;
        combined.extend_from_slice(&stderr);
        let text = String::from_utf8_lossy(&combined).to_string();
        let snippet: String = text.chars().take(200).collect();
        return Err(DecodeError::Wgrib2Failed {
            code: status.code().unwrap_or(-1),
            stderr: snippet,
        });
    }

    let inventory = String::from_utf8_lossy(&stdout);
    let messages = parse_inventory(&inventory);

    match std::fs::read(bin_path) {
        Ok(bin) => Ok(parse_lola_binary(&bin, &messages, grid_spec)),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(HashMap::new()),
        Err(e) => Err(e.into()),
    }
}

/// Parse `wgrib2`'s stdout inventory (`"n:offset:date:var:level:..."`).
/// Malformed lines are dropped.
pub fn parse_inventory(out: &str) -> Vec<Message> {
    out.lines()
        .filter_map(|line| {
            let mut parts = line.splitn(8, ':');
            let _n = parts.next()?;
            let _offset = parts.next()?;
            let _date = parts.next()?;
            let var = parts.next()?;
            let level = parts.next()?;
            Some(Message {
                var: var.to_string(),
                level: level.to_string(),
            })
        })
        .collect()
}

/// Parse the Fortran-unformatted output file from `wgrib2 -lola`.
/// Layout: per message = 4-byte LE u32 record length + N·4 bytes f32 LE
/// data + 4-byte LE u32 record length. `N = nx * ny`. Cells whose value
/// exceeds `UNDEFINED_VALUE / 2` are treated as missing (the sentinel
/// wgrib2 writes for "no data").
pub fn parse_lola_binary(bin: &[u8], messages: &[Message], grid_spec: &GridSpec) -> PointGrid {
    let nx = grid_spec.lon_count;
    let ny = grid_spec.lat_count;
    let bytes_per_msg = nx * ny * 4;
    let record_overhead = 8;
    let stride = bytes_per_msg + record_overhead;

    let mut out: PointGrid = HashMap::with_capacity(nx * ny);

    for (msg_idx, msg) in messages.iter().enumerate() {
        let data_offset = msg_idx * stride + 4;
        if data_offset + bytes_per_msg > bin.len() {
            continue;
        }
        let chunk = &bin[data_offset..data_offset + bytes_per_msg];
        let key = format!("{}:{}", msg.var, msg.level);

        for j in 0..ny {
            let lat = round3(grid_spec.lat_start + j as f64 * grid_spec.lat_step);
            let lat_key = (lat * 1000.0).round() as i32;
            for i in 0..nx {
                let cell_offset = (j * nx + i) * 4;
                let value = LittleEndian::read_f32(&chunk[cell_offset..cell_offset + 4]);
                if value > UNDEFINED_VALUE / 2.0 {
                    continue;
                }
                let raw_lon = grid_spec.lon_start + i as f64 * grid_spec.lon_step;
                let lon = round3(denormalize_lon(raw_lon));
                let lon_key = (lon * 1000.0).round() as i32;
                out.entry((lat_key, lon_key))
                    .or_default()
                    .insert(key.clone(), value);
            }
        }
    }

    out
}

pub fn key_to_latlon(key: (i32, i32)) -> (f64, f64) {
    (key.0 as f64 / 1000.0, key.1 as f64 / 1000.0)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn lon_round_trip() {
        for lon in [-125.0, -97.0, -66.0, 0.0, 179.5] {
            assert!((denormalize_lon(normalize_lon(lon)) - lon).abs() < 1e-9);
        }
    }

    #[test]
    fn inventory_parses_var_level() {
        let text = "1:0:d=2026041915:TMP:2 m above ground:anl:\n\
                    2:1234:d=2026041915:DPT:2 m above ground:anl:\n";
        let msgs = parse_inventory(text);
        assert_eq!(msgs.len(), 2);
        assert_eq!(msgs[0].var, "TMP");
        assert_eq!(msgs[1].level, "2 m above ground");
    }

    #[test]
    fn inventory_drops_short_lines() {
        let text = "garbage\n1:0:d:TMP:surface:anl:\n";
        let msgs = parse_inventory(text);
        assert_eq!(msgs.len(), 1);
        assert_eq!(msgs[0].var, "TMP");
    }

    #[test]
    fn parse_lola_binary_reconstructs_grid() {
        // Build a synthetic 2-message, 3×2 grid. Layout matches wgrib2:
        //   [4 bytes len][6 f32 values][4 bytes len] per message.
        let spec = GridSpec {
            lon_start: -100.0,
            lon_count: 3,
            lon_step: 0.5,
            lat_start: 30.0,
            lat_count: 2,
            lat_step: 0.5,
        };
        let msgs = vec![
            Message { var: "TMP".into(), level: "surface".into() },
            Message { var: "DPT".into(), level: "surface".into() },
        ];
        let mut buf = Vec::new();
        for (msg_idx, _msg) in msgs.iter().enumerate() {
            let len: u32 = 3 * 2 * 4;
            buf.extend_from_slice(&len.to_le_bytes());
            for j in 0..2 {
                for i in 0..3 {
                    let v = (msg_idx as f32) * 100.0 + j as f32 * 10.0 + i as f32;
                    buf.extend_from_slice(&v.to_le_bytes());
                }
            }
            buf.extend_from_slice(&len.to_le_bytes());
        }

        let grid = parse_lola_binary(&buf, &msgs, &spec);
        assert_eq!(grid.len(), 6);

        // Cell (30.0, -100.0) → i=0, j=0, msg0 value = 0.0, msg1 = 100.0
        let k = ((30.0 * 1000.0) as i32, (-100.0 * 1000.0) as i32);
        let cell = grid.get(&k).unwrap();
        assert_eq!(cell["TMP:surface"], 0.0);
        assert_eq!(cell["DPT:surface"], 100.0);

        // Cell (30.5, -99.0) → i=2, j=1, msg0 = 12.0
        let k2 = ((30.5 * 1000.0) as i32, (-99.0 * 1000.0) as i32);
        let cell2 = grid.get(&k2).unwrap();
        assert_eq!(cell2["TMP:surface"], 12.0);
        assert_eq!(cell2["DPT:surface"], 112.0);
    }

    #[test]
    fn parse_lola_skips_undefined_sentinel() {
        let spec = GridSpec {
            lon_start: -100.0,
            lon_count: 2,
            lon_step: 1.0,
            lat_start: 30.0,
            lat_count: 1,
            lat_step: 1.0,
        };
        let msgs = vec![Message { var: "TMP".into(), level: "surface".into() }];
        let mut buf = Vec::new();
        let len: u32 = 2 * 4;
        buf.extend_from_slice(&len.to_le_bytes());
        buf.extend_from_slice(&UNDEFINED_VALUE.to_le_bytes()); // i=0
        buf.extend_from_slice(&42.0f32.to_le_bytes()); // i=1
        buf.extend_from_slice(&len.to_le_bytes());

        let grid = parse_lola_binary(&buf, &msgs, &spec);
        assert_eq!(grid.len(), 1); // only i=1 survived
        let k = ((30.0 * 1000.0) as i32, (-99.0 * 1000.0) as i32);
        assert_eq!(grid.get(&k).unwrap()["TMP:surface"], 42.0);
    }

    #[test]
    fn key_to_latlon_roundtrip() {
        let (lat, lon) = key_to_latlon((32_500, -97_250));
        assert_eq!(lat, 32.5);
        assert_eq!(lon, -97.25);
    }
}