SpectraRust/src/math/opfrac.rs

//! Opacity Project 电离分数表和配分函数计算。
//!
//! 重构自 TLUSTY `OPFRAC` 子程序。
//!
//! # 功能
//!
//! - 读取 Opacity Project 电离分数表
//! - 计算配分函数（通过插值）
//! - 计算电离分数

use crate::io::{Result, IoError};

// ============================================================================
// 常量
// ============================================================================

const MTEMP: usize = 100;
const MELEC: usize = 60;
const MSTAG: usize = 258;

/// ln(10)
const LN10: f64 = std::f64::consts::LN_10;

// ============================================================================
// IDAT 数组：数据文件索引映射
// ============================================================================

/// IDAT 数组：数据文件索引映射
/// 索引：原子序数 (1-28)
const IDAT: [i32; 30] = [
    1, 2, 0, 0, 0, 3, 4, 5, 0, 6,
    7, 8, 9, 10, 0, 11, 0, 12, 0, 13,
    0, 0, 0, 14, 15, 16, 0, 17, 0, 0
];

// ============================================================================
// PFOPTB COMMON 块
// ============================================================================

/// Opacity Project 配分函数和电离分数表。
/// 对应 COMMON /PFOPTB/
#[derive(Debug, Clone)]
pub struct PfOptB {
    /// 配分函数表 PFOP(MTEMP, MELEC, MSTAG)
    pub pfop: Vec<f64>,
    /// H- 配分函数 PFOPHM(MTEMP, MELEC)
    pub pfophm: Vec<f64>,
    /// 电离分数 FRAC(MTEMP, MELEC, MSTAG)
    pub frac: Vec<f64>,
    /// 电离分数 OP FROP(MTEMP, MELEC, MSTAG)
    pub frop: Vec<f64>,
    /// 温度索引 ITEMP(MTEMP)
    pub itemp: Vec<i32>,
    /// 温度点数
    pub ntt: i32,
}

impl Default for PfOptB {
    fn default() -> Self {
        Self {
            pfop: vec![0.0; MTEMP * MELEC * MSTAG],
            pfophm: vec![0.0; MTEMP * MELEC],
            frac: vec![0.0; MTEMP * MELEC * MSTAG],
            frop: vec![0.0; MTEMP * MELEC * MSTAG],
            itemp: vec![0; MTEMP],
            ntt: 0,
        }
    }
}

impl PfOptB {
    /// 创建新的 PfOptB 结构。
    pub fn new() -> Self {
        Self::default()
    }

    /// 获取 PFOP 值 (3D 数组访问)。
    pub fn get_pfop(&self, it: usize, ie: usize, indx: usize) -> f64 {
        self.pfop[it + MTEMP * ie + MTEMP * MELEC * indx]
    }

    /// 设置 PFOP 值。
    pub fn set_pfop(&mut self, it: usize, ie: usize, indx: usize, value: f64) {
        self.pfop[it + MTEMP * ie + MTEMP * MELEC * indx] = value;
    }

    /// 获取 FRAC 值。
    pub fn get_frac(&self, it: usize, ie: usize, indx: usize) -> f64 {
        self.frac[it + MTEMP * ie + MTEMP * MELEC * indx]
    }

    /// 设置 FRAC 值。
    pub fn set_frac(&mut self, it: usize, ie: usize, indx: usize, value: f64) {
        self.frac[it + MTEMP * ie + MTEMP * MELEC * indx] = value;
    }

    /// 获取 PFOPHM 值。
    pub fn get_pfophm(&self, it: usize, ie: usize) -> f64 {
        self.pfophm[it + MTEMP * ie]
    }

    /// 设置 PFOPHM 值。
    pub fn set_pfophm(&mut self, it: usize, ie: usize, value: f64) {
        self.pfophm[it + MTEMP * ie] = value;
    }
}

// ============================================================================
// 输入参数结构体
// ============================================================================

/// OPFRAC 输入参数（计算模式）。
pub struct OpfracParams {
    /// 原子序数 (1-based, 0 表示初始化)
    pub iat: i32,
    /// 离子序数 (1-based)
    pub ion: i32,
    /// 温度 (K)
    pub t: f64,
    /// 电子密度 (cm⁻³)
    pub ane: f64,
}

/// OPFRAC 输出结果。
#[derive(Debug, Clone)]
pub struct OpfracOutput {
    /// 配分函数
    pub pf: f64,
    /// 电离分数
    pub fra: f64,
}

// ============================================================================
// 核心计算函数
// ============================================================================

/// 执行 OPFRAC 计算（插值模式）。
///
/// # 参数
/// - `params`: 输入参数
/// - `pfoptb`: 预计算的配分函数表
///
/// # 返回
/// 配分函数和电离分数
pub fn opfrac_pure(params: &OpfracParams, pfoptb: &PfOptB) -> OpfracOutput {
    let iat = params.iat;

    // 如果 IAT == 0，返回默认值
    if iat == 0 {
        return OpfracOutput { pf: 1.0, fra: 1.0 };
    }

    // 如果 IAT > 28 或 IDAT[iat-1] == 0，返回默认值
    if iat > 28 || IDAT[iat as usize - 1] == 0 {
        return OpfracOutput { pf: 1.0, fra: 1.0 };
    }

    let xt = params.t.log10();
    let xne = params.ane.log10();

    // 温度索引
    let kt0 = (40.0 * xt) as i32;
    let kn0 = (2.0 * xne) as i32;

    // 温度索引查找
    let kt1 = if pfoptb.ntt == 0 {
        return OpfracOutput { pf: 1.0, fra: 1.0 };
    } else if kt0 < pfoptb.itemp[0] {
        0
    } else if kt0 >= pfoptb.itemp[pfoptb.ntt as usize - 1] {
        (pfoptb.ntt - 2).max(0) as usize
    } else {
        // 查找温度索引
        let mut found = 0usize;
        for it in 0..pfoptb.ntt as usize {
            if kt0 == pfoptb.itemp[it] {
                found = it;
                break;
            } else if kt0 < pfoptb.itemp[it] && it > 0 {
                found = it - 1;
                break;
            }
        }
        found
    };

    // 电子密度索引
    let kn1 = if kn0 < 1 { 0 } else if kn0 >= 59 { 58 } else { kn0 as usize };

    // 检查索引有效性
    if kt1 + 1 >= MTEMP || kn1 + 1 >= MELEC {
        return OpfracOutput { pf: 1.0, fra: 1.0 };
    }

    // 获取 INDXAT 索引（简化版，使用原子和离子的直接映射）
    let indx = get_indxat(params.iat as usize, params.ion as usize);
    if indx == 0 {
        return OpfracOutput { pf: 1.0, fra: 1.0 };
    }
    let indx = (indx - 1) as usize;

    // 插值系数
    let xt1 = 0.025 * pfoptb.itemp[kt1] as f64;
    let dxt = 0.05;
    let at1 = (xt - xt1) / dxt;
    let xn1 = 0.5 * kn1 as f64;
    let dxn = 0.5;
    let an1 = (xne - xn1) / dxn;

    // 双线性插值
    let x11 = pfoptb.get_pfop(kt1, kn1, indx);
    let x21 = pfoptb.get_pfop(kt1 + 1, kn1, indx);
    let x12 = pfoptb.get_pfop(kt1, kn1 + 1, indx);
    let x22 = pfoptb.get_pfop(kt1 + 1, kn1 + 1, indx);

    let x1221 = x11 * x21 * x12 * x22;
    let pf = if x1221 == 0.0 {
        // 线性插值
        let xx1 = x11 + at1 * (x21 - x11);
        let xx2 = x12 + at1 * (x22 - x12);
        xx1 + an1 * (xx2 - xx1)
    } else {
        // 对数空间插值
        let x11_log = x11.log10();
        let x21_log = x21.log10();
        let x12_log = x12.log10();
        let x22_log = x22.log10();
        let xx1 = x11_log + at1 * (x21_log - x11_log);
        let xx2 = x12_log + at1 * (x22_log - x12_log);
        let rrx = xx1 + an1 * (xx2 - xx1);
        (rrx * LN10).exp()
    };

    // 电离分数
    let fra = pfoptb.get_frac(kt1, kn1, indx);

    OpfracOutput { pf, fra }
}

/// 获取 INDXAT 索引。
///
/// 这个函数将原子序数和离子序数映射到 OP 表格中的索引。
fn get_indxat(iat: usize, ion: usize) -> i32 {
    // 简化版映射：基于原子序数和离子序数
    // 完整映射需要完整的 INDXAT 数组

    // H (iat=1): 索引 1-2
    // He (iat=2): 索引 3-5
    // 等等...

    let base = match iat {
        1 => 1,    // H: 索引 1-2
        2 => 3,    // He: 索引 3-5
        3 => 6,    // Li
        4 => 13,   // Be
        5 => 21,   // B
        6 => 30,   // C
        7 => 41,   // N
        8 => 53,   // O
        9 => 66,   // F
        10 => 80,  // Ne
        11 => 95,  // Na
        12 => 112, // Mg
        13 => 131, // Al
        14 => 152, // Si
        15 => 177, // P
        16 => 203, // S
        17 => 230, // Cl
        18 => 258, // Ar
        _ => return 0,
    };

    let offset = ion.saturating_sub(1) as i32;
    let idx = base + offset;

    if idx > 0 && idx <= 258 {
        idx
    } else {
        0
    }
}

// ============================================================================
// 初始化函数（带 I/O）
// ============================================================================

/// 读取电离分数表并初始化配分函数表。
///
/// # 参数
/// - `file_path`: ioniz.dat 文件路径
///
/// # 返回
/// 初始化的 PfOptB 结构
pub fn opfrac_init(_file_path: &str) -> Result<PfOptB> {
    // TODO: 完整实现需要解析 ioniz.dat 文件
    // 当前返回默认结构
    Ok(PfOptB::new())
}

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

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

    #[test]
    fn test_pfoptb_creation() {
        let pfoptb = PfOptB::new();
        assert_eq!(pfoptb.pfop.len(), MTEMP * MELEC * MSTAG);
        assert_eq!(pfoptb.pfophm.len(), MTEMP * MELEC);
        assert_eq!(pfoptb.frac.len(), MTEMP * MELEC * MSTAG);
    }

    #[test]
    fn test_pfoptb_accessors() {
        let mut pfoptb = PfOptB::new();

        // 测试 PFOP 访问
        pfoptb.set_pfop(0, 0, 0, 1.5);
        assert!((pfoptb.get_pfop(0, 0, 0) - 1.5).abs() < 1e-10);

        // 测试 FRAC 访问
        pfoptb.set_frac(1, 2, 3, 2.5);
        assert!((pfoptb.get_frac(1, 2, 3) - 2.5).abs() < 1e-10);

        // 测试 PFOPHM 访问
        pfoptb.set_pfophm(0, 0, 3.5);
        assert!((pfoptb.get_pfophm(0, 0) - 3.5).abs() < 1e-10);
    }

    #[test]
    fn test_opfrac_pure_zero_iat() {
        let pfoptb = PfOptB::new();

        let params = OpfracParams {
            iat: 0,
            ion: 0,
            t: 10000.0,
            ane: 1.0e12,
        };

        let result = opfrac_pure(&params, &pfoptb);
        assert!((result.pf - 1.0).abs() < 1e-10);
        assert!((result.fra - 1.0).abs() < 1e-10);
    }

    #[test]
    fn test_opfrac_pure_with_data() {
        let mut pfoptb = PfOptB::new();

        // 设置一些测试数据
        pfoptb.ntt = 10;
        for i in 0..10 {
            pfoptb.itemp[i] = (i as i32 + 1) * 4;  // 温度索引
        }

        // 设置配分函数值
        let test_pf = 2.0;
        for it in 0..10 {
            for ie in 0..MELEC {
                pfoptb.set_pfop(it, ie, 0, test_pf);
                pfoptb.set_frac(it, ie, 0, 0.5);
            }
        }

        let params = OpfracParams {
            iat: 1,  // H
            ion: 1,  // H I
            t: 10000.0,
            ane: 1.0e12,
        };

        let result = opfrac_pure(&params, &pfoptb);

        assert!(result.pf > 0.0, "PF should be positive");
    }

    #[test]
    fn test_idat_array() {
        // 验证 IDAT 数组
        assert_eq!(IDAT[0], 1);   // H
        assert_eq!(IDAT[1], 2);   // He
        assert_eq!(IDAT[5], 3);   // C
        assert_eq!(IDAT[6], 4);   // N
        assert_eq!(IDAT[7], 5);   // O
    }

    #[test]
    fn test_get_indxat() {
        // H I
        assert_eq!(get_indxat(1, 1), 1);
        // He I
        assert!(get_indxat(2, 1) > 0);
        // 未知原子
        assert_eq!(get_indxat(30, 1), 0);
    }
}