DCTS/crates/common/src/conv_check.rs

use crate::models::ConvCheckResult;
use regex::Regex;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::path::Path;
use std::sync::OnceLock;

static FORT9_RE: OnceLock<Regex> = OnceLock::new();
static NAN_RE: OnceLock<Regex> = OnceLock::new();

#[derive(Debug, Clone)]
struct Fort9Row {
    depth: i32,
    maximum: f64,
}

/// Parses `fort.9` and evaluates convergence against `chmax`
pub fn check_fort9(path: &Path, chmax: f64) -> ConvCheckResult {
    let file = match File::open(path) {
        Ok(f) => f,
        Err(e) => {
            return ConvCheckResult {
                converged: false,
                max_relc: f64::INFINITY,
                worst_depth: -1,
                last_iter: None,
                n_depths: 0,
                chmax,
                error: Some(format!("Failed to open fort.9: {}", e)),
            }
        }
    };

    let reader = BufReader::new(file);
    let re = FORT9_RE.get_or_init(|| {
        Regex::new(
            r"^\s*(\d+)\s+(\d+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+([-+\dE.]+)\s+(\d+)\s+(\d+)\s*$"
        ).unwrap()
    });

    let mut last_iter: Option<i32> = None;
    let mut cur_iter: Option<i32> = None;
    let mut cur_rows: Vec<Fort9Row> = Vec::new();

    for line in reader.lines().map_while(Result::ok) {
        if let Some(caps) = re.captures(&line) {
            let iter: i32 = match caps[1].parse() {
                Ok(v) => v,
                Err(_) => continue,
            };
            let depth: i32 = match caps[2].parse() {
                Ok(v) => v,
                Err(_) => continue,
            };
            let maximum: f64 = match caps[7].parse() {
                Ok(v) => v,
                Err(_) => continue,
            };

            if cur_iter != Some(iter) {
                cur_iter = Some(iter);
                cur_rows.clear();
            }

            cur_rows.push(Fort9Row { depth, maximum });
            last_iter = Some(iter);
        }
    }

    if cur_rows.is_empty() || last_iter.is_none() {
        return ConvCheckResult {
            converged: false,
            max_relc: f64::INFINITY,
            worst_depth: -1,
            last_iter: None,
            n_depths: 0,
            chmax,
            error: Some("No valid iteration data found in fort.9".to_string()),
        };
    }

    // Safely find depth with maximum absolute change without unwrap panic on NaN
    let worst = match cur_rows.iter().max_by(|a, b| {
        a.maximum
            .abs()
            .partial_cmp(&b.maximum.abs())
            .unwrap_or(std::cmp::Ordering::Equal)
    }) {
        Some(row) => row,
        None => {
            return ConvCheckResult {
                converged: false,
                max_relc: f64::INFINITY,
                worst_depth: -1,
                last_iter,
                n_depths: 0,
                chmax,
                error: Some(
                    "No valid iteration rows found when calculating maximum change".to_string(),
                ),
            };
        }
    };

    let max_relc = worst.maximum.abs();
    let is_valid_num = max_relc.is_finite();

    ConvCheckResult {
        converged: is_valid_num && max_relc < chmax,
        max_relc,
        worst_depth: worst.depth,
        last_iter,
        n_depths: cur_rows.len(),
        chmax,
        error: if is_valid_num {
            None
        } else {
            Some("Convergence value is NaN or Inf".to_string())
        },
    }
}

/// Checks if an atmosphere file (.7) contains NaN lines (>10% NaN lines = invalid) using exact word boundary
///
/// 文件缺失时返回 `false`（语义：不存在 NaN 内容）。这与“含 NaN 导致无效”是不同语义；
/// 调用方需先自行确认文件存在性，不应将“缺失”与“含 NaN”混为一谈。
pub fn atmosphere_has_nan(path: &Path) -> bool {
    let file = match File::open(path) {
        Ok(f) => f,
        Err(_) => return false,
    };
    let reader = BufReader::new(file);
    let mut total_lines = 0;
    let mut nan_lines = 0;
    let nan_re = NAN_RE.get_or_init(|| Regex::new(r"(?i)\bnan\b").unwrap());

    for line in reader.lines().map_while(Result::ok) {
        total_lines += 1;
        if nan_re.is_match(&line) {
            nan_lines += 1;
        }
    }

    if total_lines == 0 {
        return true;
    }

    (nan_lines as f64) > (total_lines as f64 * 0.1)
}

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

    #[test]
    fn test_nan_check() {
        let dir = tempfile::tempdir().unwrap();
        let file_path = dir.path().join("test.7");
        std::fs::write(&file_path, "1 2 3\n4 5 6\n7 8 9\n").unwrap();
        assert!(!atmosphere_has_nan(&file_path));

        let nan_file_path = dir.path().join("nan.7");
        std::fs::write(&nan_file_path, "NaN 2 3\nNaN 5 6\n7 8 9\n").unwrap();
        assert!(atmosphere_has_nan(&nan_file_path));

        // Substring false positive test
        let banana_file_path = dir.path().join("banana.7");
        std::fs::write(&banana_file_path, "banana 2 3\nbanana 5 6\n7 8 9\n").unwrap();
        assert!(!atmosphere_has_nan(&banana_file_path));

        // Missing file returns false (absence != contains NaN)
        let missing_path = dir.path().join("missing.7");
        assert!(!atmosphere_has_nan(&missing_path));
    }
}