SpectraRust/src/synspec/math/hydini.rs

//! Hydrogen line profile data initialization.
//!
//! Translated from SYNSPEC54.FOR subroutine HYDINI (line 6877).
//!
//! Initializes necessary arrays for evaluating hydrogen line profiles
//! from the Lemke, Tremblay-Bergeron, or Schoening-Butler tables.

#![allow(clippy::never_loop)]
use std::fs::File;
use std::io::{BufRead, BufReader};

use super::stark0::stark0;

/// Constants for hydrogen profile arrays
pub const NLINES_MAX: usize = 22;
pub const NLEVELS: usize = 4;
pub const NWL_MAX: usize = 100;
pub const NT_MAX: usize = 20;
pub const NE_MAX: usize = 20;

/// Hydrogen line profile table data
#[derive(Debug, Clone)]
pub struct HydProfileTable {
    /// Line index (i, j)
    pub i: usize,
    pub j: usize,
    /// Central wavelength
    pub wl0: f64,
    /// Number of wavelength points
    pub nwl: usize,
    /// Number of temperature points
    pub nt: usize,
    /// Number of electron density points
    pub ne: usize,
    /// Log10 wavelength displacements [NWL_MAX]
    pub wl: [f64; NWL_MAX],
    /// Log10 temperature grid [NT_MAX]
    pub xt: [f64; NT_MAX],
    /// Log10 electron density grid [NE_MAX]
    pub xne: [f64; NE_MAX],
    /// Profile values [NWL_MAX x NT_MAX x NE_MAX]
    pub prf: [[[f64; NE_MAX]; NT_MAX]; NWL_MAX],
    /// Asymptotic profile coefficient
    pub xk: f64,
}

impl Default for HydProfileTable {
    fn default() -> Self {
        Self {
            i: 0,
            j: 0,
            wl0: 0.0,
            nwl: 0,
            nt: 0,
            ne: 0,
            wl: [0.0; NWL_MAX],
            xt: [0.0; NT_MAX],
            xne: [0.0; NE_MAX],
            prf: [[[0.0; NE_MAX]; NT_MAX]; NWL_MAX],
            xk: 0.0,
        }
    }
}

/// Hydrogen line initialization result
#[derive(Debug, Clone)]
pub struct HydInitResult {
    /// Central wavelengths for lines [NLEVELS x NLINES_MAX]
    pub wline: [[f64; NLINES_MAX]; NLEVELS],
    /// Line index mapping [NLEVELS x NLINES_MAX]
    pub ilin0: [[usize; NLINES_MAX]; NLEVELS],
    /// Profile tables
    pub tables: Vec<HydProfileTable>,
    /// Lemke mode flag
    pub ilemke: bool,
    /// Number of lines
    pub nlihyd: usize,
}

impl Default for HydInitResult {
    fn default() -> Self {
        Self {
            wline: [[0.0; NLINES_MAX]; NLEVELS],
            ilin0: [[0; NLINES_MAX]; NLEVELS],
            tables: Vec::new(),
            ilemke: false,
            nlihyd: 0,
        }
    }
}

/// Hydrogen line profile table source
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum HydTableSource {
    /// Schoening-Butler tables (ihydpr < 0)
    SchoeningButler,
    /// Lemke tables (ihydpr = 21)
    Lemke,
    /// Tremblay-Bergeron tables (ihydpr = 22)
    Tremblay,
}

/// Parameters for HYDINI
pub struct HydiniParams {
    /// Table source selection
    pub source: HydTableSource,
    /// Path to data directory
    pub data_dir: String,
    /// Model depth points
    pub nd: usize,
    /// Temperature array [nd]
    pub temp: Vec<f64>,
    /// Electron density array [nd]
    pub elec: Vec<f64>,
    /// Turbulent velocity array [nd]
    pub vturb: Vec<f64>,
}

/// Initialize hydrogen line profile data.
///
/// # Arguments
/// * `params` - Initialization parameters
///
/// # Returns
/// Hydrogen line initialization result with profile tables
pub fn hydini(params: &HydiniParams) -> std::io::Result<HydInitResult> {
    let mut result = HydInitResult::default();

    // Initialize central wavelengths using STARK0
    for i in 0..NLEVELS {
        for j in (i + 1)..NLINES_MAX {
            let stark = stark0(i as i32 + 1, j as i32 + 1, 1);
            result.wline[i][j] = stark.wl0;
        }
    }

    // Initialize line index mapping
    for i in 0..NLEVELS {
        for j in 0..NLINES_MAX {
            result.ilin0[i][j] = 0;
        }
    }

    match params.source {
        HydTableSource::SchoeningButler => {
            read_schoening_butler(params, &mut result)?;
        }
        HydTableSource::Lemke | HydTableSource::Tremblay => {
            read_lemke_tremblay(params, &mut result)?;
        }
    }

    Ok(result)
}

/// Read Schoening-Butler tables
fn read_schoening_butler(
    params: &HydiniParams,
    result: &mut HydInitResult,
) -> std::io::Result<()> {
    let filename = format!("{}/hydprf.dat", params.data_dir);
    let file = File::open(&filename)?;
    let reader = BufReader::new(file);
    let mut lines = reader.lines();

    // Skip 12 header lines
    for _ in 0..12 {
        lines.next();
    }

    let nline = 12;
    result.ilemke = false;

    for iline in 0..nline {
        // Read line indices
        let header = read_next_line(&mut lines)?;
        let (i, j) = parse_line_indices(&header)?;

        let j = if iline == 11 { 10 } else { j }; // Special case for last line
        let wl0 = result.wline[i - 1][j - 1];
        result.ilin0[i - 1][j - 1] = iline + 1;

        let mut table = HydProfileTable {
            i,
            j,
            wl0,
            ..Default::default()
        };

        // Read wavelength points
        let wl_line = read_next_line(&mut lines)?;
        let wl_parts = parse_data_line(&wl_line)?;
        let nwl = wl_parts.len() - 1; // First value is character
        table.nwl = nwl;
        for k in 0..nwl.min(NWL_MAX) {
            table.wl[k] = if wl_parts[k + 1] < 1.0e-4 {
                (1.0e-4_f64).log10()
            } else {
                wl_parts[k + 1].log10()
            };
        }

        // Read temperature points
        let xt_line = read_next_line(&mut lines)?;
        let xt_parts = parse_data_line(&xt_line)?;
        let nt = xt_parts.len() - 1;
        table.nt = nt;
        for k in 0..nt.min(NT_MAX) {
            table.xt[k] = xt_parts[k + 1];
        }

        // Read electron density points
        let xne_line = read_next_line(&mut lines)?;
        let xne_parts = parse_data_line(&xne_line)?;
        let ne = xne_parts.len() - 1;
        table.ne = ne;
        for k in 0..ne.min(NE_MAX) {
            table.xne[k] = xne_parts[k + 1];
        }

        // Skip blank line
        lines.next();

        // Read profile data
        for ie in 0..ne.min(NE_MAX) {
            for it in 0..nt.min(NT_MAX) {
                lines.next(); // Skip blank line
                let prf_line = read_next_line(&mut lines)?;
                let prf_parts = parse_data_line(&prf_line)?;
                for iwl in 0..nwl.min(NWL_MAX) {
                    if iwl < prf_parts.len() {
                        table.prf[iwl][it][ie] = prf_parts[iwl];
                    }
                }
            }
        }

        // Compute asymptotic profile coefficient
        if nwl > 0 && ne > 0 {
            let xclog = table.prf[nwl - 1][0][0]
                + 2.5 * table.wl[nwl - 1]
                + 31.5304
                - table.xne[0]
                - 2.0 * wl0.log10();
            let xklog = 0.6666667 * (xclog - 0.176);
            table.xk = (xklog * std::f64::consts::LN_10).exp();
        }

        result.tables.push(table);
    }

    Ok(())
}

/// Read Lemke or Tremblay tables
fn read_lemke_tremblay(
    params: &HydiniParams,
    result: &mut HydInitResult,
) -> std::io::Result<()> {
    let filename = match params.source {
        HydTableSource::Lemke => format!("{}/lemke.dat", params.data_dir),
        HydTableSource::Tremblay => format!("{}/tremblay.dat", params.data_dir),
        _ => unreachable!(),
    };

    let file = File::open(&filename)?;
    let reader = BufReader::new(file);
    let mut lines = reader.lines();

    result.ilemke = true;

    // Read number of tables
    let ntab_line = read_next_line(&mut lines)?;
    let ntab: usize = ntab_line.trim().parse().map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("NTAB: {}", e))
    })?;

    let mut iline = 0;

    for _ in 0..ntab {
        // Read number of lines in this table
        let nlly_line = read_next_line(&mut lines)?;
        let nlly: usize = nlly_line.trim().parse().map_err(|e| {
            std::io::Error::new(std::io::ErrorKind::InvalidData, format!("NLLY: {}", e))
        })?;

        let ilineb = iline;

        // Read line parameters
        for _ in 0..nlly {
            let param_line = read_next_line(&mut lines)?;
            let parts = parse_data_line(&param_line)?;

            if parts.len() < 11 {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidData,
                    "Invalid Lemke/Tremblay parameter line",
                ));
            }

            let i = parts[0] as usize;
            let j = parts[1] as usize;
            let almin = parts[2];
            let anemin = parts[3];
            let tmin = parts[4];
            let dla = parts[5];
            let dle = parts[6];
            let dlt = parts[7];
            let nwl = parts[8] as usize;
            let ne = parts[9] as usize;
            let nt = parts[10] as usize;

            let wl0 = result.wline[i - 1][j - 1];
            result.ilin0[i - 1][j - 1] = iline + 1;

            let mut table = HydProfileTable {
                i,
                j,
                wl0,
                nwl,
                nt,
                ne,
                ..Default::default()
            };

            // Generate wavelength grid
            for iwl in 0..nwl.min(NWL_MAX) {
                table.wl[iwl] = almin + (iwl as f64) * dla;
            }

            // Generate electron density grid
            for ie in 0..ne.min(NE_MAX) {
                table.xne[ie] = anemin + (ie as f64) * dle;
            }

            // Generate temperature grid
            for it in 0..nt.min(NT_MAX) {
                table.xt[it] = tmin + (it as f64) * dlt;
            }

            result.tables.push(table);
            iline += 1;
        }

        // Read profile data for each line
        for ili in 0..nlly {
            let ilne = ilineb + ili;
            let table = &mut result.tables[ilne];
            let nwl = table.nwl;
            let ne = table.ne;
            let nt = table.nt;

            lines.next(); // Skip blank line

            for ie in 0..ne.min(NE_MAX) {
                for it in 0..nt.min(NT_MAX) {
                    let prf_line = read_next_line(&mut lines)?;
                    let parts = parse_data_line(&prf_line)?;

                    // First value is QLT (quality factor), skip it
                    for iwl in 0..nwl.min(NWL_MAX) {
                        if iwl + 1 < parts.len() {
                            table.prf[iwl][it][ie] = parts[iwl + 1];
                        }
                    }
                }
            }

            // Compute asymptotic profile coefficient
            if nwl > 0 && ne > 0 {
                let xclog = table.prf[nwl - 1][0][0]
                    + 2.5 * table.wl[nwl - 1].log10()
                    + 31.5304
                    - table.xne[0]
                    - 2.0 * table.wl0.log10();
                let xklog = 0.6666667 * (xclog - 0.176);
                table.xk = (xklog * std::f64::consts::LN_10).exp();
            }
        }
    }

    result.nlihyd = iline;

    Ok(())
}

/// Read next non-empty line
fn read_next_line(lines: &mut impl Iterator<Item = std::io::Result<String>>) -> std::io::Result<String> {
    loop {
        match lines.next() {
            Some(Ok(line)) => return Ok(line),
            Some(Err(e)) => return Err(e),
            None => return Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "Unexpected end of file",
            )),
        }
    }
}

/// Parse line indices from header: FORMAT(12X,I1,9X,I1)
fn parse_line_indices(line: &str) -> std::io::Result<(usize, usize)> {
    let parts: Vec<&str> = line.split_whitespace().collect();
    if parts.len() < 2 {
        return Err(std::io::Error::new(
            std::io::ErrorKind::InvalidData,
            format!("Invalid line indices: {}", line),
        ));
    }

    let i = parts[0].parse::<usize>().map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("I: {}", e))
    })?;
    let j = parts[1].parse::<usize>().map_err(|e| {
        std::io::Error::new(std::io::ErrorKind::InvalidData, format!("J: {}", e))
    })?;

    Ok((i, j))
}

/// Parse data line (free format)
fn parse_data_line(line: &str) -> std::io::Result<Vec<f64>> {
    let values: Vec<f64> = line
        .split_whitespace()
        .filter_map(|s| s.parse::<f64>().ok())
        .collect();

    Ok(values)
}

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

    #[test]
    fn test_hydini_default() {
        let result = HydInitResult::default();
        assert_eq!(result.wline.len(), NLEVELS);
        assert_eq!(result.ilin0.len(), NLEVELS);
        assert!(result.tables.is_empty());
    }

    #[test]
    fn test_hyd_profile_table_default() {
        let table = HydProfileTable::default();
        assert_eq!(table.nwl, 0);
        assert_eq!(table.nt, 0);
        assert_eq!(table.ne, 0);
    }

    #[test]
    fn test_parse_line_indices() {
        let line = "  1         2";
        let result = parse_line_indices(line);
        assert!(result.is_ok());
        let (i, j) = result.unwrap();
        assert_eq!(i, 1);
        assert_eq!(j, 2);
    }

    #[test]
    fn test_parse_data_line() {
        let line = "  1.0  2.0  3.0  4.0";
        let result = parse_data_line(line);
        assert!(result.is_ok());
        let values = result.unwrap();
        assert_eq!(values.len(), 4);
        assert!((values[0] - 1.0).abs() < 1e-10);
    }
}