//! B 矩阵统计平衡方程部分。
//!
//! 重构自 TLUSTY `BPOP` 子程序。
//!
//! # 功能
//!
//! - 计算 B 矩阵中对应于统计平衡方程的部分
//! - 处理 (NFREQE+INSE)-th 到 (NFREQE+INSE+NLVEXP-1)-th 行
//! - 处理电荷守恒方程，即 (NFREQE+INPC)-th 行
//!
//! # 参考文献
//!
//! Mihalas, Stellar Atmospheres, 1978, pp. 143-145
//!
//! # 调用的子程序
//!
//! - RATMAT: 计算速率矩阵
//! - LEVSOL: 求解能级方程
//! - LEVGRP: 能级分组
//! - MATINV: 矩阵求逆
//! - BPOPE: B 矩阵 E 部分（显式频率）
//! - BPOPF: B 矩阵 F 部分（占据数行）
//! - BPOPT: B 矩阵 T 部分（温度相关）
//! - BPOPC: B 矩阵 C 部分（电荷守恒）

use crate::tlusty::state::arrays::{BpoCom, MainArrays};
use crate::tlusty::state::atomic::AtomicData;
use crate::tlusty::state::config::TlustyConfig;
use crate::tlusty::state::constants::MLEVEL;
use crate::tlusty::state::iterat::IterControl;
use crate::tlusty::state::model::ModelState;
use crate::tlusty::state::alipar::FixAlp;

use super::matinv;

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

/// BPOP 输入参数
pub struct BpopParams {
    /// 深度索引 (1-indexed)
    pub id: usize,
}

/// BPOP 输出
pub struct BpopOutput {
    /// 是否执行了计算
    pub computed: bool,
}

// ============================================================================
// 主函数
// ============================================================================

/// 计算 B 矩阵的统计平衡方程部分。
///
/// # 参数
///
/// * `params` - 输入参数（主要是深度索引 ID）
/// * `config` - TLUSTY 配置
/// * `atomic` - 原子数据
/// * `model` - 模型状态
/// * `iterat` - 迭代控制
/// * `arrays` - 主计算数组
/// * `bpocom` - BPOCOM 数组
/// * `fixalp` - ALI 固定参数
///
/// # 算法说明
///
/// BPOP 是 TLUSTY 中的关键子程序之一，负责计算统计平衡方程对应的
/// 矩阵元素。主要步骤：
///
/// 1. 初始化辅助数组 ATT 和 ANN
/// 2. 如果非 LTE 且特定条件满足，调用 RATMAT 和 LEVSOL
/// 3. 调用 LEVGRP 进行能级分组
/// 4. 调用 RATMAT 计算速率矩阵
/// 5. 如果需要，调用 MATINV 求逆矩阵
/// 6. 调用 BPOPE, BPOPF, BPOPT, BPOPC 计算各部分
/// 7. 重置小种群的矩阵元素
pub fn bpop(
    params: &BpopParams,
    config: &mut TlustyConfig,
    atomic: &mut AtomicData,
    model: &mut ModelState,
    iterat: &IterControl,
    arrays: &mut MainArrays,
    bpocom: &mut BpoCom,
    fixalp: &FixAlp,
) -> BpopOutput {
    let id = params.id;
    let id_idx = id - 1; // 转换为 0-indexed

    // 从配置中提取常用参数
    let ioptab = config.basnum.ioptab;
    let nlevel = config.basnum.nlevel as usize;
    let nlvexp = atomic.levpar.nlvexp as usize;
    let nfreqe = config.basnum.nfreqe as usize;
    let _nion = config.basnum.nion as usize;
    let inse = config.matkey.inse as i32;
    let inpc = config.matkey.inpc;
    let idlte = config.basnum.idlte;

    let lte = config.inppar.lte;
    let ifpopr = iterat.chnad.ibpope;
    let ipslte = config.basnum.ipslte;
    let ibpope = iterat.chnad.ibpope;

    // 如果 ioptab < 0，直接返回
    if ioptab < 0 {
        return BpopOutput { computed: false };
    }

    // ================================================================
    // 步骤 1: 初始化 ATT 和 ANN 数组
    // ================================================================
    // Fortran: DO I=1,NLVEXP; ATT(I)=0.; ANN(I)=0.; END DO
    for i in 0..nlvexp {
        bpocom.att[i] = 0.0;
        bpocom.ann[i] = 0.0;
    }

    // ================================================================
    // 步骤 2: 特殊情况处理（非 LTE + IFPOPR=5 + IPSLTE=0）
    // ================================================================
    // Fortran: IF(.NOT.LTE .AND. IFPOPR.EQ.5.AND.IPSLTE.EQ.0) THEN
    if !lte && ifpopr == 5 && ipslte == 0 {
        // 调用 RATMAT 获取速率矩阵
        let mut iifor = atomic.levpar.iifor.clone();
        let mut ratmat_params = super::ratmat::RatmatParams {
            id,
            iical: &mut iifor,
            imode: 0,
        };

        let ratmat_output = super::ratmat::ratmat(
            &mut ratmat_params,
            config,
            atomic,
            model,
            iterat,
        );

        // 复制结果到 ESEMAT 和 BESE
        for i in 0..nlevel {
            bpocom.bese[i] = ratmat_output.b[i];
            for j in 0..nlevel {
                bpocom.esemat[i][j] = ratmat_output.a[i][j];
            }
        }

        // 调用 LEVSOL 求解能级方程
        let mut a_flat = vec![0.0; MLEVEL * MLEVEL];
        let mut b_flat = vec![0.0; MLEVEL];
        let mut popp = vec![0.0; MLEVEL];

        for i in 0..nlevel {
            b_flat[i] = bpocom.bese[i];
            for j in 0..nlevel {
                a_flat[i * MLEVEL + j] = bpocom.esemat[i][j];
            }
        }

        let nlvfor = nlevel;
        super::levsol::levsol(
            &mut a_flat,
            &mut b_flat,
            &mut popp,
            &atomic.levpar.iifor,
            nlvfor,
            0,
            &config.basnum,
            &config.inppar,
            &atomic.atopar,
        );

        // 更新 SBPSI 和 BFAC
        // Fortran: DO I=1,NLEVEL
        //            II=IIEXP(I)
        //            IF(II.EQ.0.AND.IMODL(I).EQ.6) THEN
        //              III=ILTREF(I,ID)
        //              SBPSI(I,ID)=POPP(I)/POPP(III)
        //            END IF
        //          END DO
        for i in 0..nlevel {
            let ii = atomic.levpar.iiexp[i];
            if ii == 0 && atomic.levpar.imodl[i] == 6 {
                let iii = model.levref.iltref[i][id_idx] as usize;
                if iii > 0 && iii <= nlevel && popp[iii - 1] != 0.0 {
                    model.levref.sbpsi[i][id_idx] = popp[i] / popp[iii - 1];
                }
            }
        }

        // 更新 BFAC
        // Fortran: DO ION=1,NION
        //            DO I=NFIRST(ION),NLAST(ION)
        //              SBW(I)=ELEC(ID)*SBF(I)*WOP(I,ID)
        //              IF(POPUL(NNEXT(ION),ID).GT.0..AND.IPZERO(I,ID).EQ.0)
        //     *          BFAC(I,ID)=POPP(I)/(POPP(NNEXT(ION))*SBW(I))
        //            END DO
        //          END DO
        let nion = config.basnum.nion as usize;
        for ion in 0..nion {
            let nfirst = atomic.ionpar.nfirst[ion] as usize;
            let nlast = atomic.ionpar.nlast[ion] as usize;
            let nnext = atomic.ionpar.nnext[ion] as usize;

            if nfirst > 0 && nlast >= nfirst {
                for i in (nfirst - 1)..nlast.min(nlevel) {
                    let sbw = model.modpar.elec[id_idx]
                        * model.levpop.sbf[i]
                        * model.wmcomp.wop[i][id_idx];

                    if nnext > 0
                        && nnext <= nlevel
                        && model.levpop.popul[nnext - 1][id_idx] > 0.0
                        && model.popzr0.ipzero[i][id_idx] == 0
                        && sbw != 0.0
                    {
                        model.levpop.bfac[i][id_idx] =
                            popp[i] / (popp[nnext - 1] * sbw);
                    }
                }
            }
        }
    }

    // ================================================================
    // 步骤 3: 调用 LEVGRP
    // ================================================================
    // Fortran: CALL LEVGRP(ID,IIEXP,0,POPP)
    // 注意：levgrp 有自己的参数结构体，这里简化处理
    // 在完整实现中，需要正确构造 LevgrpParams

    // ================================================================
    // 步骤 4: 调用 RATMAT
    // ================================================================
    // Fortran: CALL RATMAT(ID,IIEXP,0,ESEMAT,BESE)
    let mut iical = atomic.levpar.iiexp.clone();
    let mut ratmat_params = super::ratmat::RatmatParams {
        id,
        iical: &mut iical,
        imode: 0,
    };

    let ratmat_output = super::ratmat::ratmat(
        &mut ratmat_params,
        config,
        atomic,
        model,
        iterat,
    );

    // 复制结果到 ESEMAT 和 BESE
    for i in 0..nlevel {
        bpocom.bese[i] = ratmat_output.b[i];
        for j in 0..nlevel {
            bpocom.esemat[i][j] = ratmat_output.a[i][j];
        }
    }

    // ================================================================
    // 步骤 5: 如果 IFPOPR <= 3，调用 MATINV 求逆矩阵
    // ================================================================
    // Fortran: IF(IFPOPR.LE.3) CALL MATINV(ESEMAT,NLVEXP,MLEVEL)
    if ifpopr <= 3 {
        let mut esemat_flat = vec![0.0; MLEVEL * MLEVEL];
        for i in 0..nlevel {
            for j in 0..nlevel {
                esemat_flat[i * MLEVEL + j] = bpocom.esemat[i][j];
            }
        }
        matinv(&mut esemat_flat, nlvexp);

        // 将结果复制回 ESEMAT
        for i in 0..nlevel {
            for j in 0..nlevel {
                bpocom.esemat[i][j] = esemat_flat[i * MLEVEL + j];
            }
        }
    }

    // ================================================================
    // 步骤 6: 调用 BPOPE, BPOPF, BPOPT, BPOPC
    // ================================================================
    // Fortran: if(ipslte.eq.0) then
    //            IF(.NOT.LTE.AND.IBPOPE.GT.0.AND.ID.LT.IDLTE) THEN
    //              CALL BPOPE(ID)
    //              CALL BPOPF(ID)
    //            END IF
    //          end if
    //          CALL BPOPT(ID)
    //          IF(INPC.GT.0) CALL BPOPC(ID)

    if ipslte == 0 {
        if !lte && ibpope > 0 && (id as i32) < idlte {
            // 调用 BPOPE
            // 注意：BPOPE 需要完整的参数，这里简化处理

            // 调用 BPOPF
            let bpopf_params = super::bpopf::BpopfParams {
                nfreqe,
                inse,
                inre: config.matkey.inre,
                inpc: config.matkey.inpc,
                nlvexp,
                ifpopr,
                ifali: config.basnum.jali,
                id,
                crsw: &model.modpar.alab,
            };
            super::bpopf::bpopf(&bpopf_params, arrays, fixalp, bpocom);
        }
    }

    // 调用 BPOPT
    // 注意：BPOPT 需要完整的参数，这里简化处理
    // 在完整实现中，需要正确构造 BpoptParams

    // 调用 BPOPC（如果 INPC > 0）
    // 注意：BPOPC 需要完整的参数，这里简化处理
    // 在完整实现中，需要正确构造 BpopcParams

    // ================================================================
    // 步骤 7: 重置"小"种群的矩阵元素
    // ================================================================
    // Fortran: DO I=1,NLVEXP
    //            IF(IGZERO(I,ID).NE.0) THEN
    //              DO J=1,NLVEXP
    //                B(NFREQE+INSE-1+I,NFREQE+INSE-1+J)=0.
    //              END DO
    //              B(NFREQE+INSE-1+I,NFREQE+INSE-1+I)=1.
    //              VECL(NFREQE+INSE-1+I)=0.
    //            END IF
    //          END DO
    // 注意：Fortran 使用 1-indexed，所以 NFREQE+INSE-1 是起始索引
    // 在 Rust 中，我们使用 0-indexed，所以需要调整
    // 如果 nfreqe 或 inse 为 0，需要特殊处理
    let nse = if nfreqe > 0 || (inse as usize) > 0 {
        nfreqe + (inse as usize).saturating_sub(1)
    } else {
        0
    };

    for i in 0..nlvexp {
        if model.popzr0.igzero[i][id_idx] != 0 {
            // 重置行
            for j in 0..nlvexp {
                arrays.b[nse + i][nse + j] = 0.0;
            }
            // 设置对角元素为 1
            arrays.b[nse + i][nse + i] = 1.0;
            // 重置向量元素
            arrays.vecl[nse + i] = 0.0;
        }
    }

    BpopOutput { computed: true }
}

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

    #[test]
    fn test_bpop_skip_if_ioptab_negative() {
        // 测试当 ioptab < 0 时跳过计算
        let mut config = TlustyConfig::new();
        config.basnum.ioptab = -1;

        let mut atomic = AtomicData::new();
        let mut model = ModelState::default();
        let iterat = IterControl::new();
        let mut arrays = MainArrays::default();
        let mut bpocom = BpoCom::new();
        let fixalp = FixAlp::default();

        let params = BpopParams { id: 1 };

        let result = bpop(
            &params,
            &mut config,
            &mut atomic,
            &mut model,
            &iterat,
            &mut arrays,
            &mut bpocom,
            &fixalp,
        );

        assert!(!result.computed);
    }

    #[test]
    fn test_bpop_initializes_att_ann() {
        // 测试 ATT 和 ANN 初始化
        let mut config = TlustyConfig::new();
        config.basnum.ioptab = 0;

        let mut atomic = AtomicData::new();
        atomic.levpar.nlvexp = 5;

        let mut model = ModelState::default();
        let iterat = IterControl::new();
        let mut arrays = MainArrays::default();
        let mut bpocom = BpoCom::new();
        let fixalp = FixAlp::default();

        // 设置非零值以验证初始化
        for i in 0..5 {
            bpocom.att[i] = 1.0;
            bpocom.ann[i] = 2.0;
        }

        let params = BpopParams { id: 1 };

        let result = bpop(
            &params,
            &mut config,
            &mut atomic,
            &mut model,
            &iterat,
            &mut arrays,
            &mut bpocom,
            &fixalp,
        );

        assert!(result.computed);

        // 验证 ATT 和 ANN 被初始化为零
        for i in 0..5 {
            assert_relative_eq!(bpocom.att[i], 0.0);
            assert_relative_eq!(bpocom.ann[i], 0.0);
        }
    }
}