//! 不透明度计算的深度无关量初始化。
//!
//! 重构自 TLUSTY `traini.f`

use crate::tlusty::state::atomic::{IonPar, LevPar, TraPar};
use crate::tlusty::state::config::BasNum;
use crate::tlusty::state::constants::MFREQ;
use crate::tlusty::state::model::{CompIf, DwnPar, LinFrq, LinOvr, ObfPar};

// ============================================================================
// TRAINI - 不透明度初始化
// ============================================================================

/// 初始化不透明度计算的深度无关量。
///
/// # 参数
///
/// - `basnum` - 基本数值计数器
/// - `ionpar` - 离子参数
/// - `levpar` - 能级参数
/// - `trapar` - 跃迁参数 (会被修改)
/// - `obfpar` - 束缚-自由参数
/// - `dwnpar` - 溶解分数参数 (会被修改)
/// - `linovr` - 谱线叠加参数 (会被修改)
/// - `linfrq` - 谱线频率参数 (会被修改)
/// - `compif` - 计算标志
///
/// # Fortran 原始代码
///
/// ```fortran
/// SUBROUTINE TRAINI
///   INCLUDE 'IMPLIC.FOR'
///   INCLUDE 'BASICS.FOR'
///   INCLUDE 'ATOMIC.FOR'
///   INCLUDE 'MODELQ.FOR'
///   INCLUDE 'ODFPAR.FOR'
///
///   do itr=1,ntrans
///      idiel(itr)=0
///   end do
///
///   NCDW=0
///   DO 10 IBFT=1,NTRANC
///      ITR=ITRBF(IBFT)
///      ii=ilow(itr)
///      if(ilk(iup(itr)).ne.0.and.nfirst(iel(ii)).eq.ii.
///    *      and.IFDIEL.NE.0) idiel(itr)=1
///      MODW=IABS(INDEXP(ITR))
///      IF(MODW.NE.5.AND.MODW.NE.15) GO TO 10
///      NCDW=NCDW+1
///      MCDW(ITR)=NCDW
///      ITRCDW(NCDW)=ITR
///  10 CONTINUE
///   IF(ISPODF.GE.1) RETURN
///
///   DO IJ=1,NFREQ
///      NLINES(IJ)=0
///   END DO
///
///   DO 100 ITR=1,NTRANS
///      IF(LINEXP(ITR)) GO TO 100
///      DO IJ=IFR0(ITR),IFR1(ITR)
///         IJLIN(IJ)=ITR
///      END DO
///  100 CONTINUE
/// END
/// ```
pub fn traini(
    basnum: &BasNum,
    ionpar: &IonPar,
    levpar: &LevPar,
    trapar: &mut TraPar,
    obfpar: &ObfPar,
    dwnpar: &mut DwnPar,
    linovr: &mut LinOvr,
    linfrq: &mut LinFrq,
    compif: &CompIf,
) {
    let ntrans = basnum.ntrans as usize;
    let ntranc = basnum.ntranc as usize;
    let nfreq = basnum.nfreq as usize;
    let ispodf = basnum.ispodf;
    let ifdiel = basnum.ifdiel;

    // 初始化 idiel
    for itr in 0..ntrans {
        trapar.idiel[itr] = 0;
    }

    // 束缚-自由跃迁处理
    let mut ncdw: i32 = 0;

    for ibft in 0..ntranc {
        let itr = (obfpar.itrbf[ibft] - 1) as usize; // Fortran 1-indexed -> Rust 0-indexed
        let ii = (trapar.ilow[itr] - 1) as usize;
        let iup_idx = (trapar.iup[itr] - 1) as usize;

        // 检查是否是电离能级
        // if(ilk(iup(itr)).ne.0.and.nfirst(iel(ii)).eq.ii.and.IFDIEL.NE.0)
        let iel_idx = (levpar.iel[ii] - 1) as usize;
        if levpar.ilk[iup_idx] != 0
            && ionpar.nfirst[iel_idx] == (ii + 1) as i32
            && ifdiel != 0
        {
            trapar.idiel[itr] = 1;
        }

        // 检查模数
        let modw = trapar.indexp[itr].abs();
        if modw != 5 && modw != 15 {
            continue;
        }

        ncdw += 1;
        dwnpar.mcdw[itr] = ncdw;
        dwnpar.itrcdw[(ncdw - 1) as usize] = (itr + 1) as i32;
    }

    dwnpar.ncdw = ncdw;

    // 如果 ISPODF >= 1，不处理束缚-束缚跃迁
    if ispodf >= 1 {
        return;
    }

    // 束缚-束缚跃迁处理
    // 初始化 nlines
    for ij in 0..nfreq {
        linfrq.nlines[ij] = 0;
    }

    // 处理跃迁
    for itr in 0..ntrans {
        // 如果是经验线，跳过
        if compif.linexp[itr] {
            continue;
        }

        let ifr0 = (trapar.ifr0[itr] - 1) as usize; // Fortran 1-indexed
        let ifr1 = (trapar.ifr1[itr] - 1) as usize;

        for ij in ifr0..=ifr1.min(MFREQ - 1) {
            linovr.ijlin[ij] = (itr + 1) as i32;
        }
    }
}

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

    fn create_test_data() -> (BasNum, IonPar, LevPar, TraPar, ObfPar, DwnPar, LinOvr, LinFrq, CompIf) {
        let mut basnum = BasNum::default();
        basnum.ntrans = 10;
        basnum.ntranc = 5;
        basnum.nfreq = 100;
        basnum.ispodf = 0;
        basnum.ifdiel = 0;

        let ionpar = IonPar::default();
        let mut levpar = LevPar::default();

        // iel[ii] 表示能级 ii 属于哪个离子 (1-indexed)
        // 在实际运行中，iel 从原子数据文件读取
        // 这里模拟有效数据：所有能级属于第一个离子
        for i in 0..10 {
            levpar.iel[i] = 1;
        }

        let mut trapar = TraPar::default();
        let mut obfpar = ObfPar::new();  // 使用 new() 而不是 default()

        // 设置一些跃迁
        for i in 0..5 {
            obfpar.itrbf[i] = (i + 1) as i32;
            trapar.ilow[i] = (i + 1) as i32;
            trapar.iup[i] = (i + 2) as i32;
            trapar.indexp[i] = 5; // 模数 5
            trapar.ifr0[i] = 1;
            trapar.ifr1[i] = 10;
        }

        let dwnpar = DwnPar::default();
        let linovr = LinOvr::default();
        let linfrq = LinFrq::default();
        let compif = CompIf::default();

        (basnum, ionpar, levpar, trapar, obfpar, dwnpar, linovr, linfrq, compif)
    }

    #[test]
    fn test_traini_basic() {
        let (basnum, ionpar, levpar, mut trapar, obfpar, mut dwnpar, mut linovr, mut linfrq, compif) =
            create_test_data();

        traini(
            &basnum, &ionpar, &levpar, &mut trapar, &obfpar, &mut dwnpar, &mut linovr, &mut linfrq,
            &compif,
        );

        // 检查 idiel 被初始化
        for itr in 0..basnum.ntrans as usize {
            assert_eq!(trapar.idiel[itr], 0);
        }
    }

    #[test]
    fn test_traini_ncdw() {
        let (basnum, ionpar, levpar, mut trapar, obfpar, mut dwnpar, mut linovr, mut linfrq, compif) =
            create_test_data();

        traini(
            &basnum, &ionpar, &levpar, &mut trapar, &obfpar, &mut dwnpar, &mut linovr, &mut linfrq,
            &compif,
        );

        // 由于所有跃迁都有 indexp=5，ncdw 应该 > 0
        assert!(dwnpar.ncdw > 0);
    }

    #[test]
    fn test_traini_ispoft() {
        let (mut basnum, ionpar, levpar, mut trapar, obfpar, mut dwnpar, mut linovr, mut linfrq, compif) =
            create_test_data();
        basnum.ispodf = 1; // 跳过束缚-束缚跃迁

        traini(
            &basnum, &ionpar, &levpar, &mut trapar, &obfpar, &mut dwnpar, &mut linovr, &mut linfrq,
            &compif,
        );

        // linfrq.nlines 应该保持为 0
        for ij in 0..basnum.nfreq as usize {
            assert_eq!(linfrq.nlines[ij], 0);
        }
    }
}