SpectraRust/src/tlusty/io/initia.rs

//! TLUSTY 初始化驱动程序。
//!
//! 重构自 TLUSTY `initia.f`。
//! 这是 TLUSTY 的主入口点，负责读取输入和初始化所有状态。
//!
//! # 功能
//!
//! - 读取基本输入参数（TEFF, GRAV, LTE 等）
//! - 设置频率网格和权重
//! - 初始化原子数据
//! - 读取和设置模型
//! - 配置各种物理参数
//!
//! # 重构策略
//!
//! 由于 INITIA 是一个大型驱动程序，大部分逻辑是调用其他子模块。
//! 本模块将纯计算部分提取为独立函数，便于测试。

use super::{Result, FortranReader, FortranWriter};
use crate::tlusty::state::constants::*;
// f2r_depends: CHANGE, CHCTAB, CORRWM, DMDER, DOPGAM, GOMINI, INIFRC, INIFRS, INIFRT, INPDIS, INPMOD, INTERP, IROSET, LEVSET, LINSET, LINSPL, LTEGR, LTEGRD, NSTOUT, NSTPAR, ODFHYS, ODFSET, OPADD0, OPAHST, QUIT, RAYINI, RDATA, RDATAX, READBF, RTEANG, SIGAVE, SRTFRQ, STATE, TABINI, TABINT, TRAINI

// ============================================================================
// 物理常数
// ============================================================================

/// h/k (Planck 常数 / Boltzmann 常数) [K·s]
const HK: f64 = 4.79927e-11;
/// h (Planck 常数) [erg·s]
const H: f64 = 6.62620e-27;
/// 电子电荷 [esu]
const ECH: f64 = 4.80298e-10;
/// 电子质量 [g]
const EMASS: f64 = 9.1091e-28;
/// Boltzmann 常数 [erg/K]
const BOLK: f64 = 1.38066e-16;
/// 氢原子质量 [g]
const HMASS: f64 = 1.6733e-24;
/// 单位转换常数 (用于 Klein-Nishina)
const XCON: f64 = 8.0935e-21;
/// Thomson 散射截面 [cm²]
const SIGE: f64 = 6.6524e-25;
/// π
const PI: f64 = std::f64::consts::PI;
/// 2π
const TWO: f64 = 2.0 * PI;
/// 1.0
const UN: f64 = 1.0;
/// 0.5
const HALF: f64 = 0.5;
/// Stefan-Boltzmann 常数 / 4
const SIG4P: f64 = 1.380835e-2;

// ============================================================================
// 辅助数据 - 统计权重
// ============================================================================

/// 统计权重数据数组（来自 DATA 语句）
const IGLE: [i32; 18] = [2, 1, 2, 1, 6, 9, 4, 9, 6, 1, 2, 1, 6, 9, 4, 9, 6, 1];
const IGMN: [i32; 25] = [
    2, 1, 2, 1, 6, 9, 4, 9, 6, 1, 2, 1, 6, 9, 4, 9, 6, 1,
    10, 21, 28, 25, 6, 7, 6
];
const IGFE: [i32; 26] = [
    2, 1, 2, 1, 6, 9, 4, 9, 6, 1, 2, 1, 6, 9, 4, 9, 6, 1,
    10, 21, 28, 25, 6, 25, 30, 25
];
const IGNI: [i32; 28] = [
    2, 1, 2, 1, 6, 9, 4, 9, 6, 1, 2, 1, 6, 9, 4, 9, 6, 1,
    10, 21, 28, 25, 6, 25, 28, 21, 10, 21
];

// ============================================================================
// 纯计算函数
// ============================================================================

/// 生成对数均匀频率网格。
///
/// 在频率范围 [frmin, frmax] 内生成 nfreq 个对数均匀分布的频率点，
/// 并计算相应的梯形积分权重。
///
/// # 参数
///
/// * `frmin` - 最小频率 (Hz)
/// * `frmax` - 最大频率 (Hz)
/// * `nfreq` - 频率点数
///
/// # 返回
///
/// (freq, w) 元组：
/// - `freq`: 频率数组（降序排列）
/// - `w`: 积分权重数组
pub fn generate_log_frequency_grid(
    frmin: f64,
    frmax: f64,
    nfreq: usize,
) -> (Vec<f64>, Vec<f64>) {
    if nfreq == 0 {
        return (Vec::new(), Vec::new());
    }

    if nfreq == 1 {
        // 单点特殊情况
        return (vec![frmin], vec![1.0]);
    }

    let log_frmin = frmin.ln();
    let log_frmax = frmax.ln();
    let delta = (log_frmax - log_frmin) / (nfreq - 1) as f64;

    // 生成对数均匀网格（升序）
    let mut freq_ascending: Vec<f64> = (0..nfreq)
        .map(|i| (log_frmin + delta * i as f64).exp())
        .collect();

    // 反转为降序（Fortran 原始行为）
    freq_ascending.reverse();
    let freq = freq_ascending;

    // 计算梯形积分权重
    let mut w = vec![0.0; nfreq];
    w[0] = 0.5 * (freq[0] - freq[1]);
    w[nfreq - 1] = 0.5 * (freq[nfreq - 2] - freq[nfreq - 1]);
    for ij in 1..nfreq - 1 {
        w[ij] = 0.5 * (freq[ij - 1] - freq[ij + 1]);
    }

    (freq, w)
}

/// 计算 Klein-Nishina 散射截面。
///
/// 根据 Rybicki & Lightman (1975) 的公式计算 Compton 散射截面。
/// 对于低能光子（xf << 1），使用泰勒展开；
/// 对于高能光子（xf >> 1），使用渐近公式。
///
/// # 参数
///
/// * `freq` - 频率 (Hz)
/// * `knish` - 是否使用完整的 Klein-Nishina 公式
///             0: 一阶近似
///             1: 完整公式
///
/// # 返回
///
/// 散射截面 (cm²)
pub fn klein_nishina_cross_section(freq: f64, knish: i32) -> f64 {
    if knish == 0 {
        // 一阶近似
        return SIGE * (UN - TWO * freq * XCON);
    }

    let xf = XCON * freq;

    if xf < 1e-1 {
        // 泰勒展开（低能极限）
        SIGE * (1.0 - xf * (2.0 - xf * (26.0 / 5.0 - xf * (13.3
            - xf * (1144.0 / 35.0 - xf * (544.0 / 7.0 - xf * (3784.0 / 21.0
            - xf * (6148.0 / 15.0 - xf * (151552.0 / 165.0
            - xf * 111872.0 / 55.0)))))))))
    } else if xf > 1e3 {
        // 渐近公式（高能极限）
        SIGE * 3.0 / 8.0 / xf * (2.0 * xf).ln_1p() + 0.5
    } else {
        // 完整 Klein-Nishina 公式
        SIGE * 0.75 * ((1.0 + xf) / xf.powi(3)
            * (2.0 * xf * (1.0 + xf) / (1.0 + 2.0 * xf)
               - (1.0 + 2.0 * xf).ln_1p())
            + 0.5 * (1.0 + 2.0 * xf).ln_1p() / xf
            - (1.0 + 3.0 * xf) / (1.0 + 2.0 * xf).powi(2))
    }
}

/// 计算 Planck 函数 B_ν(T)。
///
/// # 参数
///
/// * `freq` - 频率 (Hz)
/// * `temp` - 温度 (K)
///
/// # 返回
///
/// Planck 函数值 [erg/(s·cm²·Hz·sr)]
pub fn planck_function(freq: f64, temp: f64) -> f64 {
    let x = HK * freq / temp;
    if x > 700.0 {
        return 0.0; // 避免溢出
    }
    let hp = H * freq;
    let exp_x = x.exp();
    // CAS = 2.997925e18 Å/s = 2.997925e10 cm/s
    let cl = 2.997925e10;
    2.0 * hp * freq.powi(3) / (cl * cl) / (exp_x - UN)
}

/// 计算外部辐照强度。
///
/// # 参数
///
/// * `freq` - 频率数组 (Hz)
/// * `w` - 频率权重
/// * `trad` - 辐射温度 (K)，> 0 使用黑体，< 0 从文件读取，= 0 无辐照
/// * `wdil` - 稀释因子
///
/// # 返回
///
/// (extrad, extot) 元组：
/// - `extrad`: 各频率的外部辐射强度
/// - `extot`: 总外部辐射能量
pub fn compute_external_irradiation(
    freq: &[f64],
    w: &[f64],
    trad: f64,
    wdil: f64,
) -> (Vec<f64>, f64) {
    let nfreq = freq.len();
    let mut extrad = vec![0.0; nfreq];
    let mut extot = 0.0;

    if trad == 0.0 {
        // 无外部辐照
        return (extrad, extot);
    }

    if trad > 0.0 {
        // 使用黑体辐射
        for ij in 0..nfreq {
            let bnue = planck_function(freq[ij], trad);
            extrad[ij] = bnue / ((HK * freq[ij] / trad).exp() - UN) * wdil;
            extot += w[ij] * extrad[ij];
        }
    }

    (extrad, extot)
}

/// 初始化 1/i² 和 1/i³ 数组。
///
/// # 参数
///
/// * `nlmx` - 最大氢能级数
///
/// # 返回
///
/// (xi2, xi3) 元组
pub fn init_reciprocal_powers(nlmx: usize) -> (Vec<f64>, Vec<f64>) {
    let mut xi2 = vec![0.0; nlmx + 1];
    let mut xi3 = vec![0.0; nlmx + 1];

    for i in 1..=nlmx {
        let x = i as f64;
        xi2[i] = UN / (x * x);
        xi3[i] = xi2[i] / x;
    }

    (xi2, xi3)
}

/// 获取元素的统计权重。
///
/// # 参数
///
/// * `iatii` - 原子序数
/// * `izii` - 离子电荷
///
/// # 返回
///
/// 统计权重值
pub fn get_statistical_weight(iatii: i32, izii: i32) -> f64 {
    if iatii <= izii {
        return 1.0;
    }

    let idx = (iatii - izii) as usize;
    match iatii {
        1..=24 if idx <= 18 => IGLE[idx - 1] as f64,
        25 if idx <= 25 => IGMN[idx - 1] as f64,
        26 if idx <= 26 => IGFE[idx - 1] as f64,
        28 if idx <= 28 => IGNI[idx - 1] as f64,
        _ => 1.0,
    }
}

// ============================================================================
// 参数结构体
// ============================================================================

/// 初始化配置参数。
#[derive(Debug, Clone)]
pub struct InitiaConfig {
    /// 最大频率点数
    pub mfreq: usize,
    /// 最大深度点数
    pub mdepth: usize,
    /// 最大离子数
    pub mion: usize,
    /// 最大能级数
    pub mlevel: usize,
    /// 最大跃迁数
    pub mtrans: usize,
    /// 最大原子数
    pub matom: usize,
    /// 最大连续频率点数
    pub mfreqc: usize,
    /// 最大线性化频率数
    pub mfrex: usize,
    /// 最大线性化能级数
    pub mlvexp: usize,
    /// 最大总参数数
    pub mtot: usize,
}

impl Default for InitiaConfig {
    fn default() -> Self {
        Self {
            mfreq: MFREQ,
            mdepth: MDEPTH,
            mion: MION,
            mlevel: MLEVEL,
            mtrans: MTRANS,
            matom: MATOM,
            mfreqc: MFREQC,
            mfrex: MFREX,
            mlvexp: MLVEXP,
            mtot: MTOT,
        }
    }
}

/// 频率网格参数。
#[derive(Debug, Clone)]
pub struct FrequencyGridParams {
    /// 最小频率 (Hz)
    pub frmin: f64,
    /// 最大频率 (Hz)
    pub frmax: f64,
    /// 频率点数
    pub nfreq: usize,
    /// 是否使用预设频率
    pub ifrset: i32,
}

/// 频率网格输出。
#[derive(Debug, Clone)]
pub struct FrequencyGridOutput {
    /// 频率数组 (Hz)
    pub freq: Vec<f64>,
    /// 权重数组
    pub w: Vec<f64>,
    /// ALI 索引数组
    pub ijali: Vec<i32>,
}

/// INITIA 主参数结构体。
#[derive(Debug, Clone)]
pub struct InitiaParams {
    /// 配置
    pub config: InitiaConfig,
    /// 有效温度 (K)
    pub teff: f64,
    /// 表面重力 (cm/s², log10)
    pub grav: f64,
    /// 是否 LTE
    pub lte: bool,
    /// 是否灰大气
    pub ltgrey: bool,
    /// 湍流速度 (cm/s)
    pub vtb: f64,
    /// 是否处理湍流压力
    pub ipturb: i32,
    /// 深度点数
    pub nd: usize,
}

impl Default for InitiaParams {
    fn default() -> Self {
        Self {
            config: InitiaConfig::default(),
            teff: 10000.0,
            grav: 4.0,
            lte: false,
            ltgrey: false,
            vtb: 0.0,
            ipturb: 0,
            nd: 50,
        }
    }
}

/// INITIA 输出结构体。
#[derive(Debug, Clone)]
pub struct InitiaOutput {
    /// 频率网格
    pub freq_grid: FrequencyGridOutput,
    /// 1/i² 数组
    pub xi2: Vec<f64>,
    /// 1/i³ 数组
    pub xi3: Vec<f64>,
    /// 湍流速度数组
    pub vturb: Vec<f64>,
    /// Compton 散射截面
    pub sigec: Vec<f64>,
    /// 外部辐照强度
    pub extrad: Vec<f64>,
}

// ============================================================================
// 主初始化函数
// ============================================================================

/// 执行初始化的纯计算部分。
///
/// 这个函数实现 INITIA 中不涉及 I/O 的纯计算逻辑，
/// 包括频率网格生成、Klein-Nishina 截面计算等。
///
/// # 参数
///
/// * `params` - 初始化参数
/// * `grid_params` - 频率网格参数
/// * `icompt` - Compton 散射模式 (0: 不变截面, 1: 可变)
/// * `knish` - Klein-Nishina 模式 (0: 一阶近似, 1: 完整公式)
///
/// # 返回
///
/// 初始化输出结构体
pub fn initia_pure(
    params: &InitiaParams,
    grid_params: &FrequencyGridParams,
    icompt: i32,
    knish: i32,
) -> InitiaOutput {
    // 1. 生成频率网格
    let (freq, w) = generate_log_frequency_grid(
        grid_params.frmin,
        grid_params.frmax,
        grid_params.nfreq,
    );

    // 初始化 ALI 索引
    let ijali = vec![1; grid_params.nfreq];

    // 2. 初始化 1/i² 和 1/i³
    let (xi2, xi3) = init_reciprocal_powers(NLMX);

    // 3. 初始化湍流速度
    let mut vturb = vec![0.0; params.nd];
    let mut vturbs = vec![0.0; params.nd];

    let vtb_actual = if params.vtb.abs() < 1e3 {
        params.vtb * 1e5 // 从 km/s 转换为 cm/s
    } else {
        params.vtb
    };

    for id in 0..params.nd {
        if vtb_actual > 0.0 {
            vturb[id] = vtb_actual;
        }
        if params.ipturb == 0 {
            vturb[id] = 0.0;
        }
        vturbs[id] = vtb_actual.abs();
    }

    // 4. 计算 Compton 散射截面
    let mut sigec = vec![SIGE; grid_params.nfreq];
    if icompt != 0 {
        for ij in 0..grid_params.nfreq {
            sigec[ij] = klein_nishina_cross_section(freq[ij], knish);
        }
    }

    // 5. 初始化外部辐照
    let extrad = vec![0.0; grid_params.nfreq];

    InitiaOutput {
        freq_grid: FrequencyGridOutput {
            freq,
            w,
            ijali,
        },
        xi2,
        xi3,
        vturb,
        sigec,
        extrad,
    }
}

// ============================================================================
// 测试
// ============================================================================

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

    #[test]
    fn test_generate_log_frequency_grid() {
        // 测试单点
        let (freq, w) = generate_log_frequency_grid(1e14, 1e15, 1);
        assert_eq!(freq.len(), 1);
        assert_relative_eq!(freq[0], 1e14);
        assert_relative_eq!(w[0], 1.0);

        // 测试多点
        let (freq, w) = generate_log_frequency_grid(1e14, 1e15, 5);
        assert_eq!(freq.len(), 5);

        // 频率应该是降序的
        for i in 0..freq.len() - 1 {
            assert!(freq[i] > freq[i + 1]);
        }

        // 频率应该在范围内
        assert!(freq[0] <= 1e15);
        assert!(freq[freq.len() - 1] >= 1e14);

        // 权重总和检查（近似）
        let total_w: f64 = w.iter().sum();
        assert!(total_w > 0.0);
    }

    #[test]
    fn test_klein_nishina_cross_section() {
        // 低能极限：截面接近 Thomson 截面
        let low_energy = klein_nishina_cross_section(1e14, 1);
        assert_relative_eq!(low_energy / SIGE, 1.0, epsilon = 1e-4);

        // 一阶近似
        let first_order = klein_nishina_cross_section(1e18, 0);
        assert!(first_order < SIGE);

        // 高能极限：截面应该更小
        let high_energy = klein_nishina_cross_section(1e22, 1);
        assert!(high_energy < SIGE);
        assert!(high_energy > 0.0);
    }

    #[test]
    fn test_planck_function() {
        // 测试 Wien 极限（高频/低温）
        let bnue = planck_function(1e15, 5000.0);
        assert!(bnue > 0.0);

        // 测试 Rayleigh-Jeans 极限（低频/高温）
        let bnue_rj = planck_function(1e12, 50000.0);
        assert!(bnue_rj > 0.0);

        // 测试极高温（避免溢出）
        let bnue_hot = planck_function(1e15, 100000.0);
        assert!(bnue_hot > 0.0);
    }

    #[test]
    fn test_init_reciprocal_powers() {
        let (xi2, xi3) = init_reciprocal_powers(10);

        // 检查 i=1
        assert_relative_eq!(xi2[1], 1.0);
        assert_relative_eq!(xi3[1], 1.0);

        // 检查 i=2
        assert_relative_eq!(xi2[2], 0.25);
        assert_relative_eq!(xi3[2], 0.125);

        // 检查 i=10
        assert_relative_eq!(xi2[10], 0.01);
        assert_relative_eq!(xi3[10], 0.001);
    }

    #[test]
    fn test_get_statistical_weight() {
        // 测试 H I (Z=1, ion=0)
        let g_h1 = get_statistical_weight(1, 0);
        assert_eq!(g_h1, 2.0);

        // 测试 He I (Z=2, ion=0)
        let g_he1 = get_statistical_weight(2, 0);
        assert_eq!(g_he1, 1.0);

        // 测试 He II (Z=2, ion=1)
        let g_he2 = get_statistical_weight(2, 1);
        assert_eq!(g_he2, 2.0);

        // 测试完全电离（应该返回 1.0）
        let g_ionized = get_statistical_weight(1, 1);
        assert_eq!(g_ionized, 1.0);
    }

    #[test]
    fn test_compute_external_irradiation() {
        let freq = vec![1e14, 5e14, 1e15];
        let w = vec![0.5e14, 0.5e14, 0.5e14];

        // 无外部辐照
        let (extrad, extot) = compute_external_irradiation(&freq, &w, 0.0, 1.0);
        assert_eq!(extrad, vec![0.0, 0.0, 0.0]);
        assert_relative_eq!(extot, 0.0);

        // 有外部辐照
        let (extrad, extot) = compute_external_irradiation(&freq, &w, 10000.0, 0.5);
        assert!(extrad.iter().all(|&x| x >= 0.0));
        assert!(extot > 0.0);
    }

    #[test]
    fn test_initia_pure() {
        let config = InitiaConfig::default();
        let params = InitiaParams {
            config: config.clone(),
            teff: 35000.0,
            grav: 4.5,
            lte: false,
            ltgrey: false,
            vtb: 10.0, // 10 km/s
            ipturb: 0,
            nd: 50,
        };

        let grid_params = FrequencyGridParams {
            frmin: 1e14,
            frmax: 1e16,
            nfreq: 100,
            ifrset: 0,
        };

        let output = initia_pure(&params, &grid_params, 0, 0);

        assert_eq!(output.freq_grid.freq.len(), 100);
        assert_eq!(output.freq_grid.w.len(), 100);
        assert_eq!(output.freq_grid.ijali.len(), 100);
        assert_eq!(output.vturb.len(), 50);
        assert_eq!(output.sigec.len(), 100);

        // 检查 Compton 截面（不变模式）
        for &s in &output.sigec {
            assert_relative_eq!(s, SIGE);
        }
    }
}