SpectraRust/src/math/prnt.rs

//! 特定能级的辐射率和碰撞率平衡计算。
//!
//! 重构自 TLUSTY `PRNT` 子程序。
//!
//! # 功能
//!
//! - 计算指定能级的辐射率和碰撞率流入/流出
//! - 用于调试和分析能级占据数的速率平衡

use crate::state::atomic::AtomicData;
use crate::state::config::InpPar;
use crate::state::constants::HK;
use crate::state::model::{CraTes, LevPop, ModPar, RrRates, WmComp};

use super::sabolf::{sabolf_pure, SabolfParams};

// ============================================================================
// 输出结构体
// ============================================================================

/// 单个能级的速率平衡结果。
#[derive(Debug, Clone)]
pub struct RateBalance {
    /// 深度索引
    pub id: usize,
    /// 能级索引 (Fortran 1-indexed)
    pub ii: i32,
    /// 流出率 (辐射 + 碰撞)
    pub rou: f64,
    /// 流入率 (辐射 + 碰撞)
    pub rin: f64,
    /// 相对不平衡度 (rou - rin) / rin
    pub imbalance: f64,
}

/// PRNT 输出结果。
#[derive(Debug, Clone)]
pub struct PrntOutput {
    /// 各能级的速率平衡
    pub balances: Vec<RateBalance>,
}

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

/// PRNT 输入参数。
pub struct PrntParams<'a> {
    /// 模型基本参数
    pub modpar: &'a ModPar,
    /// 能级占据数
    pub levpop: &'a LevPop,
    /// 能级权重和占据概率
    pub wmcomp: &'a WmComp,
    /// 辐射率
    pub rrrates: &'a RrRates,
    /// 碰撞率
    pub crates: &'a CraTes,
    /// 原子数据
    pub atomic: &'a AtomicData,
    /// 配置参数
    pub inppar: &'a InpPar,
    /// 要分析的能级索引列表 (Fortran 1-indexed)
    pub ipop: &'a [i32],
}

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

/// 计算特定能级的速率平衡（纯计算部分）。
///
/// # 参数
/// - `params`: 输入参数
///
/// # 返回
/// 各能级的速率平衡结果
pub fn prnt_pure(params: &PrntParams) -> PrntOutput {
    let mut balances = Vec::new();

    let nd = params.modpar.temp.len();
    let atomic = params.atomic;
    let levpop = params.levpop;
    let wmcomp = params.wmcomp;
    let rrrates = params.rrrates;
    let crates = params.crates;
    let inppar = params.inppar;

    // 遍历深度点，步长 69
    for id in (0..nd).step_by(69) {
        let temp = params.modpar.temp[id];
        let ane = params.modpar.elec[id];
        let hkt = HK / temp;

        // 调用 sabolf 计算 Saha-Boltzmann 因子
        let sabolf_params = SabolfParams {
            id,
            t: temp,
            ane,
            atomic,
            wnhint: None,
            ioptab: 0,
        };
        let sabolf_result = sabolf_pure(&sabolf_params);
        let sbf = &sabolf_result.sbf;
        let usum = &sabolf_result.usum;

        // 遍历要分析的能级
        for &ii_1idx in params.ipop {
            // ii_1idx 是 Fortran 1-indexed，转换为 0-indexed
            let ii = (ii_1idx - 3) as usize; // Fortran: ii = ipop(k) - 3

            // 获取原子和离子索引
            let iat = if ii < atomic.levpar.iatm.len() {
                atomic.levpar.iatm[ii] as usize
            } else {
                continue;
            };
            let ie = if ii < atomic.levpar.iel.len() {
                atomic.levpar.iel[ii] as usize
            } else {
                continue;
            };

            // 计算该原子所有能级的占据数之和
            let n0a = if iat < atomic.atopar.n0a.len() {
                atomic.atopar.n0a[iat]
            } else {
                continue;
            };
            let nka = if iat < atomic.atopar.nka.len() {
                atomic.atopar.nka[iat]
            } else {
                continue;
            };

            // 验证 ii 是否在 [n0a, nka] 范围内
            if (ii + 1) < n0a as usize || (ii + 1) > nka as usize {
                continue;
            }

            let mut psum = 0.0_f64;
            let mut psuu = 0.0_f64;

            for j in (n0a as usize - 1)..nka as usize {
                // j 是 0-indexed
                psum += levpop.popul[j][id];

                let ilk_j = if j < atomic.levpar.ilk.len() {
                    atomic.levpar.ilk[j]
                } else {
                    0
                };

                if ilk_j > 0 {
                    let ilk_idx = (ilk_j - 1) as usize;
                    let usum_val = if ilk_idx < usum.len() {
                        usum[ilk_idx]
                    } else {
                        0.0
                    };
                    psuu += usum_val * ane * levpop.popul[j][id];
                }
            }

            // 计算 BB = DENS(ID)/WMM(ID)/YTOT(ID)*ABUND(IAT,ID)
            let dens_id = params.modpar.dens[id];
            let wmm_id = inppar.wmm[id];
            let ytot_id = inppar.ytot[id];
            let abund_iat = if iat < atomic.atopar.abund.len() {
                atomic.atopar.abund[iat][id]
            } else {
                0.0
            };

            let _bb = if wmm_id != 0.0 && ytot_id != 0.0 {
                dens_id / wmm_id / ytot_id * abund_iat
            } else {
                0.0
            };

            // 获取离子参数
            let nfirst = if ie < atomic.ionpar.nfirst.len() {
                atomic.ionpar.nfirst[ie]
            } else {
                continue;
            };
            let nlast = if ie < atomic.ionpar.nlast.len() {
                atomic.ionpar.nlast[ie]
            } else {
                continue;
            };
            let nnext = if ie < atomic.ionpar.nnext.len() {
                atomic.ionpar.nnext[ie]
            } else {
                continue;
            };

            let mut rin = 0.0_f64;
            let mut rou = 0.0_f64;

            // 遍历较低能级 (jj < ii)
            // Fortran: do jj = nfirst(ie), ii-1
            for jj_1idx in nfirst..((ii + 1) as i32) {
                let jj = (jj_1idx - 1) as usize; // 转换为 0-indexed

                // 获取跃迁索引
                let itr = if jj < atomic.trapar.itra.len()
                    && ii < atomic.trapar.itra[jj].len()
                {
                    atomic.trapar.itra[jj][ii]
                } else {
                    continue;
                };

                if itr <= 0 {
                    continue;
                }

                let itr_idx = (itr - 1) as usize;

                // 获取跃迁率
                let rru_val = if itr_idx < rrrates.rru.len() && id < rrrates.rru[itr_idx].len() {
                    rrrates.rru[itr_idx][id]
                } else {
                    0.0
                };
                let colrat_val =
                    if itr_idx < crates.colrat.len() && id < crates.colrat[itr_idx].len() {
                        crates.colrat[itr_idx][id]
                    } else {
                        0.0
                    };
                let rrd_val = if itr_idx < rrrates.rrd.len() && id < rrrates.rrd[itr_idx].len() {
                    rrrates.rrd[itr_idx][id]
                } else {
                    0.0
                };
                let coltar_val =
                    if itr_idx < crates.coltar.len() && id < crates.coltar[itr_idx].len() {
                        crates.coltar[itr_idx][id]
                    } else {
                        0.0
                    };

                let wop_ii = if ii < wmcomp.wop.len() && id < wmcomp.wop[ii].len() {
                    wmcomp.wop[ii][id]
                } else {
                    1.0
                };
                let wop_jj = if jj < wmcomp.wop.len() && id < wmcomp.wop[jj].len() {
                    wmcomp.wop[jj][id]
                } else {
                    1.0
                };

                let g_jj = if jj < atomic.levpar.g.len() {
                    atomic.levpar.g[jj]
                } else {
                    1.0
                };
                let g_ii = if ii < atomic.levpar.g.len() {
                    atomic.levpar.g[ii]
                } else {
                    1.0
                };

                let fr0_val = if itr_idx < atomic.trapar.fr0.len() {
                    atomic.trapar.fr0[itr_idx]
                } else {
                    0.0
                };

                // 上跃迁率 (jj -> ii)
                let ru = rru_val * wop_ii;
                let cu = colrat_val * wop_ii;

                // 下跃迁率 (ii -> jj)
                let (rd, cd) = if (ii + 1) as i32 <= nlast {
                    // 束缚-束缚跃迁
                    let rd = rrd_val * g_jj / g_ii * (hkt * fr0_val).exp() * wop_jj;
                    let cd = coltar_val * wop_jj;
                    (rd, cd)
                } else {
                    // 束缚-自由跃迁
                    let sbf_jj = if jj < sbf.len() { sbf[jj] } else { 1.0 };
                    let rd = rrd_val * sbf_jj * ane * wop_jj;
                    let cd = coltar_val * wop_jj;
                    (rd, cd)
                };

                let popul_jj = if jj < levpop.popul.len() && id < levpop.popul[jj].len() {
                    levpop.popul[jj][id]
                } else {
                    0.0
                };
                let popul_ii = if ii < levpop.popul.len() && id < levpop.popul[ii].len() {
                    levpop.popul[ii][id]
                } else {
                    0.0
                };

                rin += (ru + cu) * popul_jj;
                rou += (rd + cd) * popul_ii;
            }

            // 遍历较高能级 (jj > ii)
            // Fortran: do jj = ii+1, nnext(ie)
            for jj_1idx in ((ii + 2) as i32)..=nnext {
                let jj = (jj_1idx - 1) as usize; // 转换为 0-indexed

                // 获取跃迁索引
                let itr = if ii < atomic.trapar.itra.len()
                    && jj < atomic.trapar.itra[ii].len()
                {
                    atomic.trapar.itra[ii][jj]
                } else {
                    continue;
                };

                if itr <= 0 {
                    continue;
                }

                let itr_idx = (itr - 1) as usize;

                // 获取跃迁率
                let rru_val = if itr_idx < rrrates.rru.len() && id < rrrates.rru[itr_idx].len() {
                    rrrates.rru[itr_idx][id]
                } else {
                    0.0
                };
                let colrat_val =
                    if itr_idx < crates.colrat.len() && id < crates.colrat[itr_idx].len() {
                        crates.colrat[itr_idx][id]
                    } else {
                        0.0
                    };
                let rrd_val = if itr_idx < rrrates.rrd.len() && id < rrrates.rrd[itr_idx].len() {
                    rrrates.rrd[itr_idx][id]
                } else {
                    0.0
                };
                let coltar_val =
                    if itr_idx < crates.coltar.len() && id < crates.coltar[itr_idx].len() {
                        crates.coltar[itr_idx][id]
                    } else {
                        0.0
                    };

                let wop_ii = if ii < wmcomp.wop.len() && id < wmcomp.wop[ii].len() {
                    wmcomp.wop[ii][id]
                } else {
                    1.0
                };
                let wop_jj = if jj < wmcomp.wop.len() && id < wmcomp.wop[jj].len() {
                    wmcomp.wop[jj][id]
                } else {
                    1.0
                };

                let g_jj = if jj < atomic.levpar.g.len() {
                    atomic.levpar.g[jj]
                } else {
                    1.0
                };
                let g_ii = if ii < atomic.levpar.g.len() {
                    atomic.levpar.g[ii]
                } else {
                    1.0
                };

                let fr0_val = if itr_idx < atomic.trapar.fr0.len() {
                    atomic.trapar.fr0[itr_idx]
                } else {
                    0.0
                };

                // 上跃迁率 (ii -> jj)
                let ru = rru_val * wop_jj;
                let cu = colrat_val * wop_jj;

                // 下跃迁率 (jj -> ii)
                let (rd, cd) = if jj_1idx <= nlast {
                    // 束缚-束缚跃迁
                    let rd = rrd_val * g_ii / g_jj * (hkt * fr0_val).exp() * wop_ii;
                    let cd = coltar_val * wop_ii;
                    (rd, cd)
                } else {
                    // 束缚-自由跃迁
                    let sbf_ii = if ii < sbf.len() { sbf[ii] } else { 1.0 };
                    let rd = rrd_val * sbf_ii * ane * wop_ii;
                    let cd = coltar_val * wop_ii;
                    (rd, cd)
                };

                let popul_jj = if jj < levpop.popul.len() && id < levpop.popul[jj].len() {
                    levpop.popul[jj][id]
                } else {
                    0.0
                };
                let popul_ii = if ii < levpop.popul.len() && id < levpop.popul[ii].len() {
                    levpop.popul[ii][id]
                } else {
                    0.0
                };

                rou += (ru + cu) * popul_ii;
                rin += (rd + cd) * popul_jj;
            }

            // 计算相对不平衡度
            let imbalance = if rin != 0.0 { (rou - rin) / rin } else { 0.0 };

            balances.push(RateBalance {
                id,
                ii: ii_1idx,
                rou,
                rin,
                imbalance,
            });
        }
    }

    PrntOutput { balances }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::state::atomic::{AtoPar, IonPar, LevPar, TraPar};
    use crate::state::config::InpPar;
    use crate::state::constants::{MDEPTH, MION, MLEVEL, MTRANS};
    use crate::state::model::{CraTes, LevPop, ModPar, RrRates, WmComp};

    fn create_test_modpar() -> ModPar {
        let mut modpar = ModPar::default();
        modpar.temp[0] = 10000.0;
        modpar.temp[1] = 9000.0;
        modpar.elec[0] = 1.0e12;
        modpar.elec[1] = 1.0e11;
        modpar.dens[0] = 1.0e14;
        modpar.dens[1] = 1.0e13;
        modpar
    }

    fn create_test_levpop() -> LevPop {
        LevPop::default()
    }

    fn create_test_atomic() -> AtomicData {
        let mut atomic = AtomicData::default();

        // 设置能级数据
        for i in 0..100 {
            atomic.levpar.g[i] = 2.0;
            atomic.levpar.iatm[i] = 1;
            atomic.levpar.iel[i] = 1;
            atomic.levpar.ilk[i] = 0;
            atomic.levpar.enion[i] = 10.0 - i as f64 * 0.1;
        }

        // 设置原子数据
        atomic.atopar.n0a[0] = 1;
        atomic.atopar.nka[0] = 50;
        atomic.atopar.abund[0][0] = 0.1;

        // 设置离子数据 - 只设置第一个离子，并确保 nfirst >= 1
        atomic.ionpar.nfirst[0] = 1;
        atomic.ionpar.nlast[0] = 50;
        atomic.ionpar.nnext[0] = 51;
        atomic.ionpar.iz[0] = 1;

        // 清空其他离子的数据，避免 sabolf 处理无效离子
        // sabolf 基于 iz.len() 遍历，而 iz.len() = MION
        // 为了避免处理无效离子，设置 nfirst > nlast 使循环跳过
        // 同时确保 nlast >= 1 避免 nlst = nlast - 1 下溢
        for i in 1..MION {
            atomic.ionpar.nfirst[i] = 2; // nfirst > nlast
            atomic.ionpar.nlast[i] = 1;
            atomic.ionpar.nnext[i] = 0;
        }

        atomic
    }

    fn create_test_wmcomp() -> WmComp {
        let mut wmcomp = WmComp::default();
        for i in 0..10 {
            for j in 0..MDEPTH {
                wmcomp.wop[i][j] = 1.0;
            }
        }
        wmcomp
    }

    fn create_test_rrrates() -> RrRates {
        RrRates::default()
    }

    fn create_test_crates() -> CraTes {
        CraTes::default()
    }

    fn create_test_inppar() -> InpPar {
        let mut inppar = InpPar::default();
        inppar.wmm[0] = 1.0;
        inppar.ytot[0] = 1.0;
        inppar
    }

    #[test]
    fn test_prnt_basic() {
        let modpar = create_test_modpar();
        let levpop = create_test_levpop();
        let atomic = create_test_atomic();
        let wmcomp = create_test_wmcomp();
        let rrrates = create_test_rrrates();
        let crates = create_test_crates();
        let inppar = create_test_inppar();

        // 测试能级索引 (Fortran 1-indexed)
        let ipop = [98, 99, 100, 115];

        let params = PrntParams {
            modpar: &modpar,
            levpop: &levpop,
            wmcomp: &wmcomp,
            rrrates: &rrrates,
            crates: &crates,
            atomic: &atomic,
            inppar: &inppar,
            ipop: &ipop,
        };

        let result = prnt_pure(&params);

        // 由于测试数据是空的，结果应该为空或只有有限的结果
        println!("Number of balances: {}", result.balances.len());
    }

    #[test]
    fn test_prnt_with_populations() {
        let _modpar = create_test_modpar();
        let mut levpop = create_test_levpop();
        let atomic = create_test_atomic();
        let wmcomp = create_test_wmcomp();
        let rrrates = create_test_rrrates();
        let crates = create_test_crates();
        let inppar = create_test_inppar();

        // 设置一些非零占据数
        for i in 0..50 {
            levpop.popul[i][0] = 0.01;
        }

        // 测试能级索引
        let ipop = [100]; // 只测试能级 100

        let params = PrntParams {
            modpar: &_modpar,
            levpop: &levpop,
            wmcomp: &wmcomp,
            rrrates: &rrrates,
            crates: &crates,
            atomic: &atomic,
            inppar: &inppar,
            ipop: &ipop,
        };

        let result = prnt_pure(&params);

        // 验证结果
        for balance in &result.balances {
            println!(
                "id={}, ii={}, rou={}, rin={}, imbalance={}",
                balance.id, balance.ii, balance.rou, balance.rin, balance.imbalance
            );
        }
    }

    #[test]
    fn test_rate_balance_structure() {
        let balance = RateBalance {
            id: 0,
            ii: 100,
            rou: 1.0e10,
            rin: 1.0e10,
            imbalance: 0.0,
        };

        assert_eq!(balance.id, 0);
        assert_eq!(balance.ii, 100);
        assert!((balance.rou - 1.0e10).abs() < 1e5);
        assert!((balance.rin - 1.0e10).abs() < 1e5);
        assert!(balance.imbalance.abs() < 1e-10);
    }
}