SpectraRust/src/tlusty/math/radiative/radpre.rs

//! 辐射加速度计算 - RADPRE。
//!
//! 重构自 TLUSTY `radpre.f`
//!
//! 计算辐射加速度，自动排除对总辐射压力贡献最强的线。
//! 使用深度相关准则进行频率筛选。

use crate::tlusty::state::constants::{BOLK, HALF, MDEPTH, MFREQ, MLEVEL, MTRANS};
use crate::tlusty::math::indexx;
use crate::tlusty::math::quit;
// f2r_depends: OPACF1, RTEFR1

// ============================================================================
// 常量定义
// ============================================================================

/// PGRD 常量: 4.1916825e-10
const PGRD: f64 = 4.1916825e-10;

// ============================================================================
// XGRD 预设数组
// ============================================================================

/// XGRD0: 10 元素预设数组（用于 XGRAD = 0）
const XGRD0: [f64; 10] = [
    0.1, 0.3, 0.5, 0.7, 0.9, 0.92, 0.94, 0.96, 0.98, 0.99,
];

/// XGRD1: 20 元素预设数组（用于 XGRAD = -1）
const XGRD1: [f64; 20] = [
    0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.65, 0.7, 0.75, 0.8,
    0.85, 0.9, 0.92, 0.94, 0.96, 0.98, 0.99, 0.99, 0.99, 0.99,
];

/// XGRD2: 20 元素预设数组（用于 XGRAD = -2）
const XGRD2: [f64; 20] = [
    0.1, 0.2, 0.3, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7,
    0.75, 0.8, 0.84, 0.87, 0.9, 0.93, 0.95, 0.97, 0.98, 0.99,
];

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

/// RADPRE 配置参数。
#[derive(Debug, Clone)]
pub struct RadpreConfig {
    /// XGRAD 参数（控制频率筛选）
    /// = 0: 使用 XGRD0 预设
    /// = -1: 使用 XGRD1 预设
    /// = -2: 使用 XGRD2 预设
    /// > 0: 使用固定值
    pub xgrad: f64,
    /// 重力加速度 (cm/s²)
    pub grav: f64,
    /// 辐射压力标志 (0: 不计算, >0: 计算)
    pub ifprad: i32,
    /// ODF 采样标志
    pub ispodf: i32,
    /// 最大显式频率数
    pub mfrex: usize,
}

impl Default for RadpreConfig {
    fn default() -> Self {
        Self {
            xgrad: 0.0,
            grav: 0.0,
            ifprad: 0,
            ispodf: 0,
            mfrex: 100,
        }
    }
}

/// RADPRE 模型状态参数。
pub struct RadpreModelState<'a> {
    /// 深度点数
    pub nd: usize,
    /// 温度 [nd]
    pub temp: &'a [f64],
    /// 电子密度 [nd]
    pub elec: &'a [f64],
    /// 总粒子密度 [nd]
    pub dens: &'a [f64],
    /// 柱质量 [nd]
    pub dm: &'a [f64],
    /// 平均分子量倒数 [nd]
    pub wmm: &'a [f64],
    /// 湍流速度 [nd]
    pub vturb: &'a [f64],
    /// 深度间隔 [nd-1]
    pub deldm: &'a [f64],
    /// 柱质量梯度 [nd]
    pub dedm1: f64,
    /// Roseland 不透明度 [nd]
    pub abrosd: &'a [f64],
    /// 总吸收系数 [nd]
    pub absot: &'a [f64],
    /// 密度 [nd] (用于 dens1)
    pub dens1: &'a [f64],
}

/// RADPRE 频率相关参数（可变）。
pub struct RadpreFreqParamsMut<'a> {
    /// 频率总数
    pub nfreq: usize,
    /// 频率 [nfreq]
    pub freq: &'a [f64],
    /// 频率权重 [nfreq]
    pub w: &'a [f64],
    /// 主谱线索引 [nfreq], 0 表示无
    pub ijlin: &'a [i32],
    /// 每个频率的线数 [nfreq]
    pub nlines: &'a [i32],
    /// 跃迁中心频率 [mtrans]
    pub fr0: &'a [f64],
    /// 跃迁显式频率索引 [mtrans]
    pub indexp: &'a mut [i32],
    /// 跃迁显式频率标志 [mtrans]
    pub lexp: &'a mut [bool],
}

/// RADPRE ALI 相关参数（可变）。
pub struct RadpreAliParamsMut<'a> {
    /// ALI 索引 [nfreq]
    pub ijali: &'a mut [i32],
    /// 显式频率索引 [nfreq]
    pub ijex: &'a mut [i32],
    /// 显式频率对应的原始频率索引 [mfrex]
    pub ijfr: &'a mut [i32],
    /// 显式频率标志 [nfreq]
    pub ijx: &'a mut [i32],
    /// 连续性权重 [nfreq]
    pub wc: &'a mut [f64],
    /// 当前显式频率计数
    pub nfreqe: &'a mut i32,
    /// 跃迁线索引 [max_lines][nfreq]
    pub itrlin: &'a [Vec<i32>],
}

/// RADPRE 辐射场参数（由 OPACF1 和 RTEFR1 计算）。
pub struct RadpreRadField<'a> {
    /// 辐射强度 [nd]
    pub rad1: &'a [f64],
    /// Eddington 因子 [nd]
    pub fak1: &'a [f64],
    /// 吸收系数 [nd]
    pub abso1: &'a [f64],
    /// 表面 Eddington 因子 [nfreq]
    pub fh: &'a [f64],
    /// 外部辐射 [nfreq]
    pub hextrd: &'a [f64],
}

/// RADPRE 输出状态（可变）。
pub struct RadpreOutputStateMut<'a> {
    /// 跳过频率标志 [nd][nfreq]
    pub lskip: &'a mut [Vec<bool>],
    /// 辐射压力 [nd]
    pub pradt: &'a mut [f64],
    /// 累积辐射压力 [nd]
    pub prada: &'a mut [f64],
    /// 频率相关辐射加速度 [nd][nfreq]
    pub gradf: &'a mut [Vec<f64>],
}

/// RADPRE 输出结果。
#[derive(Debug, Clone)]
pub struct RadpreOutput {
    /// 深度相关阈值 [nd]
    pub xgrd: Vec<f64>,
    /// 气体+湍流压力加速度 [nd]
    pub ggrt: Vec<f64>,
    /// 辐射加速度 [nd]
    pub grad: Vec<f64>,
    /// 累积辐射加速度 [nd]
    pub grada: Vec<f64>,
    /// 表面辐射压力
    pub prd0: f64,
    /// 新增显式频率数
    pub nfe: i32,
}

// ============================================================================
// 主函数
// ============================================================================

/// 计算辐射加速度（RADPRE 主函数）。
///
/// # 参数
/// - `config`: 配置参数
/// - `model`: 模型状态
/// - `freq`: 频率参数（可变）
/// - `ali`: ALI 参数（可变）
/// - `output`: 输出状态（可变）
/// - `nn`: 跃迁计数（会被修改）
///
/// # 返回值
/// RADPRE 输出结果
#[allow(clippy::too_many_arguments)]
pub fn radpre(
    config: &RadpreConfig,
    model: &RadpreModelState,
    freq: &mut RadpreFreqParamsMut,
    ali: &mut RadpreAliParamsMut,
    output: &mut RadpreOutputStateMut,
    nn: &mut i32,
) -> RadpreOutput {
    let nd = model.nd;
    let nfreq = freq.nfreq;

    // 输出数组初始化
    let mut xgrd = vec![0.0; nd];
    let mut ggrt = vec![0.0; nd];
    let mut grad = vec![0.0; nd];
    let mut grada = vec![0.0; nd];
    let mut prid = vec![0.0; nd];
    let mut pgt = vec![0.0; nd];

    // 工作数组
    let mut gradi = vec![0.0; nfreq];

    // ========================================================================
    // 步骤 1: 设置深度相关阈值 XGRD
    // ========================================================================
    setup_xgrd(config.xgrad, nd, &mut xgrd);

    // ========================================================================
    // 步骤 2: 计算气体和湍流压力的加速度
    // ========================================================================
    // PGAS = (DENS/WMM + ELEC) * BOLK * TEMP
    for id in 0..nd {
        let pgas = (model.dens[id] / model.wmm[id] + model.elec[id]) * BOLK * model.temp[id];
        pgt[id] = pgas + HALF * model.dens[id] * model.vturb[id] * model.vturb[id];
    }

    // GGRT(ID) = (PGT(ID) - PGT(ID-1)) / (DM(ID) - DM(ID-1))
    ggrt[0] = 0.0; // 会在下面设置为 ggrt[1]
    for id in 1..nd {
        let dm_diff = model.dm[id] - model.dm[id - 1];
        if dm_diff.abs() > 1e-30 {
            ggrt[id] = (pgt[id] - pgt[id - 1]) / dm_diff;
        } else {
            ggrt[id] = 0.0;
        }
    }
    ggrt[0] = ggrt[1];

    // ========================================================================
    // 步骤 3: 初始化辐射相关数组
    // ========================================================================
    for id in 0..nd {
        grad[id] = 0.0;
        grada[id] = 0.0;
        output.pradt[id] = 0.0;
    }

    // PRID(ID) = PGRD / (DM(ID) - DM(ID-1))
    for id in 1..nd {
        let dm_diff = model.dm[id] - model.dm[id - 1];
        if dm_diff.abs() > 1e-30 {
            prid[id] = PGRD / dm_diff;
        } else {
            prid[id] = 0.0;
        }
    }

    let pgrd1 = PGRD / model.dens1[0];
    let mut prd0 = 0.0;

    // ========================================================================
    // 步骤 4: 遍历所有频率，计算辐射加速度
    // ========================================================================
    // 注意：实际的 OPACF1 和 RTEFR1 调用需要在外部完成
    // 这里只是框架，grada 需要在外部通过 radpre_accumulate_frequency 累积

    // ========================================================================
    // 步骤 5: 深度相关的频率拒绝
    // ========================================================================
    let mut nfe = 0;

    // 累积 Roseland 光学深度
    let mut taur = HALF * model.dedm1 * model.abrosd[0] * model.dens[0];

    for id in 0..nd {
        // 更新光学深度
        if id > 0 {
            let dtaur = model.deldm[id - 1] * (model.abrosd[id] + model.abrosd[id - 1]);
            taur += dtaur;
        }

        // 计算阈值
        let xgr0 = config.grav * xgrd[id].abs();

        // 初始化 gradi 数组
        for ij in 0..nfreq {
            gradi[ij] = output.gradf[id][ij];
            // 如果不计算辐射压力，跳过所有频率
            if config.ifprad == 0 {
                output.lskip[id][ij] = true;
            } else {
                output.lskip[id][ij] = false;
            }
        }

        // 对辐射加速度排序
        let iigr = indexx(&gradi);

        grad[id] = grada[id];
        let mut ggrt0 = ggrt[id];

        // 如果 XGRAD > 0 且 ID > 1，继承上一层的 LSKIP
        if config.xgrad > 0.0 && id > 0 {
            for ij in 0..nfreq {
                output.lskip[id][ij] = output.lskip[0][ij];
                if output.lskip[id][ij] {
                    ggrt0 -= gradi[ij];
                    grad[id] -= gradi[ij];
                }
            }
            continue;
        }

        // 对于 ID >= ND-1，跳过频率拒绝
        if id >= nd - 2 {
            continue;
        }

        // 频率拒绝循环
        let mut ijr = nfreq as i32 - 1;

        while grad[id] > xgr0 && ijr >= 0 {
            let ij = iigr[ijr as usize];

            // 检查是否是连续谱或无谱线
            if freq.ijlin[ij] == 0 && freq.nlines[ij] == 0 {
                ijr -= 1;
                continue;
            }

            // 标记跳过
            output.lskip[id][ij] = true;
            ggrt0 -= gradi[ij];
            grad[id] -= gradi[ij];

            // 处理显式频率
            if xgrd[id] < 0.0 && nfe < 10 {
                process_explicit_frequency(
                    id, ij, freq, ali, nn, &mut nfe, config.mfrex, config.ispodf,
                );
            }

            ijr -= 1;
        }
    }

    // 辐射压力单位转换
    // PRADT(ID) = PRADT(ID) * PCK
    // 注意：PCK 常量需要从 constants 获取
    // 这里暂时省略单位转换

    RadpreOutput {
        xgrd,
        ggrt,
        grad,
        grada,
        prd0,
        nfe,
    }
}

// ============================================================================
// 辅助函数
// ============================================================================

/// 设置深度相关阈值 XGRD。
fn setup_xgrd(xgrad: f64, nd: usize, xgrd: &mut [f64]) {
    if xgrad == 0.0 {
        // 使用 XGRD0 预设
        for id in 0..nd.min(10) {
            xgrd[id] = XGRD0[id];
        }
        for id in 10..nd {
            xgrd[id] = xgrd[id - 1];
        }
    } else if xgrad == -1.0 {
        // 使用 XGRD1 预设
        for id in 0..nd.min(20) {
            xgrd[id] = XGRD1[id];
        }
        for id in 20..nd {
            xgrd[id] = xgrd[id - 1];
        }
    } else if xgrad == -2.0 {
        // 使用 XGRD2 预设
        for id in 0..nd.min(20) {
            xgrd[id] = XGRD2[id];
        }
        for id in 20..nd {
            xgrd[id] = xgrd[id - 1];
        }
    } else {
        // 使用固定值
        for id in 0..nd {
            xgrd[id] = xgrad;
        }
    }
}

/// 处理显式频率。
#[allow(clippy::too_many_arguments)]
fn process_explicit_frequency(
    id: usize,
    ij: usize,
    freq: &mut RadpreFreqParamsMut,
    ali: &mut RadpreAliParamsMut,
    nn: &mut i32,
    nfe: &mut i32,
    mfrex: usize,
    ispodf: i32,
) {
    if ispodf == 0 {
        // 单谱线情况
        let itr = freq.ijlin[ij];
        if itr == 0 {
            return;
        }

        let itr_idx = (itr.abs() - 1) as usize;
        let indxpa = freq.indexp[itr_idx].abs();

        let dx = (freq.freq[ij] - freq.freq.get(ij + 1).unwrap_or(&0.0)) * 0.25;
        let dz = (freq.freq[ij] - freq.fr0[itr_idx]).abs();

        if dz < dx && indxpa == 1 {
            if freq.indexp[itr_idx] < 0 {
                freq.indexp[itr_idx] = -9;
            } else {
                freq.indexp[itr_idx] = 9;
            }

            if !freq.lexp[itr_idx] {
                freq.lexp[itr_idx] = true;
                *ali.nfreqe += 1;

                if *ali.nfreqe as usize > mfrex {
                    quit("nfreqe.gt.mfrex", *ali.nfreqe, mfrex as i32);
                }

                *nn += 1;
                ali.ijali[ij] = 0;
                ali.ijex[ij] = *ali.nfreqe;
                ali.ijfr[(*ali.nfreqe - 1) as usize] = ij as i32;
                ali.ijx[ij] = 1;
                ali.wc[ij] = 0.0;
                // Fortran: WRITE(10,612) FREQ(IJ),ITR,IJ,NFREQE
                eprintln!(" AUTOMATIC EXPLICIT FREQ.  {:12.6e}{:8}{:8}{:8}", freq.freq[ij], freq.ijlin[ij], ij, *ali.nfreqe);
                *nfe += 1;
            }
        }
    } else {
        // 多谱线情况
        let nlines_ij = freq.nlines[ij];
        for ilint in 0..nlines_ij as usize {
            let itr = ali.itrlin[ilint][ij];
            if itr <= 0 {
                continue;
            }

            let itr_idx = (itr - 1) as usize;
            let indxpa = freq.indexp[itr_idx].abs();

            let dx = (freq.freq[ij] - freq.freq.get(ij + 1).unwrap_or(&0.0)) * 0.25;
            let dz = (freq.freq[ij] - freq.fr0[itr_idx]).abs();

            if dz > dx || indxpa != 1 {
                continue;
            }

            if freq.indexp[itr_idx] < 0 {
                freq.indexp[itr_idx] = -9;
            } else {
                freq.indexp[itr_idx] = 9;
            }

            if !freq.lexp[itr_idx] {
                freq.lexp[itr_idx] = true;
                *ali.nfreqe += 1;

                if *ali.nfreqe as usize > mfrex {
                    quit("nfreqe.gt.mfrex", *ali.nfreqe, mfrex as i32);
                }

                *nn += 1;
                ali.ijali[ij] = 0;
                ali.ijex[ij] = *ali.nfreqe;
                ali.ijfr[(*ali.nfreqe - 1) as usize] = ij as i32;
                ali.ijx[ij] = 1;
                ali.wc[ij] = 0.0;
                // Fortran: WRITE(10,612) FREQ(IJ),ITR,IJ,NFREQE
                eprintln!(" AUTOMATIC EXPLICIT FREQ.  {:12.6e}{:8}{:8}{:8}", freq.freq[ij], itr, ij, *ali.nfreqe);
                *nfe += 1;
            }
        }
    }
}

/// 计算单个频率的辐射加速度贡献。
///
/// 这个函数在频率循环中调用，用于累积辐射加速度。
pub fn radpre_accumulate_frequency(
    ij: usize,
    nd: usize,
    rad: &RadpreRadField,
    model: &RadpreModelState,
    freq: &RadpreFreqParamsMut,
    output: &mut RadpreOutputStateMut,
    grada: &mut [f64],
    prd0: &mut f64,
) {
    let pgrd1 = PGRD / model.dens1[0];
    let mut prid = vec![0.0; nd];

    // 计算 PRID
    for id in 1..nd {
        let dm_diff = model.dm[id] - model.dm[id - 1];
        if dm_diff.abs() > 1e-30 {
            prid[id] = PGRD / dm_diff;
        }
    }

    // 表面辐射加速度
    let fluxw = freq.w[ij] * (rad.rad1[0] * rad.fh[ij] - rad.hextrd[ij]);
    output.gradf[0][ij] = fluxw * rad.abso1[0] * pgrd1;
    grada[0] += output.gradf[0][ij];

    // 深度点辐射加速度
    for id in 1..nd {
        let frd = rad.fak1[id] * rad.rad1[id] - rad.fak1[id - 1] * rad.rad1[id - 1];
        output.gradf[id][ij] = freq.w[ij] * frd * prid[id];
        grada[id] += output.gradf[id][ij];
    }

    // 更新表面辐射压力
    *prd0 += rad.abso1[0] * freq.w[ij] * (rad.rad1[0] * rad.fh[ij] - rad.hextrd[ij]);
}

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

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

    #[test]
    fn test_setup_xgrd_default() {
        let nd = 15;
        let mut xgrd = vec![0.0; nd];

        setup_xgrd(0.0, nd, &mut xgrd);

        // 检查前 10 个元素
        for i in 0..10 {
            assert!((xgrd[i] - XGRD0[i]).abs() < 1e-15);
        }
        // 检查后面的元素继承
        for i in 10..nd {
            assert!((xgrd[i] - xgrd[i - 1]).abs() < 1e-15);
        }
    }

    #[test]
    fn test_setup_xgrd_minus1() {
        let nd = 25;
        let mut xgrd = vec![0.0; nd];

        setup_xgrd(-1.0, nd, &mut xgrd);

        // 检查前 20 个元素
        for i in 0..20 {
            assert!((xgrd[i] - XGRD1[i]).abs() < 1e-15);
        }
        // 检查后面的元素继承
        for i in 20..nd {
            assert!((xgrd[i] - xgrd[i - 1]).abs() < 1e-15);
        }
    }

    #[test]
    fn test_setup_xgrd_minus2() {
        let nd = 25;
        let mut xgrd = vec![0.0; nd];

        setup_xgrd(-2.0, nd, &mut xgrd);

        // 检查前 20 个元素
        for i in 0..20 {
            assert!((xgrd[i] - XGRD2[i]).abs() < 1e-15);
        }
    }

    #[test]
    fn test_setup_xgrd_fixed() {
        let nd = 10;
        let mut xgrd = vec![0.0; nd];

        setup_xgrd(0.5, nd, &mut xgrd);

        // 所有元素应该等于固定值
        for i in 0..nd {
            assert!((xgrd[i] - 0.5).abs() < 1e-15);
        }
    }

    #[test]
    fn test_pgrd_constant() {
        // 验证 PGRD 常量值
        assert!((PGRD - 4.1916825e-10).abs() < 1e-20);
    }

    #[test]
    fn test_xgrd_arrays() {
        // 验证预设数组长度
        assert_eq!(XGRD0.len(), 10);
        assert_eq!(XGRD1.len(), 20);
        assert_eq!(XGRD2.len(), 20);

        // 验证数组是递增的
        for i in 1..XGRD0.len() {
            assert!(XGRD0[i] >= XGRD0[i - 1]);
        }
        for i in 1..XGRD1.len() {
            assert!(XGRD1[i] >= XGRD1[i - 1]);
        }
        for i in 1..XGRD2.len() {
            assert!(XGRD2[i] >= XGRD2[i - 1]);
        }
    }
}