prop/rust/prop_grid_rs/src/native_duct.rs

//! HRRR native hybrid-sigma duct detection for the CONUS grid.
//!
//! Port of `Microwaveprop.Weather.HrrrNativeClient.fetch_native_duct_grid`.
//! The native-level HRRR file (`wrfnatf{fh}.grib2`) is ~566 MB with TMP /
//! SPFH / HGT / PRES on all 50 hybrid-sigma levels — the ~10-50 m vertical
//! resolution needed to resolve surface ducts and boundary-layer
//! inversions. `HrrrClient` only covers surface + pressure levels, which
//! is ~250 m resolution near the surface and loses the ducts entirely.
//!
//! Flow: fetch idx → select duct byte ranges (4 vars × 50 levels = 200
//! messages, ~300 MB subset) → download via merged byte ranges → run
//! wgrib2 `-lola` to grid → per-cell build a `NativeProfile` and run
//! `duct::analyze`. Returns a HashMap keyed by (lat, lon) to per-cell
//! duct metrics.

use std::collections::HashMap;
use std::path::Path;

use chrono::NaiveDate;

use crate::decoder::{self, DecodeError};
use crate::duct::{self, NativeProfile};
use crate::fetcher::{self, HrrrClient};
use crate::field_grid::{FieldGrid, PlaneId};
use crate::grid::GridSpec;

/// The 4 variables needed for duct detection. UGRD/VGRD/TKE were
/// features we never used; skipping them saves ~230 MB per file.
pub const DUCT_VARIABLES: [&str; 4] = ["TMP", "SPFH", "HGT", "PRES"];

/// HRRR's 50 native hybrid-sigma levels. Level 1 is surface.
pub const NATIVE_LEVEL_COUNT: u8 = 50;

#[derive(Debug, Clone, Copy)]
pub struct DuctMetrics {
    pub native_min_gradient: f64,
    pub best_duct_freq_ghz: Option<f64>,
    pub max_duct_thickness_m: Option<f64>,
    pub duct_count: u32,
}

#[derive(Debug, thiserror::Error)]
pub enum NativeDuctError {
    #[error("fetch: {0}")]
    Fetch(#[from] fetcher::FetchError),
    #[error("decode: {0}")]
    Decode(#[from] DecodeError),
    #[error("io: {0}")]
    Io(#[from] std::io::Error),
}

/// Build the (var, level) tuples for the duct-variable messages. 4
/// vars × 50 levels = 200 messages, matching the Elixir
/// `HrrrNativeClient.duct_messages/0`.
pub fn duct_messages() -> Vec<(String, String)> {
    let mut out = Vec::with_capacity((NATIVE_LEVEL_COUNT as usize) * DUCT_VARIABLES.len());
    for level in 1..=NATIVE_LEVEL_COUNT {
        for var in DUCT_VARIABLES {
            out.push((var.to_string(), format!("{level} hybrid level")));
        }
    }
    out
}

/// wgrib2 `-match` regex: "`:(TMP|SPFH|HGT|PRES):.*hybrid level:`".
pub fn match_pattern() -> String {
    format!(":({}):.*hybrid level:", DUCT_VARIABLES.join("|"))
}

/// URL builder for native-level HRRR files. Re-export of
/// `fetcher::hrrr_native_url` kept for external callers.
pub use crate::fetcher::hrrr_native_url;

/// Fetch the native-level duct grid for a (date, hour, forecast_hour)
/// triple. Downloads the byte-range subset, decodes via wgrib2, and
/// computes duct metrics per cell.
pub async fn fetch_native_duct_grid(
    client: &HrrrClient,
    date: NaiveDate,
    hour: u8,
    forecast_hour: u8,
    grid_spec: GridSpec,
) -> Result<HashMap<usize, DuctMetrics>, NativeDuctError> {
    let messages = duct_messages();
    let blob = client
        .fetch_native_blob(date, hour, forecast_hour, &messages)
        .await?;

    // wgrib2 decode must run on the blocking pool — it forks a
    // subprocess and the stdout read is synchronous IO.
    let pattern = match_pattern();
    let grid: FieldGrid =
        tokio::task::spawn_blocking(move || decoder::extract_grid(&blob, &pattern, grid_spec))
            .await
            .expect("blocking join")?;

    Ok(reduce_grid_to_ducts(&grid))
}

/// Plane ids for the 4 duct variables across all 50 hybrid levels,
/// resolved once per grid instead of formatting 200 lookup keys per
/// cell (the old code built four `String`s per level per cell — 19 M
/// allocations across a CONUS grid).
struct DuctPlanes {
    levels: Vec<(PlaneId, PlaneId, Option<PlaneId>, Option<PlaneId>)>,
}

impl DuctPlanes {
    fn resolve(grid: &FieldGrid) -> Self {
        let mut levels = Vec::with_capacity(NATIVE_LEVEL_COUNT as usize);
        for level in 1..=NATIVE_LEVEL_COUNT {
            let lvl = format!("{level} hybrid level");
            // HGT and TMP are required; a level missing either is junk
            // and was skipped by the old per-cell filter.
            let (Some(hgt), Some(tmp)) = (
                grid.plane_id(&format!("HGT:{lvl}")),
                grid.plane_id(&format!("TMP:{lvl}")),
            ) else {
                continue;
            };
            levels.push((
                hgt,
                tmp,
                grid.plane_id(&format!("SPFH:{lvl}")),
                grid.plane_id(&format!("PRES:{lvl}")),
            ));
        }
        Self { levels }
    }
}

/// Per-cell reducer: build a `NativeProfile` from the dense planes and
/// run the duct analyser. Keyed by cell index rather than millidegree
/// lat/lon so the merge into the surface grid is a direct index write.
pub fn reduce_grid_to_ducts(grid: &FieldGrid) -> HashMap<usize, DuctMetrics> {
    let planes = DuctPlanes::resolve(grid);
    let mut out = HashMap::new();
    let mut scratch: Vec<(f64, f64, f64, f64)> = Vec::with_capacity(planes.levels.len());

    for cell in 0..grid.n_cells() {
        scratch.clear();
        for &(hgt_p, tmp_p, spfh_p, pres_p) in &planes.levels {
            let (Some(hgt), Some(tmp)) = (grid.at(hgt_p, cell), grid.at(tmp_p, cell)) else {
                continue;
            };
            // SPFH / PRES default to 0.0 if absent — the duct math handles
            // degenerate rows (returns 0.0 gradient / no duct) without
            // raising, matching Elixir's nil-tolerance.
            let spfh = grid.at_opt(spfh_p, cell).unwrap_or(0.0);
            let pres = grid.at_opt(pres_p, cell).unwrap_or(0.0);
            scratch.push((hgt as f64, tmp as f64, spfh as f64, pres as f64));
        }
        if let Some(profile) = profile_from_levels(&mut scratch) {
            let analysis = duct::analyze(&profile);
            out.insert(
                cell,
                DuctMetrics {
                    native_min_gradient: duct::min_m_gradient(&profile),
                    best_duct_freq_ghz: analysis.best_duct_band_ghz,
                    max_duct_thickness_m: duct::max_duct_thickness_m(&analysis.ducts),
                    duct_count: analysis.ducts.len() as u32,
                },
            );
        }
    }
    out
}

/// Sort one cell's collected `(height, tmp, spfh, pres)` rows ascending
/// by height and turn them into a `NativeProfile`. Mirrors the Elixir
/// `build_native_profile/1`: returns `None` if fewer than 3 levels
/// survive (too sparse to analyse).
///
/// Takes `&mut` so the caller can reuse one scratch buffer across all
/// ~95 k cells rather than allocating per cell.
fn profile_from_levels(levels: &mut [(f64, f64, f64, f64)]) -> Option<NativeProfile> {
    if levels.len() < 3 {
        return None;
    }

    levels.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

    let heights_m = levels.iter().map(|x| x.0).collect();
    let temp_k = levels.iter().map(|x| x.1).collect();
    let spfh = levels.iter().map(|x| x.2).collect();
    let pressure_pa = levels.iter().map(|x| x.3).collect();
    Some(NativeProfile {
        heights_m,
        temp_k,
        spfh,
        pressure_pa,
    })
}

/// Merge native duct metrics into an existing surface+pressure grid.
/// Mirrors Elixir's `apply_duct_grid` — cells missing from the duct
/// map keep NaN (i.e. "absent", exactly as the old code left the key
/// out); cells present get the four duct planes filled in.
pub fn merge_duct_grid(base: &mut FieldGrid, ducts: &HashMap<usize, DuctMetrics>) {
    // Collect into dense planes, then push once each. Writing four
    // planes in one pass beats four passes over the duct map.
    let n = base.n_cells();
    let mut min_grad = vec![f32::NAN; n];
    let mut best_freq = vec![f32::NAN; n];
    let mut max_thick = vec![f32::NAN; n];
    let mut count = vec![f32::NAN; n];

    for (&cell, metrics) in ducts {
        if cell >= n {
            continue;
        }
        min_grad[cell] = metrics.native_min_gradient as f32;
        if let Some(f) = metrics.best_duct_freq_ghz {
            best_freq[cell] = f as f32;
        }
        if let Some(t) = metrics.max_duct_thickness_m {
            max_thick[cell] = t as f32;
        }
        count[cell] = metrics.duct_count as f32;
    }

    base.push_plane("native_min_gradient", min_grad);
    base.push_plane("best_duct_freq_ghz", best_freq);
    base.push_plane("max_duct_thickness_m", max_thick);
    base.push_plane("duct_count", count);
}

/// Legacy entry used by goldens/tests: decode a GRIB2 blob already
/// on disk with the duct match pattern and reduce.
pub fn duct_grid_from_file(
    grib_path: &Path,
    grid_spec: GridSpec,
) -> Result<HashMap<usize, DuctMetrics>, NativeDuctError> {
    let grid = decoder::extract_grid_from_file(grib_path, &match_pattern(), grid_spec)?;
    Ok(reduce_grid_to_ducts(&grid))
}

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

    #[test]
    fn duct_messages_counts_200() {
        let m = duct_messages();
        assert_eq!(m.len(), 200);
        // First and last samples.
        assert_eq!(m[0], ("TMP".into(), "1 hybrid level".into()));
        assert_eq!(m[m.len() - 1], ("PRES".into(), "50 hybrid level".into()));
    }

    #[test]
    fn match_pattern_includes_all_duct_vars() {
        let p = match_pattern();
        assert_eq!(p, ":(TMP|SPFH|HGT|PRES):.*hybrid level:");
    }

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

    /// Grid where cell 0 carries `cell0_levels` hybrid levels and cell 1
    /// carries `cell1_levels`. Levels beyond a cell's count are NaN, which
    /// is how the dense representation says "absent".
    fn grid_with_levels(cell0_levels: usize, cell1_levels: usize) -> FieldGrid {
        let mut g = FieldGrid::new(two_cell_spec());
        let max = cell0_levels.max(cell1_levels);
        for i in 1..=max {
            let present = |c: usize| if i <= c { 1.0 } else { f32::NAN };
            g.push_plane(
                &format!("HGT:{i} hybrid level"),
                vec![
                    present(cell0_levels) * (i as f32) * 50.0,
                    present(cell1_levels) * (i as f32) * 50.0,
                ],
            );
            g.push_plane(
                &format!("TMP:{i} hybrid level"),
                vec![
                    present(cell0_levels) * (290.0 - (i as f32) * 0.5),
                    present(cell1_levels) * (290.0 - (i as f32) * 0.5),
                ],
            );
            g.push_plane(
                &format!("SPFH:{i} hybrid level"),
                vec![
                    present(cell0_levels) * (0.008 - (i as f32) * 0.0001),
                    present(cell1_levels) * (0.008 - (i as f32) * 0.0001),
                ],
            );
            g.push_plane(
                &format!("PRES:{i} hybrid level"),
                vec![
                    present(cell0_levels) * (101_000.0 - (i as f32) * 600.0),
                    present(cell1_levels) * (101_000.0 - (i as f32) * 600.0),
                ],
            );
        }
        g
    }

    #[test]
    fn profile_from_levels_drops_cells_with_fewer_than_3_levels() {
        assert!(profile_from_levels(&mut []).is_none());
        assert!(profile_from_levels(&mut [(1.0, 2.0, 3.0, 4.0); 2]).is_none());
        assert!(profile_from_levels(&mut [(1.0, 2.0, 3.0, 4.0); 3]).is_some());
    }

    #[test]
    fn profile_from_levels_sorts_by_height() {
        // Interleaved: level 3 first in collection order, level 1 last.
        let mut levels = [
            (150.0, 287.0, 0.0, 0.0),
            (10.0, 291.0, 0.0, 0.0),
            (50.0, 289.0, 0.0, 0.0),
        ];
        let p = profile_from_levels(&mut levels).unwrap();
        assert_eq!(p.heights_m, vec![10.0, 50.0, 150.0]);
        assert_eq!(p.temp_k, vec![291.0, 289.0, 287.0]);
    }

    #[test]
    fn reduce_grid_builds_metrics_for_cells_with_enough_levels() {
        let grid = grid_with_levels(50, 2); // cell 1 too sparse
        let out = reduce_grid_to_ducts(&grid);
        assert_eq!(out.len(), 1);
        let m = &out[&0];
        // Temperature decreases with height and moisture decreases too
        // — constructed profile has no strong inversion, so the
        // count should be 0 but min_gradient is still computed.
        assert!(m.native_min_gradient.is_finite());
    }

    #[test]
    fn merge_duct_grid_fills_only_cells_present_in_the_duct_map() {
        let mut base = FieldGrid::new(two_cell_spec());
        base.push_plane("PRES:surface", vec![101_000.0, 101_000.0]);
        let ducts = HashMap::from([(
            0usize,
            DuctMetrics {
                native_min_gradient: -150.0,
                best_duct_freq_ghz: Some(24.0),
                max_duct_thickness_m: Some(80.0),
                duct_count: 1,
            },
        )]);

        merge_duct_grid(&mut base, &ducts);

        let pres = base.plane_id("PRES:surface").unwrap();
        let grad = base.plane_id("native_min_gradient").unwrap();
        let freq = base.plane_id("best_duct_freq_ghz").unwrap();
        let count = base.plane_id("duct_count").unwrap();

        assert_eq!(base.at(pres, 0), Some(101_000.0));
        assert_eq!(base.at(grad, 0), Some(-150.0));
        assert_eq!(base.at(freq, 0), Some(24.0));
        assert_eq!(base.at(count, 0), Some(1.0));

        // Cell 1 had no duct entry — every duct plane stays absent, which
        // is what the old code achieved by never inserting the key.
        assert_eq!(base.at(grad, 1), None);
        assert_eq!(base.at(count, 1), None);
        assert_eq!(base.at(pres, 1), Some(101_000.0));
    }
}