SpectraRust/src/tlusty/math/solvers/rybchn.rs

//! 处理 Rybicki 公式中的相对变化。
//!
//! 重构自 TLUSTY `rybchn.f`。
//!
//! 功能：
//! - 处理温度变化并限制变化幅度
//! - 更新温度、压力和密度分布
//! - 使用 ELDENS 计算电子密度
//! - 使用 PGSET 迭代计算气体压力

use crate::tlusty::math::{eldens, EldensConfig, EldensOutput, EldensParams};
use crate::tlusty::math::{pgset, PgsetParams, PgsetOutput};
use crate::tlusty::state::constants::{BOLK, HALF, TWO, UN};

/// 最大深度点数（从 pgset 导入）
use crate::tlusty::math::MDEPTH;

/// RYBCHN 配置参数
#[derive(Debug, Clone)]
pub struct RybchnConfig {
    /// 迭代次数
    pub iter: i32,
    /// 最大迭代次数
    pub niter: i32,
    /// 最大允许变化
    pub chmax: f64,
    /// DPSILT 参数（变化上限）
    pub dpsilt: f64,
    /// NFREQE + 1（用于 PSI0 数组）
    pub nfreqe_p1: usize,
    /// NRETC 参数（负值表示特殊处理）
    pub nretc: i32,
    /// IDISK 参数（0=平面平行，其他=球对称）
    pub idisk: i32,
    /// IFPRAD 参数（辐射压力标志）
    pub ifprad: i32,
    /// IOPTAB 参数
    pub ioptab: i32,
    /// QGRAV 常数
    pub qgrav: f64,
    /// GRAV 常数
    pub grav: f64,
    /// PCK 常数
    pub pck: f64,
}

impl Default for RybchnConfig {
    fn default() -> Self {
        Self {
            iter: 1,
            niter: 100,
            chmax: 0.1,
            dpsilt: 2.0,
            nfreqe_p1: 1,
            nretc: 0,
            idisk: 0,
            ifprad: 0,
            ioptab: 0,
            qgrav: 1.0,
            grav: 1.0,
            pck: 1.0,
        }
    }
}

/// RYBCHN 输入参数
#[derive(Debug, Clone)]
pub struct RybchnParams {
    /// 配置参数
    pub config: RybchnConfig,
    /// 深度点数
    pub nd: usize,
    /// 温度变化数组 CHANGT
    pub changt: Vec<f64>,
    /// 温度数组 TEMP
    pub temp: Vec<f64>,
    /// 柱质量密度数组 DM
    pub dm: Vec<f64>,
    /// 深度变量 ZD
    pub zd: Vec<f64>,
    /// 平均分子量数组 WMM
    pub wmm: Vec<f64>,
    /// 辐射压力梯度 GRD
    pub grd: Vec<f64>,
    /// 气体压力 PGS
    pub pgs: Vec<f64>,
    /// CS 数组（用于 PGSET）
    pub cs: Vec<f64>,
    /// PRAD2D 数组（用于 PGSET）
    pub prad2d: Vec<f64>,
    /// F1HE 参数（用于 PGSET）
    pub f1he: f64,
    /// ELDENS 配置
    pub eldens_config: EldensConfig,
    /// YTOT 参数（总氢丰度因子）
    pub ytot: f64,
    /// QREF 参数
    pub qref: f64,
    /// DQNR 参数
    pub dqnr: f64,
    /// WMY 参数
    pub wmy: f64,
    /// PGSET 温度迭代次数
    pub ntemp: usize,
}

/// RYBCHN 输出结果
#[derive(Debug, Clone)]
pub struct RybchnOutput {
    /// 更新后的温度
    pub temp: Vec<f64>,
    /// 更新后的物质密度
    pub dens: Vec<f64>,
    /// 更新后的电子密度
    pub elec: Vec<f64>,
    /// 更新后的气体压力
    pub pgs: Vec<f64>,
    /// 更新后的 PGS0
    pub pgs0: Vec<f64>,
    /// 更新后的 ANTP
    pub antp: Vec<f64>,
    /// 更新后的深度变量 ZD
    pub zd: Vec<f64>,
    /// 更新后的 CS
    pub cs: Vec<f64>,
    /// 更新后的 F1HE
    pub f1he: f64,
    /// 最大变化
    pub chmx: f64,
    /// 是否最后一次迭代
    pub lfin: bool,
    /// PSI0(NRE) 值
    pub psi0_nre: f64,
    /// PSY0(NRE, :) 数组
    pub psy0_nre: Vec<f64>,
}

/// 处理 Rybicki 公式中的相对变化（纯计算函数）。
///
/// # 参数
/// * `params` - 输入参数
///
/// # 返回值
/// 包含更新后的温度、密度、压力等的输出结构体
pub fn rybchn(params: &RybchnParams) -> RybchnOutput {
    let nd = params.nd;
    let config = &params.config;

    // 变化限制
    let dplp = config.dpsilt - UN;
    let dplm = UN / config.dpsilt - UN;
    let nre = config.nfreqe_p1;

    // WRITE(9,800) - header on first iteration
    if config.iter == 1 {
        eprintln!(" RELATIVE CHANGES OF VECTOR PSI");
        eprintln!("  ITER  ID    TEMP        NE       POP     RAD      MAXIMUM   ilev   ifr");
    }

    // 临时数组
    let mut temp = params.temp.clone();
    let mut pgs = params.pgs.clone();
    let mut dens = vec![0.0; nd];
    let mut elec = vec![0.0; nd];
    let mut pgs0 = vec![0.0; nd];
    let mut antp = vec![0.0; nd];
    let mut zd = params.zd.clone();
    let mut cs = params.cs.clone();
    let mut f1he = params.f1he;
    let mut psy0_nre = vec![0.0; nd];
    let mut psi0_nre = 0.0;

    let mut chmx = 0.0_f64;

    // 步骤 1：处理温度变化
    let tmpold = temp.clone();
    let mut chn: f64 = 0.0;
    let mut jjr: i32 = 0;

    for id in (0..nd).rev() {
        let cht = params.changt[id] / temp[id];
        let mut chan = cht;

        // 限制变化幅度
        if chan <= dplm {
            chan = dplm;
        }
        if chan > dplp {
            chan = dplp;
        }

        temp[id] = temp[id] * (chan + UN);

        // 更新 PSI0 和 PSY0
        psi0_nre = temp[id];
        psy0_nre[id] = psi0_nre;

        // WRITE(9,801) ITER,ID,CHT,chn,chN,CHN,chT,jjr,jjr
        eprintln!("{:5}{:5}{:10.2e}{:10.2e}{:10.2e}{:10.2e}{:10.2e}{:5}{:5}",
            config.iter, id + 1, cht, chan, chn, chn, cht, jjr, jjr);

        // 跟踪最大变化
        if cht.abs() > chmx {
            chmx = cht.abs();
        }
    }

    // 步骤 2：处理 NRETC < 0 的情况（特殊边界处理）
    if config.nretc < 0 {
        let start = (-config.nretc) as usize;
        if start < nd {
            for id in (0..start).rev() {
                temp[id] = temp[id + 1];
                psy0_nre[id] = psy0_nre[id + 1];
            }
        }
    }

    // 步骤 3：根据 IOPTAB 和 IDISK 处理压力和密度
    if config.ioptab > -2 {
        if config.idisk == 0 {
            // 平面平行几何
            if config.ifprad > 0 {
                // 包含辐射压力
                for id in 0..nd {
                    let t = temp[id];
                    let dtod = temp[id] / tmpold[id] - UN;

                    // 限制温度变化
                    let dtod = if dtod > 0.2 { 0.2 } else if dtod < -0.2 { -0.2 } else { dtod };

                    let gfac = UN + 4.0 * dtod;

                    if id == 0 {
                        pgs[id] = params.dm[id] * (config.grav - params.grd[id] * gfac);
                    } else {
                        pgs[id] = pgs[id - 1]
                            + config.grav * (params.dm[id] - params.dm[id - 1])
                            - config.pck * params.grd[id] * gfac;
                    }

                    let an = pgs[id] / BOLK / t;
                    let eldens_out = compute_eldens(params, id, t, an);

                    dens[id] = params.wmm[id] * (an - eldens_out.ane);
                    elec[id] = eldens_out.ane;
                }
            } else {
                // 不包含辐射压力
                for id in 0..nd {
                    let t = temp[id];
                    pgs[id] = params.dm[id] * config.grav;
                    let an = pgs[id] / BOLK / t;

                    let eldens_out = compute_eldens(params, id, t, an);

                    dens[id] = params.wmm[id] * (an - eldens_out.ane);
                    elec[id] = eldens_out.ane;
                }
            }
        } else {
            // 球对称几何
            let _pgpre = pgs[0];

            for id in 1..nd {
                let dtod = temp[id] / tmpold[id] - UN;
                let dtod = if dtod > 0.2 { 0.2 } else if dtod < -0.2 { -0.2 } else { dtod };

                let gfac = UN + 4.0 * dtod;
                let grv = (params.dm[id] - params.dm[id - 1])
                    * config.qgrav
                    * (zd[id] + zd[id - 1])
                    * HALF;

                pgs[id] = pgs[id - 1] - config.pck * params.grd[id] * gfac + grv;
                pgs0[id] = pgs[id];
            }
            pgs0[0] = pgs[0];

            // 迭代计算压力
            let mut itpg = 0;
            let mut z1 = zd[0];

            loop {
                itpg += 1;

                // 调用 PGSET
                let pgset_params = PgsetParams {
                    nd,
                    ntemp: params.ntemp,
                    dm: params.dm.clone(),
                    temp: temp.clone(),
                    pgs0: pgs0.clone(),
                    cs: cs.clone(),
                    prad2d: params.prad2d.clone(),
                    f1he,
                    qgrav: config.qgrav,
                    bolk: BOLK,
                };

                let pgset_out = pgset(&pgset_params);

                // 更新压力和温度
                for id in 0..nd {
                    pgs0[id] = pgset_out.pgs0[id];
                    antp[id] = pgset_out.antp[id];
                }

                // 使用更新后的 ANTP 计算 ELDENS
                for id in 0..nd {
                    let t = temp[id];
                    let an = antp[id];

                    let eldens_out = compute_eldens(params, id, t, an);

                    dens[id] = params.wmm[id] * (an - eldens_out.ane);
                    elec[id] = eldens_out.ane;
                    pgs[id] = BOLK * t * an;
                    pgs0[id] = pgs[id];
                }

                // 重新计算深度变量 ZD
                for id in 0..(nd - 1) {
                    let ddp = (params.dm[id + 1] - params.dm[id]) * HALF;
                    zd[id] = zd[id + 1] + ddp / dens[id + 1] + ddp / dens[id];
                }

                // 检查收敛
                if ((zd[0] - z1) / z1).abs() < 1e-3 || itpg > 5 {
                    break;
                }

                // 更新 CS
                for id in 0..nd {
                    cs[id] = pgs[id] / dens[id] / temp[id];
                }

                // 计算 F1HE
                let hr1 = params.grd[0] / config.qgrav;
                let hg1 = (TWO * cs[0] * temp[0] / config.qgrav).sqrt();
                let x = (zd[0] - hr1) / hg1;

                f1he = if x < 3.0 {
                    let x = if x < 0.0 { 0.0 } else { x };
                    // 8.86226925D-1 * EXP(X*X) * ERFCX(X)
                    0.886226925 * (x * x).exp() * erfcx_approx(x)
                } else {
                    HALF * (UN - HALF / x / x) / x
                };

                if ((zd[0] - z1) / z1).abs() < 1e-4 || itpg > 5 {
                    break;
                }

                z1 = zd[0];
            }
        }
    }

    // 步骤 4：检查是否最后一次迭代
    // STOP if changes become too large
    if config.iter != 1 && chmx.abs() > 1e16 {
        eprintln!(" **** STOP in RYBSOL after ITER{:4}", config.iter);
        eprintln!(" Max change:{:12.2e}", chmx);
        panic!("RYBSOL: max change too large");
    }

    let lfin = chmx.abs() <= config.chmax || config.iter >= config.niter;

    RybchnOutput {
        temp,
        dens,
        elec,
        pgs,
        pgs0,
        antp,
        zd,
        cs,
        f1he,
        chmx,
        lfin,
        psi0_nre,
        psy0_nre,
    }
}

/// 辅助函数：计算电子密度
fn compute_eldens(params: &RybchnParams, id: usize, t: f64, an: f64) -> EldensOutput {
    let eldens_params = EldensParams {
        id: id + 1, // Fortran 1-indexed
        t,
        an,
        ytot: params.ytot,
        qref: params.qref,
        dqnr: params.dqnr,
        wmy: params.wmy,
        config: params.eldens_config.clone(),
        state_params: None,
        molecule_data: None,
        anato_data: None,
    };

    eldens(&eldens_params, 1)
}

/// ERFCX 近似函数（缩放互补误差函数）
fn erfcx_approx(x: f64) -> f64 {
    // 使用简单的近似
    // 对于小 x: erfcx(x) ≈ exp(x²) * erfc(x)
    // 对于大 x: erfcx(x) ≈ 1/(sqrt(π)*x)
    if x < 0.0 {
        return 2.0 * (x * x).exp() - erfcx_approx(-x);
    }

    if x < 0.5 {
        // 小 x 近似
        let a = 0.886226925; // sqrt(π)/2
        let x2 = x * x;
        a * (1.0 - x * (1.128379167 - x * (0.376126389 - x * 0.09647576)))
    } else if x < 3.0 {
        // 中等 x 使用表格或更精确近似
        // 简化为 exp(x²)*erfc(x) 的近似
        let t = 1.0 / (1.0 + 0.5 * x);
        let tau = t
            * (0.17087211
                + t * (-0.32684114 + t * (0.36039897 + t * (-0.25734667 + t * (0.14029546 + t * -0.04159884)))));
        (x * x).exp() * tau * 0.5641895835 // sqrt(1/π) = 0.5641895835
    } else {
        // 大 x 近似
        1.0 / (x * std::f64::consts::PI.sqrt())
    }
}

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

    fn create_test_params() -> RybchnParams {
        let nd = 5;

        RybchnParams {
            config: RybchnConfig::default(),
            nd,
            changt: vec![0.01, 0.02, 0.03, 0.02, 0.01], // 小的温度变化
            temp: vec![10000.0, 9500.0, 9000.0, 8500.0, 8000.0],
            dm: vec![0.1, 0.2, 0.3, 0.4, 0.5],
            zd: vec![1.0, 2.0, 3.0, 4.0, 5.0],
            wmm: vec![1.0; nd],
            grd: vec![0.1; nd],
            pgs: vec![1e5; nd],
            cs: vec![1e8; nd],
            prad2d: vec![0.0; nd],
            f1he: 1.0,
            eldens_config: EldensConfig::default(),
            ytot: 1.0,
            qref: 0.0,
            dqnr: 0.0,
            wmy: 1.0,
            ntemp: 1,
        }
    }

    #[test]
    fn test_rybchn_basic() {
        let params = create_test_params();
        let output = rybchn(&params);

        // 检查输出维度
        assert_eq!(output.temp.len(), params.nd);
        assert_eq!(output.dens.len(), params.nd);
        assert_eq!(output.elec.len(), params.nd);

        // 检查温度更新（应该增加）
        for id in 0..params.nd {
            assert!(output.temp[id] > params.temp[id] * 0.9);
            assert!(output.temp[id] < params.temp[id] * 1.1);
        }

        // 检查收敛标志
        // 由于变化很小，应该收敛
        println!("chmx = {}", output.chmx);
        println!("lfin = {}", output.lfin);
    }

    #[test]
    fn test_rybchn_large_change() {
        let mut params = create_test_params();
        // 设置较大的温度变化
        params.changt = vec![0.5, 0.5, 0.5, 0.5, 0.5];

        let output = rybchn(&params);

        // 检查温度变化被限制
        // 由于 dplp = dpsilt - 1 = 2 - 1 = 1, dplm = 1/dpsilt - 1 = 0.5 - 1 = -0.5
        // 所以变化限制在 -0.5 到 1.0 之间
        for id in 0..params.nd {
            let expected_max = params.temp[id] * (1.0 + 1.0); // temp * (dplp + 1)
            let expected_min = params.temp[id] * (1.0 - 0.5); // temp * (dplm + 1)
            assert!(output.temp[id] <= expected_max * 1.01);
            assert!(output.temp[id] >= expected_min * 0.99);
        }
    }

    #[test]
    fn test_erfcx_approx() {
        // 测试 ERFCX 近似函数
        let test_values = [0.0, 0.1, 0.5, 1.0, 2.0, 3.0];

        for x in test_values {
            let result = erfcx_approx(x);
            println!("erfcx({}) = {}", x, result);
            // 基本检查：erfcx 应该是正数
            assert!(result > 0.0);
        }
    }
}