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, CellValues, DecodeError, PointGrid};
use crate::duct::{self, NativeProfile};
use crate::fetcher::{self, HrrrClient};
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<(i32, i32), 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: PointGrid =
        tokio::task::spawn_blocking(move || decoder::extract_grid(&blob, &pattern, grid_spec))
            .await
            .expect("blocking join")?;

    Ok(reduce_grid_to_ducts(grid))
}

/// Per-cell reducer: turn the `CellValues` map into a
/// `NativeProfile` and run the duct analyser.
pub fn reduce_grid_to_ducts(grid: PointGrid) -> HashMap<(i32, i32), DuctMetrics> {
    let mut out = HashMap::with_capacity(grid.len());
    for (key, cell) in grid {
        if let Some(profile) = build_native_profile(&cell) {
            let analysis = duct::analyze(&profile);
            out.insert(
                key,
                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
}

/// Build a `NativeProfile` from one cell's raw values. Mirrors the
/// Elixir `build_native_profile/1`: walk all 50 hybrid levels, keep
/// only those with both HGT and TMP present, then sort ascending by
/// height. Returns `None` if fewer than 3 levels survive (too sparse
/// to analyse).
pub fn build_native_profile(cell: &CellValues) -> Option<NativeProfile> {
    // (height, tmp, spfh, pres) per level. Filter out rows where HGT
    // or TMP is absent — those levels are junk.
    let mut levels: Vec<(f64, f64, f64, f64)> = Vec::with_capacity(NATIVE_LEVEL_COUNT as usize);
    for level in 1..=NATIVE_LEVEL_COUNT {
        // `cell.get` dispatches via Arc<str>'s Borrow<str>, so pass
        // &str refs built by `format!(…).as_str()` (or inline &str
        // concat) — passing `&String` triggers the wrong Borrow impl.
        let lvl_str = format!("{level} hybrid level");
        let hgt_key = format!("HGT:{lvl_str}");
        let tmp_key = format!("TMP:{lvl_str}");
        let spfh_key = format!("SPFH:{lvl_str}");
        let pres_key = format!("PRES:{lvl_str}");
        let Some(&hgt) = cell.get(hgt_key.as_str()) else {
            continue;
        };
        let Some(&tmp) = cell.get(tmp_key.as_str()) 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 = cell.get(spfh_key.as_str()).copied().unwrap_or(0.0);
        let pres = cell.get(pres_key.as_str()).copied().unwrap_or(0.0);
        levels.push((hgt as f64, tmp as f64, spfh as f64, pres as f64));
    }

    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 their previous contents; cells present overwrite the
/// `native_min_gradient` / `best_duct_freq_ghz` / etc. fields.
///
/// Consumers that only need the raw duct HashMap can skip this; the
/// merge helper exists for symmetry with the Elixir call site.
pub fn merge_duct_grid(
    mut base: HashMap<(i32, i32), CellValues>,
    ducts: &HashMap<(i32, i32), DuctMetrics>,
) -> HashMap<(i32, i32), CellValues> {
    for (key, metrics) in ducts {
        if let Some(cell) = base.get_mut(key) {
            cell.insert(
                "native_min_gradient".into(),
                metrics.native_min_gradient as f32,
            );
            if let Some(f) = metrics.best_duct_freq_ghz {
                cell.insert("best_duct_freq_ghz".into(), f as f32);
            }
            if let Some(t) = metrics.max_duct_thickness_m {
                cell.insert("max_duct_thickness_m".into(), t as f32);
            }
            cell.insert("duct_count".into(), metrics.duct_count as f32);
        }
    }
    base
}

/// 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<(i32, i32), 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 cell_with_levels(count: usize) -> CellValues {
        let mut c = CellValues::new();
        for i in 1..=count {
            c.insert(format!("HGT:{i} hybrid level").into(), (i as f32) * 50.0);
            c.insert(
                format!("TMP:{i} hybrid level").into(),
                290.0 - (i as f32) * 0.5,
            );
            c.insert(
                format!("SPFH:{i} hybrid level").into(),
                0.008 - (i as f32) * 0.0001,
            );
            c.insert(
                format!("PRES:{i} hybrid level").into(),
                101_000.0 - (i as f32) * 600.0,
            );
        }
        c
    }

    #[test]
    fn build_profile_drops_cells_with_fewer_than_3_levels() {
        assert!(build_native_profile(&cell_with_levels(0)).is_none());
        assert!(build_native_profile(&cell_with_levels(2)).is_none());
        assert!(build_native_profile(&cell_with_levels(3)).is_some());
    }

    #[test]
    fn build_profile_sorts_by_height() {
        let mut cell = CellValues::new();
        // Interleaved: level 3 comes first in insertion order, level 1 last.
        cell.insert("HGT:3 hybrid level".into(), 150.0);
        cell.insert("TMP:3 hybrid level".into(), 287.0);
        cell.insert("HGT:1 hybrid level".into(), 10.0);
        cell.insert("TMP:1 hybrid level".into(), 291.0);
        cell.insert("HGT:2 hybrid level".into(), 50.0);
        cell.insert("TMP:2 hybrid level".into(), 289.0);
        let p = build_native_profile(&cell).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 mut grid: PointGrid = HashMap::new();
        grid.insert((0, 0), cell_with_levels(50));
        grid.insert((1, 0), cell_with_levels(2)); // too sparse
        let out = reduce_grid_to_ducts(grid);
        assert_eq!(out.len(), 1);
        let m = &out[&(0, 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_only_touches_shared_keys() {
        let mut base: HashMap<(i32, i32), CellValues> = HashMap::new();
        base.insert((0, 0), {
            let mut c = CellValues::new();
            c.insert("PRES:surface".into(), 101_000.0);
            c
        });
        let ducts = HashMap::from([(
            (0, 0),
            DuctMetrics {
                native_min_gradient: -150.0,
                best_duct_freq_ghz: Some(24.0),
                max_duct_thickness_m: Some(80.0),
                duct_count: 1,
            },
        )]);
        let merged = merge_duct_grid(base, &ducts);
        let cell = &merged[&(0, 0)];
        assert_eq!(cell.get("PRES:surface").copied(), Some(101_000.0));
        assert_eq!(cell.get("native_min_gradient").copied(), Some(-150.0));
        assert_eq!(cell.get("best_duct_freq_ghz").copied(), Some(24.0));
        assert_eq!(cell.get("duct_count").copied(), Some(1.0));
    }
}