//! 氢碰撞速率计算。
//!
//! 重构自 TLUSTY `colh.f`
//!
//! 计算氢的碰撞电离和碰撞激发速率。
//! 标准表达式来自 Mihalas, Heasley, and Auer (1975)。

use crate::tlusty::math::butler;
use crate::tlusty::math::ceh12;
use crate::tlusty::math::cspec;
use crate::tlusty::math::irc;
use crate::tlusty::data::{COLH_CCOOL, COLH_CHOT};
use crate::tlusty::state::constants::{EH, HK, MLEVEL, TWO, UN};
// f2r_depends: BUTLER, CSPEC, IRC

// 物理常量
const CC0: f64 = 5.465e-11;
const CEX1: f64 = -30.20581;
const CEX2: f64 = 3.8608704;
const CEX3: f64 = 305.63574;
const ALF0: f64 = 1.8;
const ALF1: f64 = 0.4;
const BET0: f64 = 3.0;
const BET1: f64 = 1.2;
const O148: f64 = 0.148;
const CHMI: f64 = 5.59e-15;

// 指数积分系数
const EXPIA1: f64 = -0.57721566;
const EXPIA2: f64 = 0.99999193;
const EXPIA3: f64 = -0.24991055;
const EXPIA4: f64 = 0.05519968;
const EXPIA5: f64 = -0.00976004;
const EXPIA6: f64 = 0.00107857;
const EXPIB1: f64 = 0.2677734343;
const EXPIB2: f64 = 8.6347608925;
const EXPIB3: f64 = 18.059016973;
const EXPIB4: f64 = 8.5733287401;
const EXPIC1: f64 = 3.9584969228;
const EXPIC2: f64 = 21.0996530827;
const EXPIC3: f64 = 25.6329561486;
const EXPIC4: f64 = 9.5733223454;

/// COLH 输入参数
pub struct ColhParams {
    /// 深度索引 (1-indexed)
    pub id: usize,
    /// 温度 (K)
    pub t: f64,
    /// 碰撞速率选择标志
    pub icolhn: i32,
}

/// COLH 原子数据
pub struct ColhAtomicData<'a> {
    /// 氢元素索引
    pub ielh: usize,
    /// H- 元素索引 (0 表示没有 H-)
    pub ielhm: usize,
    /// 氢原子索引
    pub iath: usize,
    /// 能级的第一个索引 (nelem)
    pub nfirst: &'a [i32],
    /// 能级的最后一个索引 (nelem)
    pub nlast: &'a [i32],
    /// 电离能级索引 (natom)
    pub nka: &'a [i32],
    /// 主量子数 (nlevel)
    pub nquant: &'a [i32],
    /// 上能级截止 (nelem)
    pub icup: &'a [i32],
    /// 跃迁索引数组 (nlevel × nlevel)
    pub itra: &'a [i32],
    /// 碰撞速率标志 (ntrans)
    pub icol: &'a [i32],
    /// 跃迁频率 (ntrans)
    pub fr0: &'a [f64],
    /// 振子强度 (ntrans)
    pub osc0: &'a [f64],
    /// 碰撞参数 (ntrans)
    pub cpar: &'a [f64],
    /// 电离能 (nlevel)
    pub enion: &'a [f64],
    /// 能级宽度选项 (nlevel)
    pub ifwop: &'a [i32],
    /// WNHINT 数组 (nlmx × nd)
    pub wnhint: &'a [f64],
    /// OSH 振子强度 (10 × 20)
    pub osh: &'a [f64],
    /// 最大谱线数
    pub nlmx: usize,
}

/// COLH 输出
pub struct ColhOutput<'a> {
    /// 碰撞速率数组 (ntrans)
    pub col: &'a mut [f64],
}

// A 系数数组（用于高 n 碰撞电离）
static A: [[f64; 10]; 6] = [
    [-86.7633398, 2632.8369, 7478.9556, -4202.8442, -47995.930,
     -120942.89, -202300.81, -261373.03, -266337.91, -192293.20],
    [100.919188, -2738.7485, -8495.4590, 1937.3763, 45825.371,
     122209.39, 211928.67, 285044.75, 309455.47, 258802.22],
    [-45.7813807, 1121.3976, 3794.6826, 340.35764, -16617.055,
     -47390.313, -84973.688, -117833.95, -133243.61, -120363.95],
    [10.1978559, -224.30670, -822.83636, -290.10489, 2905.7393,
     8944.6025, 16556.992, 23544.543, 27419.742, 26002.143],
    [-1.11223557, 21.923729, 86.619110, 48.840523, -246.99014,
     -828.41028, -1581.2722, -2297.9321, -2738.1743, -2686.4087],
    [0.0474198818, -0.83974838, -3.5534720, -2.6097214, 8.1972208,
     30.267115, 59.521984, 88.178680, 107.05288, 107.73775],
];

/// 计算氢的碰撞速率。
///
/// # 参数
///
/// * `params` - 输入参数
/// * `atomic` - 原子数据
/// * `output` - 输出碰撞速率
pub fn colh(params: &ColhParams, atomic: &ColhAtomicData, output: &mut ColhOutput) {
    let t = params.t;
    let id = params.id;
    let id_idx = id - 1;

    let hkt = HK / t;
    let ct = CC0 * t.sqrt();
    let tk = hkt / EH;
    let x = t.log10();
    let x2 = x * x;
    let x3 = x * x2;
    let x4 = x2 * x2;
    let x5 = x3 * x2;
    let sqt = t.sqrt();

    let xtt = [1.0, t, t * t, t * t * t];

    let n0hn = atomic.nfirst[atomic.ielh] as usize - 1;
    let n1h = atomic.nlast[atomic.ielh] as usize - 1;
    let nkh = atomic.nka[atomic.iath] as usize - 1;
    let n1q = if atomic.icup[atomic.ielh] > 0 {
        atomic.icup[atomic.ielh] as usize
    } else {
        0
    };
    let n1hc = if atomic.ifwop[n1h] < 0 { n1h - 1 } else { n1h };
    let mut nhl = if n1q > 0 {
        n1q
    } else {
        n1h - n0hn + 1
    };
    if atomic.ifwop[n1h] < 0 {
        nhl = atomic.nlmx;
    }

    for ii in n0hn..=n1h {
        let i = ii - n0hn + 1;
        let it = atomic.itra[ii + MLEVEL * nkh] as usize;
        if it == 0 {
            continue;
        }

        // 碰撞电离
        if t > 1e6 {
            // 高温使用 XSTAR 公式
            let rno = 16.0;
            let izc = 1;
            let cs = irc(i as i32, t, izc, rno);
            output.col[it - 1] = cs;
        } else {
            let ic = atomic.icol[it - 1];
            let u0 = atomic.fr0[it - 1] * hkt;

            if ic < 0 {
                // 非标准公式
                output.col[it - 1] = cspec(
                    ii as i32, nkh as i32, ic,
                    atomic.osc0[it - 1], atomic.cpar[it - 1], u0, t
                );
            } else if atomic.ifwop[ii] < 0 {
                // 合并态电离
                let ehk = EH / tk;
                let n00q = atomic.nquant[n1h - 1] as usize + 1;
                let mut sum1 = 0.0;
                let mut sum2 = 0.0;
                for img in n00q..=atomic.nlmx {
                    let xi = img as f64;
                    let xii = xi * xi;
                    sum1 += xii * xii * xi * atomic.wnhint[(img - 1) + atomic.nlmx * id_idx];
                    sum2 += xii * atomic.wnhint[(img - 1) + atomic.nlmx * id_idx] * (ehk / xii).exp();
                }
                output.col[it - 1] = ct * sum1 / sum2;
            } else {
                // 标准公式
                let gam = if i > 10 {
                    (i * i * i) as f64
                } else {
                    A[0][i - 1] + A[1][i - 1] * x + A[2][i - 1] * x2
                        + A[3][i - 1] * x3 + A[4][i - 1] * x4 + A[5][i - 1] * x5
                };
                output.col[it - 1] = ct * (-u0).exp() * gam;
            }
        }

        // 碰撞激发
        let i1 = i + 1;
        if i1 > nhl {
            continue;
        }

        let xi = i as f64;
        let vi = xi * xi;
        let alf = ALF0 - ALF1 / vi;
        let bet = BET0 - BET1 / xi;
        let csca = 8.63e-6 / 2.0 / vi / sqt;
        let mut csum = 0.0;

        for j in i1..nhl {
            let xj = j as f64;
            let vj = xj * xj;
            let jj = j + n0hn;

            let (cs, is_lumped) = if jj > n1hc {
                // 非显式能级，归入碰撞电离
                let e = UN / vi - UN / vj;
                let u0 = EH * e * tk;
                let c1 = if j <= 20 {
                    atomic.osh[(i - 1) + 20 * (j - 1)]
                } else {
                    atomic.osh[(i - 1) + 20 * 19] * ((400.0 - vi) / 20.0 * xj / (vj - vi)).powi(3)
                };
                (compute_collision_rate(params.icolhn, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt), true)
            } else {
                let ict = atomic.itra[ii + MLEVEL * (jj - 1)] as usize;
                if ict == 0 {
                    continue;
                }
                let ic = atomic.icol[ict - 1];
                let u0 = atomic.fr0[ict - 1] * hkt;
                let c1 = atomic.osc0[ict - 1];

                let cs = if ic < 0 {
                    cspec(ii as i32, jj as i32, ic, c1, atomic.cpar[ict - 1], u0, t)
                } else if ic == 0 {
                    compute_collision_rate(params.icolhn, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt)
                } else if ic == 1 {
                    ct * (-u0).exp() * (CEX1 + CEX2 * x + CEX3 / x / x)
                } else {
                    ceh12(t)
                };
                (cs, false)
            };

            if is_lumped {
                csum += cs;
            } else {
                let ict = atomic.itra[ii + MLEVEL * (jj - 1)] as usize;
                if ict > 0 {
                    output.col[ict - 1] = cs;
                }
            }
        }

        // 添加累积碰撞速率
        let it_main = atomic.itra[ii + MLEVEL * nkh] as usize;
        if it_main > 0 && n1q > 0 {
            output.col[it_main - 1] += csum;
        }
        let ith = atomic.itra[ii + MLEVEL * n1h] as usize;
        if atomic.ifwop[n1h] < 0 && ith > 0 {
            output.col[ith - 1] = csum;
        }
    }

    // H- 碰撞电离
    if atomic.ielhm > 0 {
        let it_hm = atomic.itra[(atomic.nfirst[atomic.ielhm] as usize - 1) + MLEVEL * n0hn] as usize;
        if it_hm > 0 {
            let ic = atomic.icol[it_hm - 1];
            if ic >= 0 {
                output.col[it_hm - 1] = CHMI * t * t.sqrt();
            } else {
                let u0 = atomic.enion[atomic.nfirst[atomic.ielhm] as usize - 1] * tk;
                output.col[it_hm - 1] = cspec(
                    atomic.nfirst[atomic.ielhm] as i32,
                    n0hn as i32,
                    ic,
                    atomic.osc0[it_hm - 1],
                    atomic.cpar[it_hm - 1],
                    u0,
                    t,
                );
            }
        }
    }
}

/// 计算碰撞激发速率
fn compute_collision_rate(
    icolhn: i32,
    i: usize,
    j: usize,
    t: f64,
    u0: f64,
    c1: f64,
    csca: f64,
    alf: f64,
    bet: f64,
    xi: f64,
    xj: f64,
    xtt: &[f64],
) -> f64 {
    // 使用 Butler 新计算
    if icolhn == 1 && j <= 7 {
        let (cs, ierr) = butler(i as i32, j as i32, t, u0);
        if ierr == 0 {
            return cs;
        }
    }

    // Giovanardi 公式
    if icolhn == 2 && j <= 15 {
        let cs = if t <= 60000.0 {
            get_ccool(i, j, xtt)
        } else {
            get_chot(i, j, xtt)
        };
        return csca * cs * (-u0).exp();
    }

    // 标准公式 (Mihalas et al 1975)
    let ct = CC0 * t.sqrt();
    let e = u0 / (HK / t);
    let cs = 4.0 * ct * c1 / (e * e);
    let ex = (-u0).exp();

    let e1 = if u0 <= UN {
        -u0.ln() + EXPIA1 + u0 * (EXPIA2 + u0 * (EXPIA3 + u0 * (EXPIA4 + u0 * (EXPIA5 + u0 * EXPIA6))))
    } else {
        let u0_inv = 1.0 / u0;
        (-u0).exp() * ((EXPIB1 + u0 * (EXPIB2 + u0 * (EXPIB3 + u0 * EXPIB4)))
            / (EXPIC1 + u0 * (EXPIC2 + u0 * (EXPIC3 + u0 * EXPIC4))))
            * u0_inv
    };

    let mut e5 = e1;
    for ix in 1..=4 {
        e5 = (ex - u0 * e5) / ix as f64;
    }

    let mut cs = cs * u0 * (e1 + O148 * u0 * e5);
    if j - i != 1 {
        cs *= bet + TWO * (alf - bet) / (xj - xi);
    }
    cs
}

/// 获取 CCOOL 系数（从 data.rs 获取）
/// CCOOL(I, J, K) - I=1..4, J=1..14, K=1..15
/// Fortran 列优先: idx = (I-1) + 4*((J-1) + 14*(K-1))
fn get_ccool(i: usize, j: usize, xtt: &[f64]) -> f64 {
    // Giovanardi 公式: CS = sum_{ica=1}^{4} CCOOL(ica, i, j) * XTT(ica)
    let mut cs = 0.0;
    for ica in 0..4 {
        // Fortran: CCOOL(ica+1, i, j)
        // 列优先索引: (ica) + 4*((i-1) + 14*(j-1))
        let idx = ica + 4 * ((i - 1) + 14 * (j - 1));
        if idx < COLH_CCOOL.len() {
            cs += COLH_CCOOL[idx] * xtt[ica];
        }
    }
    cs
}

/// 获取 CHOT 系数（从 data.rs 获取）
/// CHOT(I, J, K) - I=1..4, J=1..14, K=1..15
/// Fortran 列优先: idx = (I-1) + 4*((J-1) + 14*(K-1))
fn get_chot(i: usize, j: usize, xtt: &[f64]) -> f64 {
    // Giovanardi 公式: CS = sum_{ica=1}^{4} CHOT(ica, i, j) * XTT(ica)
    let mut cs = 0.0;
    for ica in 0..4 {
        // Fortran: CHOT(ica+1, i, j)
        // 列优先索引: (ica) + 4*((i-1) + 14*(j-1))
        let idx = ica + 4 * ((i - 1) + 14 * (j - 1));
        if idx < COLH_CHOT.len() {
            cs += COLH_CHOT[idx] * xtt[ica];
        }
    }
    cs
}

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

    #[test]
    fn test_colh_constants() {
        assert!((CC0 - 5.465e-11).abs() < 1e-20);
        assert!((ALF0 - 1.8).abs() < 1e-10);
        assert!((BET0 - 3.0).abs() < 1e-10);
    }

    #[test]
    fn test_get_ccool_returns_data() {
        // 测试 get_ccool 函数是否正确返回数据
        let xtt = [1.0, 10000.0, 1e8, 1e12]; // T=10000K 的幂次

        // 测试几个有效索引
        for i in 1..=14 {
            for j in 1..=15 {
                let cs = get_ccool(i, j, &xtt);
                // 返回值应该是有限数
                assert!(cs.is_finite(), "get_ccool({}, {}) returned non-finite: {}", i, j, cs);
            }
        }
    }

    #[test]
    fn test_get_chot_returns_data() {
        // 测试 get_chot 函数是否正确返回数据
        let xtt = [1.0, 100000.0, 1e10, 1e15]; // T=100000K 的幂次

        // 测试几个有效索引
        for i in 1..=14 {
            for j in 1..=15 {
                let cs = get_chot(i, j, &xtt);
                // 返回值应该是有限数
                assert!(cs.is_finite(), "get_chot({}, {}) returned non-finite: {}", i, j, cs);
            }
        }
    }

    #[test]
    fn test_get_ccool_temperature_scaling() {
        // 验证 CCOOL 系数随温度的缩放
        let xtt_low = [1.0, 5000.0, 2.5e7, 1.25e11];
        let xtt_high = [1.0, 50000.0, 2.5e9, 1.25e14];

        for i in 1..=5 {
            for j in 1..=5 {
                let cs_low = get_ccool(i, j, &xtt_low);
                let cs_high = get_ccool(i, j, &xtt_high);

                // 高温应该有更大的系数（大部分情况）
                // 这里只验证都是有限数
                assert!(cs_low.is_finite() && cs_high.is_finite());
            }
        }
    }

    #[test]
    fn test_compute_collision_rate_standard() {
        // 测试标准 Mihalas 公式
        let t: f64 = 10000.0;
        let i = 1;
        let j = 2;
        let u0: f64 = 0.5; // 激发能量 / kT
        let c1: f64 = 0.4162; // Lyman-alpha 振子强度
        let csca = 8.63e-6 / 2.0 / 1.0 / t.sqrt();
        let alf = ALF0 - ALF1 / 1.0;
        let bet = BET0 - BET1 / 1.0;
        let xi = 1.0;
        let xj = 2.0;
        let xtt = [1.0, t, t*t, t*t*t];

        // 使用 icolhn=0 强制使用标准公式
        let cs = compute_collision_rate(0, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt);

        // 碰撞速率应该是有限的（可能很小但应该有效）
        assert!(cs.is_finite(), "Collision rate should be finite");
    }

    #[test]
    fn test_compute_collision_rate_butler() {
        // 测试 Butler 公式 (icolhn=1, j<=7)
        let t: f64 = 10000.0;
        let i = 1;
        let j = 2;
        let u0: f64 = 0.5;
        let c1: f64 = 0.4162;
        let csca = 8.63e-6 / 2.0 / 1.0 / t.sqrt();
        let alf = ALF0 - ALF1 / 1.0;
        let bet = BET0 - BET1 / 1.0;
        let xi = 1.0;
        let xj = 2.0;
        let xtt = [1.0, t, t*t, t*t*t];

        let cs = compute_collision_rate(1, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt);

        assert!(cs > 0.0, "Butler collision rate should be positive");
        assert!(cs.is_finite(), "Butler collision rate should be finite");
    }

    #[test]
    fn test_compute_collision_rate_giovanardi_cool() {
        // 测试 Giovanardi 冷公式 (icolhn=2, j<=15, T<=60000K)
        let t: f64 = 10000.0;
        let i = 1;
        let j = 2;
        let u0: f64 = 0.5;
        let c1: f64 = 0.4162;
        let csca = 8.63e-6 / 2.0 / 1.0 / t.sqrt();
        let alf = ALF0 - ALF1 / 1.0;
        let bet = BET0 - BET1 / 1.0;
        let xi = 1.0;
        let xj = 2.0;
        let xtt = [1.0, t, t*t, t*t*t];

        let cs = compute_collision_rate(2, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt);

        assert!(cs > 0.0, "Giovanardi cool collision rate should be positive");
        assert!(cs.is_finite(), "Giovanardi cool collision rate should be finite");
    }

    #[test]
    fn test_compute_collision_rate_giovanardi_hot() {
        // 测试 Giovanardi 热公式 (icolhn=2, j<=15, T>60000K)
        let t: f64 = 100000.0;
        let i = 1;
        let j = 2;
        let u0: f64 = 0.1; // 更小的激发能量比
        let c1: f64 = 0.4162;
        let csca = 8.63e-6 / 2.0 / 1.0 / t.sqrt();
        let alf = ALF0 - ALF1 / 1.0;
        let bet = BET0 - BET1 / 1.0;
        let xi = 1.0;
        let xj = 2.0;
        let xtt = [1.0, t, t*t, t*t*t];

        let cs = compute_collision_rate(2, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt);

        assert!(cs > 0.0, "Giovanardi hot collision rate should be positive");
        assert!(cs.is_finite(), "Giovanardi hot collision rate should be finite");
    }

    #[test]
    fn test_collision_rate_temperature_dependence() {
        // 验证碰撞速率随温度的变化
        let i = 1;
        let j = 2;
        let u0_base: f64 = 10.0; // 跃迁能量
        let c1: f64 = 0.4162;
        let alf = ALF0 - ALF1 / 1.0;
        let bet = BET0 - BET1 / 1.0;
        let xi = 1.0;
        let xj = 2.0;

        let temperatures: [f64; 4] = [5000.0, 10000.0, 20000.0, 50000.0];
        let mut rates = Vec::new();

        for t in &temperatures {
            let hkt = HK / t;
            let u0 = u0_base * hkt / (HK / 10000.0) * 0.5; // 归一化激发能量
            let csca = 8.63e-6 / 2.0 / 1.0 / t.sqrt();
            let xtt = [1.0, *t, (*t)*(*t), (*t)*(*t)*(*t)];

            let cs = compute_collision_rate(0, i, j, *t, u0, c1, csca, alf, bet, xi, xj, &xtt);
            rates.push(cs);
        }

        // 碰撞速率应该是有限的
        for rate in &rates {
            assert!(rate.is_finite(), "Collision rate should be finite");
        }
    }

    #[test]
    fn test_collision_rate_level_dependence() {
        // 验证碰撞速率随能级的变化
        let t: f64 = 10000.0;
        let c1: f64 = 0.1;
        let xtt = [1.0, t, t*t, t*t*t];

        let mut rates = Vec::new();

        for j in 2..=10 {
            let i = j - 1;
            let xi = i as f64;
            let xj = j as f64;
            let u0 = 0.5; // 简化
            let csca = 8.63e-6 / 2.0 / ((i*i) as f64) / t.sqrt();
            let alf = ALF0 - ALF1 / ((i*i) as f64);
            let bet = BET0 - BET1 / (i as f64);

            let cs = compute_collision_rate(0, i, j, t, u0, c1, csca, alf, bet, xi, xj, &xtt);
            rates.push(cs);
        }

        // 所有速率应该是正的有限数
        for (j, rate) in rates.iter().enumerate() {
            assert!(rate.is_finite() && rate > &0.0,
                "Rate for transition {}->{} should be positive finite", j+1, j+2);
        }
    }

    #[test]
    fn test_exponential_integral_approximation() {
        // 验证指数积分近似
        // 小 u0 使用级数展开，大 u0 使用连分式
        let small_u0 = 0.5_f64;
        let large_u0 = 5.0_f64;

        // 小 u0 的级数展开
        let e1_small = -small_u0.ln() + EXPIA1
            + small_u0 * (EXPIA2 + small_u0 * (EXPIA3
            + small_u0 * (EXPIA4 + small_u0 * (EXPIA5 + small_u0 * EXPIA6))));

        assert!(e1_small > 0.0, "E1 for small u0 should be positive");

        // 大 u0 的连分式
        let u0_inv = 1.0 / large_u0;
        let e1_large = (-large_u0).exp()
            * ((EXPIB1 + large_u0 * (EXPIB2 + large_u0 * (EXPIB3 + large_u0 * EXPIB4)))
            / (EXPIC1 + large_u0 * (EXPIC2 + large_u0 * (EXPIC3 + large_u0 * EXPIC4))))
            * u0_inv;

        assert!(e1_large > 0.0 && e1_large < 1.0, "E1 for large u0 should be between 0 and 1");
    }

    #[test]
    fn test_hm_ionization_rate() {
        // 验证 H- 碰撞电离速率公式
        let t = 10000.0_f64;
        let rate = CHMI * t * t.sqrt();

        // H- 电离速率应该在合理范围内
        assert!(rate > 0.0, "H- ionization rate should be positive");
        assert!(rate.is_finite(), "H- ionization rate should be finite");

        // CHMI = 5.59e-15, 所以 rate ≈ 5.59e-15 * 10000 * 100 = 5.59e-9
        assert!(rate > 1e-12 && rate < 1e-6,
            "H- ionization rate {:.2e} seems unreasonable", rate);
    }
}