//! NLTE 控制过程。
//!
//! 重构自 SYNSPEC `NLTE` 函数。
//!
//! 计算线中心不透明度 (ABCENT) 和线源函数 (SLIN)。

use crate::synspec::state::constants::{BOLK, MDEPTH, MLEVEL};
use crate::synspec::math::{eps, xk2dop};

/// NLTE 输入参数。
pub struct NlteParams<'a> {
    /// 线索引
    pub il: usize,
    /// 下能级索引 (0 = 非显式能级)
    pub ilw: usize,
    /// 上能级索引 (0 = 非显式能级)
    pub iun: usize,
    /// 下能级统计权重
    pub gi: f64,
    /// 上能级统计权重
    pub gj: f64,
    /// 深度点数
    pub nd: usize,
    /// 温度数组 (K)
    pub temp: &'a [f64; MDEPTH],
    /// 电子密度数组 (cm^-3)
    pub elec: &'a [f64; MDEPTH],
    /// 密度数组 (g/cm^3)
    pub dens: &'a [f64; MDEPTH],
    /// 质量深度数组 (g/cm^2)
    pub dm: &'a [f64; MDEPTH],
    /// 湍流速度平方 (cm/s)^2
    pub vturb: &'a [f64; MDEPTH],
    /// 能级 populations
    pub popul: &'a [[f64; MDEPTH]; MLEVEL],
    /// 统计权重
    pub g: &'a [f64; MLEVEL],
    /// 电离能
    pub enion: &'a [f64; MLEVEL],
    /// 相对 populations (N/U)
    pub rrr: &'a [[[f64; MDEPTH]; 90]; 30],
    /// 标准丰度
    pub abstd: &'a [f64; MDEPTH],
    /// 窗口标准丰度
    pub abstdw: &'a [[f64; MDEPTH]],
    /// NLTE populations
    pub pnlt: &'a [[[f64; MDEPTH]; 90]; 30],
    /// 原子质量
    pub amas: &'a [f64],
    /// 电离势 (eV)
    pub enev: &'a [[f64; 90]; 30],
    /// 下一个离子索引
    pub nnext: &'a [i32],
    /// 元素索引
    pub iel: &'a [i32],
    /// 线频率 (Hz)
    pub freq0: f64,
    /// 下能级能量 (K)
    pub excl0: f64,
    /// 上能级能量 (K)
    pub excu0: f64,
    /// gf 值的对数
    pub gf0: f64,
    /// 原子索引 (1-based)
    pub iat: usize,
    /// 离子索引 (1-based)
    pub ion: usize,
    /// NLTE 线索引
    pub ilnlt: usize,
    /// 窗口标志
    pub ifwin: i32,
    /// 窗口连续谱索引
    pub ijcont: usize,
}

/// NLTE 输出结果。
pub struct NlteOutput {
    /// 线中心不透明度 (每深度点)
    pub abcent: [f64; MDEPTH],
    /// 线源函数 (每深度点)
    pub slin: [f64; MDEPTH],
}

impl Default for NlteOutput {
    fn default() -> Self {
        Self {
            abcent: [0.0; MDEPTH],
            slin: [0.0; MDEPTH],
        }
    }
}

/// NLTE 控制过程。
///
/// 计算线中心不透明度 (ABCENT) 和线源函数 (SLIN)。
///
/// # 参数
///
/// * `params` - NLTE 参数
///
/// # 返回值
///
/// NLTE 输出结果
pub fn nlte(params: &NlteParams) -> NlteOutput {
    let NlteParams {
        il,
        ilw,
        iun,
        gi,
        gj,
        nd,
        temp,
        elec,
        dens,
        dm,
        vturb,
        popul,
        g,
        enion,
        rrr,
        abstd,
        abstdw,
        pnlt,
        amas,
        enev,
        nnext,
        iel,
        freq0,
        excl0,
        excu0,
        gf0,
        iat,
        ion,
        ilnlt,
        ifwin,
        ijcont,
    } = *params;

    let mut output = NlteOutput::default();

    // 检查统计权重
    if gi <= 0.0 || gj <= 0.0 {
        return output;
    }

    // 常量
    let bn = 1.4743e-2; // 2*h/c^3
    let hk = 4.79928144e-11; // h/k
    let un = 1.0;

    let egf = gf0.exp();
    let bnu = bn * (freq0 * 1.0e-15).powi(3);
    let dp0 = 3.33564e-11 * freq0;
    let dp1 = 1.651e8 / amas[iat - 1];

    if ilw > 0 {
        // 显式能级之间的跃迁
        let nki = nnext[iel[ilw - 1] as usize - 1] as usize;

        for id in 0..nd {
            let t = temp[id];
            let mut cor = 1.0;
            let pp = pnlt[iat - 1][ion - 1][id];

            // 下能级 population
            let pi = if ilw > 0 {
                popul[ilw - 1][id] / g[ilw - 1]
            } else {
                pp * ((enev[iat - 1][ion - 1] * 8067.6 * 1.4387886 - excl0) / t).exp()
            };

            // 上能级 population
            let pj = if iun > 0 {
                let p = popul[iun - 1][id] / g[iun - 1];
                cor = ((excu0 - excl0
                    + (enion[iun - 1] - enion[ilw - 1]) / BOLK)
                    / t)
                    .exp();
                p
            } else {
                pp * ((enev[iat - 1][ion - 1] * 8067.6 * 1.4387886 - excu0) / t).exp()
            };

            let x = if pj > 0.0 {
                pi / pj * cor
            } else {
                un
            };

            let x = if x == un {
                (4.79928e-11 * freq0 / t).exp()
            } else {
                x
            };

            let dop = dp0 * (dp1 * t + vturb[id]).sqrt();
            output.slin[id] = bnu / (x - un);

            if pi > 0.0 {
                output.abcent[id] = pi * (un - un / x) * egf / dop;
            }
        }
    } else {
        // 近似 NLTE 共振线 - 二阶逃逸概率理论
        let alm = 2.997925e17 / freq0;
        let hkf = hk * freq0;

        for id in 0..nd {
            let t = temp[id];
            let dop = dp0 * (dp1 * t + vturb[id]).sqrt();
            let x = (hkf / t).exp();

            output.abcent[id] = egf * (-excl0 / t).exp() * rrr[iat - 1][ion - 1][id]
                / dop
                * (1.0 - 1.0 / x);

            let mut ab = abstd[id] + output.abcent[id] * 1.77245;
            if ifwin > 0 {
                ab = abstdw[ijcont][id] + output.abcent[id] * 1.77245;
            }

            // 计算光学深度
            let abm;
            if id == 0 {
                abm = ab / dens[0];
                let tau = 0.5 * dm[0] * abm;
                // 近似 epsilon (Kastner 方法)
                let e = eps(t, elec[id], alm, ion as i32, iun as i32);
                let xk2 = xk2dop(tau);
                output.slin[id] = (e / (e + (1.0 - e) * xk2)).sqrt() * bnu / (x - 1.0);
            } else {
                abm = ab / dens[id];
                let tau = 0.5 * (dm[id] - dm[id - 1]) * (abm + ab / dens[id - 1]);
                // 近似 epsilon (Kastner 方法)
                let e = eps(t, elec[id], alm, ion as i32, iun as i32);
                let xk2 = xk2dop(tau);
                output.slin[id] = (e / (e + (1.0 - e) * xk2)).sqrt() * bnu / (x - 1.0);
            }
        }
    }

    output
}

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

    #[test]
    fn test_nlte_zero_weights() {
        let params = NlteParams {
            il: 1,
            ilw: 1,
            iun: 2,
            gi: 0.0, // 零统计权重
            gj: 2.0,
            nd: 3,
            temp: &[10000.0; MDEPTH],
            elec: &[1.0e14; MDEPTH],
            dens: &[1.0e-10; MDEPTH],
            dm: &[1.0; MDEPTH],
            vturb: &[0.0; MDEPTH],
            popul: &[[0.0; MDEPTH]; MLEVEL],
            g: &[1.0; MLEVEL],
            enion: &[0.0; MLEVEL],
            rrr: &[[[0.0; MDEPTH]; 90]; 30],
            abstd: &[0.0; MDEPTH],
            abstdw: &[[0.0; MDEPTH]],
            pnlt: &[[[0.0; MDEPTH]; 90]; 30],
            amas: &[1.0],
            enev: &[[0.0; 90]; 30],
            nnext: &[0],
            iel: &[0],
            freq0: 3.0e15,
            excl0: 0.0,
            excu0: 0.0,
            gf0: 0.0,
            iat: 1,
            ion: 1,
            ilnlt: 1,
            ifwin: 0,
            ijcont: 0,
        };

        let output = nlte(&params);

        // 当统计权重为零时，应该返回零值
        for id in 0..3 {
            assert_eq!(output.abcent[id], 0.0);
            assert_eq!(output.slin[id], 0.0);
        }
    }
}