//! 电子密度计算器。
//!
//! 重构自 TLUSTY `eldens.f`。
//!
//! 功能：
//! - 计算电子密度和总氢数密度
//! - 使用 Newton-Raphson 方法求解 Saha 方程组
//! - 计算电荷守恒和粒子守恒
//! - 计算内能和熵

use crate::tlusty::math::lineqs;
use crate::tlusty::math::{moleq_pure, MoleqParams, MoleculeEqData};
use crate::tlusty::math::{mpartf, MpartfResult};
use crate::tlusty::math::{state_pure, StateParams};
use crate::tlusty::math::{entene, EnteneParams, EnteneOutput};
use crate::tlusty::state::constants::{BOLK, HMASS, UN, TWO, HALF};

/// ELDENS 配置参数
#[derive(Debug, Clone)]
pub struct EldensConfig {
    /// 是否包含分子 (ifmol)
    pub ifmol: i32,
    /// 分子温度上限 (tmolim)
    pub tmolim: f64,
    /// ioptab 标志
    pub ioptab: i32,
    /// 氢原子索引 (iath)
    pub iath: usize,
    /// 参考原子索引 (iatref)
    pub iatref: usize,
    /// 是否处理 H- (ihm)
    pub ihm: i32,
    /// 是否处理 H2 (ih2)
    pub ih2: i32,
    /// 是否处理 H2+ (ih2p)
    pub ih2p: i32,
    /// 氢配分函数 (pfhyd)
    pub pfhyd: f64,
}

impl Default for EldensConfig {
    fn default() -> Self {
        Self {
            ifmol: 0,
            tmolim: 1e10,
            ioptab: 0,
            iath: 1,
            iatref: 1,
            ihm: 1,
            ih2: 1,
            ih2p: 1,
            pfhyd: 2.0,
        }
    }
}

/// ELDENS 输入参数
pub struct EldensParams<'a> {
    /// 深度点索引 (1-indexed)
    pub id: usize,
    /// 温度 [K]
    pub t: f64,
    /// 总粒子数密度 [cm⁻³]
    pub an: f64,
    /// 总氢丰度因子 (ytot)
    pub ytot: f64,
    /// 参考原子电荷 (qref)
    pub qref: f64,
    /// 参考原子电荷导数 (dqnr)
    pub dqnr: f64,
    /// 平均分子量 (wmy)
    pub wmy: f64,
    /// 配置
    pub config: EldensConfig,
    /// STATE 函数所需的原子数据（简化）
    pub state_params: Option<StateParams<'a>>,
    /// 分子数据（可选）
    pub molecule_data: Option<&'a MoleculeEqData>,
    /// ANATO 计算所需的原子数密度数据
    pub anato_data: Option<AnatoData>,
}

/// ANATO 计算所需的输入数据（对应 Fortran lines 245-256）
#[derive(Debug, Clone)]
pub struct AnatoData {
    /// 氢基态能级索引 n0hn (>0 表示存在)
    pub n0hn: i32,
    /// 氦基态能级索引 n0a(iathe) (>0 表示存在)
    pub n0a_iathe: i32,
    /// 氦是否为显式元素 iathe (>0 表示是)
    pub iathe: i32,
    /// popul(n0hn, id) - 氢基态粒子数（如果 n0hn > 0）
    pub popul_h: f64,
    /// popul(n0a_iathe, id) - 氦基态粒子数（如果 iathe > 0）
    pub popul_he: f64,
    /// dens(id) / wmm(id) / ytot(id) - LTE 估计
    pub lte_estimate: f64,
    /// abndd(2, id) - 氦丰度
    pub abndd_he: f64,
}

/// ELDENS 输出结果
#[derive(Debug, Clone)]
pub struct EldensOutput {
    /// 电子密度 [cm⁻³]
    pub ane: f64,
    /// 质子数密度 [cm⁻³]
    pub anp: f64,
    /// 总氢数密度 [cm⁻³]
    pub ahtot: f64,
    /// H2 分子相对数量
    pub ahmol: f64,
    /// H- 数密度 [cm⁻³]
    pub anhm: f64,
    /// 内能 [erg]
    pub energ: f64,
    /// 熵 [erg/K]
    pub entt: f64,
    /// 平均分子量
    pub wm: f64,
    /// 质量密度 [g/cm³]
    pub rhoter: f64,
    /// 电子密度与总粒子数密度之比
    pub anerel: f64,
    /// ANATO(1,ID) - 氢的原子数密度（可选，仅当 ifmol<=0 或 t>=tmolim）
    pub anato_h: Option<f64>,
    /// ANATO(2,ID) - 氦的原子数密度（可选，仅当 ifmol<=0 或 t>=tmolim）
    pub anato_he: Option<f64>,
    /// WMM 更新值: rhoter/(an-ane)，当 ipri>0 时调用者应写入 wmm(id)
    pub wmm_update: f64,
}

/// 计算电子密度（纯计算函数）。
///
/// # 参数
/// * `params` - 输入参数
/// * `ipri` - 是否更新全局状态（>0 表示更新）
///
/// # 返回值
/// 包含电子密度、内能、熵等的输出结构体
pub fn eldens_pure(params: &EldensParams, ipri: i32) -> EldensOutput {
    // 检查 ioptab 标志
    if params.config.ioptab < -1 {
        return EldensOutput {
            ane: 0.0,
            anp: 0.0,
            ahtot: 0.0,
            ahmol: 0.0,
            anhm: 0.0,
            energ: 0.0,
            entt: 0.0,
            wm: 0.0,
            rhoter: 0.0,
            anerel: 0.0,
            anato_h: None,
            anato_he: None,
            wmm_update: 0.0,
        };
    }

    let t = params.t;
    let an = params.an;

    // 初始电子密度比例估计
    let mut anerel = if t < 4000.0 {
        1e-6
    } else if t < 5000.0 {
        1e-5
    } else if t < 5500.0 {
        1e-4
    } else if t < 6000.0 {
        1e-3
    } else if t < 7000.0 {
        0.01
    } else if t < 8000.0 {
        0.1
    } else if t < 9000.0 {
        0.4
    } else {
        0.5
    };

    // 如果包含分子且温度低于分子温度上限，调用 MOLEQ
    if params.config.ifmol > 0 && t < params.config.tmolim {
        let aein = an * anerel;
        // 调用 moleq_pure 计算分子平衡
        let default_mol_data: MoleculeEqData = Default::default();
        let moleq_params = MoleqParams {
            id: params.id,
            tt: t,
            an,
            aein,
            abundances: &[1.0], // 简化
            ionization_energies: &[13.6],
            ionization_energies2: &[0.0],
            atomic_masses: &[1.0],
            nelemx: &[],
            nmetal: 0,
            molecule_data: params.molecule_data.as_ref()
                .map(|d| d as &MoleculeEqData)
                .unwrap_or(&default_mol_data),
            heh: 0.0,
            ipri: 0,
            ifmol: params.config.ifmol,
            moltab: 0,
        };
        let moleq_output = moleq_pure(&moleq_params);
        let ane = moleq_output.ane;
        anerel = ane / an;

        return EldensOutput {
            ane,
            anp: 0.0,
            ahtot: an / params.ytot,
            ahmol: moleq_output.anhm,
            anhm: moleq_output.anhm,
            energ: moleq_output.energ,
            entt: moleq_output.entt,
            wm: moleq_output.wm,
            rhoter: params.wmy * (an / params.ytot) * HMASS,
            anerel,
            anato_h: None,
            anato_he: None,
            wmm_update: 0.0,
        };
    }

    let tk = BOLK * t;
    let thet = 5040.4 / t;

    // 初始化系数
    let mut qmi = 0.0;
    let mut qp = 0.0;
    let mut q2 = 0.0;
    let mut uh2 = 1.0;
    let mut duh2 = 0.0;

    // 低温时的分子系数
    if t <= 9000.0 {
        // H- 复合系数
        qmi = 1.0353e-16 / t / t.sqrt() * (8762.9 / t).exp();

        // H2 解离系数
        qp = tk * ((-11.206998 + thet * (2.7942767 + thet * (0.079196803 - 0.024790744 * thet))) * 2.30258509299405).exp();

        // H 配分函数
        let MpartfResult { u: uh, .. } = mpartf(1, 1, 0, t);
        let uh = uh.max(TWO);

        // H2 配分函数
        let MpartfResult { u: uh2_val, dulog: duh2_val } = mpartf(0, 0, 2, t);
        uh2 = uh2_val;
        duh2 = duh2_val;

        // H2+ 电离系数
        q2 = 1.47e-20 / (t * t.sqrt()) * uh2 / uh / uh * (51951.8 / t).exp();
    }

    let tkln15 = (BOLK * t).ln() * 1.5;

    // H 电离系数
    let qh0 = ((15.38287 + 1.5 * t.log10() - 13.595 * thet) * 2.30258509299405).exp() * TWO;
    let qh = qh0 / params.config.pfhyd;

    // 初始值
    let mut ane = an * anerel;
    let mut ah = 0.0;
    let mut anh = 0.0;
    let mut q = 0.0;
    let mut dqn = 0.0;

    // Newton-Raphson 迭代
    for it in 1..=10 {
        // 调用 STATE 计算总电荷
        // 注意：STATE 需要当前的电子密度 ane，而不是初始估计值
        if let Some(ref state_params) = params.state_params {
            // 创建新的 StateParams，使用当前迭代的 ane 值
            let updated_params = StateParams {
                mode: state_params.mode,
                id: state_params.id,
                t: state_params.t,
                ane,  // 使用当前迭代的电子密度
                natoms: state_params.natoms,
                hpop: state_params.hpop,
                dens: state_params.dens,
                wmm: state_params.wmm,
                ytot: state_params.ytot,
                abndd: state_params.abndd,
                ioniz: state_params.ioniz,
                irefa: state_params.irefa,
                lgr: state_params.lgr,
                ifoppf: state_params.ifoppf,
                lrm: state_params.lrm,
            };
            let state_output = state_pure(&updated_params);
            q = state_output.q;
            dqn = state_output.dqn;
        }

        // 判断氢是否为参考原子
        if params.config.iatref == params.config.iath || params.config.ioptab >= -1 {
            // 氢是参考原子
            let g2 = qh / ane;
            let g3 = qmi * ane;
            let a_val = UN + g2 + g3;
            let d_val = g2 - g3;

            // 第一次迭代的初始估计
            if it == 1 {
                if t > 9000.0 {
                    let f1 = UN / a_val;
                    let fe = d_val / a_val + q;
                    ah = ane / fe;
                    anh = ah * f1;
                    if params.id == 1 {
                        eprintln!("DEBUG ELDENS init T>9000: f1={:.6e}, fe={:.6e}, q={:.6e}, ah={:.6e}, anh={:.6e}",
                            f1, fe, q, ah, anh);
                    }
                } else if t > 4000.0 {
                    let e_val = g2 * qp / q2;
                    let b_val = TWO * (UN + e_val);
                    let gg = ane * q2;
                    let c1 = b_val * (gg * b_val + a_val * d_val) - e_val * a_val * a_val;
                    let c2 = a_val * (TWO * e_val + b_val * q) - d_val * b_val;
                    let c3 = -e_val - b_val * q;
                    let disc = (c2 * c2 - 4.0 * c1 * c3).sqrt();
                    let f1 = (disc - c2) * HALF / c1;
                    let fe = f1 * d_val + e_val * (UN - a_val * f1) / b_val + q;
                    ah = ane / fe;
                    anh = ah * f1;
                    if params.id == 1 {
                        eprintln!("DEBUG ELDENS init 4000<T<9000: ah={:.6e}, anh={:.6e}", ah, anh);
                    }
                } else {
                    // 低温情况
                    let c1 = q2 * (TWO * params.ytot - UN);
                    let c2 = params.ytot;
                    let c3 = -an;
                    let disc = (c2 * c2 - 4.0 * c1 * c3).sqrt();
                    anh = (disc - c2) * HALF / c1;
                    ah = anh * (UN + TWO * anh * q2);

                    let c1 = UN + qmi * anh;
                    let c2 = -q * ah;
                    let c3 = -qh * anh;
                    let disc = (c2 * c2 - 4.0 * c1 * c3).sqrt();
                    ane = (disc - c2) * HALF / c1;
                }
            }

            let ae = anh / ane;
            let gg = ae * qp;
            let e_val = anh * q2;
            let b_val = anh * qmi;

            if params.id == 1 && it <= 2 {
                eprintln!("DEBUG ELDENS coeffs iter {}: ae={:.6e}, gg={:.6e}, e_val={:.6e}, b_val={:.6e}",
                    it, ae, gg, e_val, b_val);
                eprintln!("DEBUG ELDENS params iter {}: d_val={:.6e}, g2={:.6e}, a_val={:.6e}",
                    it, d_val, g2, a_val);
            }

            // 构建线性方程组
            let mut r = [[0.0; 3]; 3];
            let mut s = [0.0; 3];

            if params.config.ifmol == 0 || t >= params.config.tmolim {
                r[0][0] = params.ytot;
                r[0][1] = 0.0;
                r[0][2] = UN;
                s[0] = an - ane - params.ytot * ah;
                if params.id == 1 && it <= 2 {
                    eprintln!("DEBUG ELDENS rhs0: an={:.6e}, ane={:.6e}, ytot={:.6e}, ah={:.6e}, s0={:.6e}",
                        an, ane, params.ytot, ah, s[0]);
                }
            } else {
                r[0][0] = params.ytot - UN;
                r[0][1] = a_val + e_val + gg;
                r[0][2] = UN;
                s[0] = an - ane - anh * (a_val + e_val + gg) - (params.ytot - UN) * ah;
            }

            r[1][0] = -q;
            r[1][1] = -d_val - TWO * gg;
            r[1][2] = UN + b_val + ae * (g2 + gg) - dqn * ah;

            r[2][0] = -UN;
            r[2][1] = a_val + 4.0 * (e_val + gg);
            r[2][2] = b_val - ae * (g2 + TWO * gg);

            s[1] = anh * (d_val + gg) + q * ah - ane;
            s[2] = ah - anh * (a_val + TWO * (e_val + gg));

            if params.id == 1 && it <= 2 {
                eprintln!("DEBUG ELDENS rhs12: anh={:.6e}, d_val={:.6e}, gg={:.6e}, q={:.6e}, ah={:.6e}, ane={:.6e}",
                    anh, d_val, gg, q, ah, ane);
                eprintln!("DEBUG ELDENS rhs12: s1={:.6e}, s2={:.6e}", s[1], s[2]);
            }

            // Debug: print matrix
            if params.id == 1 && it <= 2 {
                eprintln!("DEBUG ELDENS matrix iter {}: R=[{:.15e},{:.15e},{:.15e}; {:.15e},{:.15e},{:.15e}; {:.15e},{:.15e},{:.15e}]",
                    it, r[0][0], r[0][1], r[0][2], r[1][0], r[1][1], r[1][2], r[2][0], r[2][1], r[2][2]);
                eprintln!("DEBUG ELDENS rhs iter {}: S=[{:.15e},{:.15e},{:.15e}]",
                    it, s[0], s[1], s[2]);
            }

            // 求解线性方程组
            // 注意：LINEQS 需要 Fortran 列优先存储
            let mut r_work = vec![0.0; 9];
            for j in 0..3 {
                for i in 0..3 {
                    r_work[i + j * 3] = r[i][j];  // 列优先存储
                }
            }
            let mut s_work = s.to_vec();
            let mut p = [0.0; 3];

            lineqs(&mut r_work, &mut s_work, &mut p, 3);

            // 验证解是否正确
            if params.id == 1 && it <= 2 {
                let check0 = r[0][0]*p[0] + r[0][1]*p[1] + r[0][2]*p[2];
                let check1 = r[1][0]*p[0] + r[1][1]*p[1] + r[1][2]*p[2];
                let check2 = r[2][0]*p[0] + r[2][1]*p[1] + r[2][2]*p[2];
                eprintln!("DEBUG ELDENS verify: R*p=[{:.6e}, {:.6e}, {:.6e}], S=[{:.6e}, {:.6e}, {:.6e}]",
                    check0, check1, check2, s[0], s[1], s[2]);
            }

            if params.id == 1 {
                eprintln!("DEBUG ELDENS lineqs: p=[{:.6e}, {:.6e}, {:.6e}]", p[0], p[1], p[2]);
            }

            // 更新值
            ah = ah + p[0];
            anh = anh + p[1];
            let delne = p[2];
            ane = ane + delne;

            // 收敛检查（与 Fortran 一致）
            // 注意：Fortran 只检查 ANE <= 0，不检查 ANH
            if ane <= 0.0 {
                ane = 1e-6 * an;
            }
            if (delne / ane).abs() <= 1e-3 {
                break;
            }
        } else {
            // 氢不是参考原子
            if it == 1 {
                ane = an * HALF;
                ah = ane / params.ytot;
            }

            let mut r = [[0.0; 2]; 2];
            let mut s = [0.0; 2];
            let mut p = [0.0; 2];

            r[0][0] = params.ytot;
            r[0][1] = UN;
            r[1][0] = -q - params.qref;
            r[1][1] = UN - (dqn + params.dqnr) * ah;

            s[0] = an - ane - params.ytot * ah;
            s[1] = (q + params.qref) * ah - ane;

            // 求解 2x2 系统（列优先存储，与 LINEQS 一致）
            let mut r_work = vec![0.0; 4];
            for j in 0..2 {
                for i in 0..2 {
                    r_work[i + j * 2] = r[i][j];  // 列优先存储
                }
            }
            let mut s_work = s.to_vec();

            lineqs(&mut r_work, &mut s_work, &mut p, 2);

            ah = ah + p[0];
            let delne = p[1];
            ane = ane + delne;

            // 收敛检查
            if (delne / ane).abs() <= 1e-3 {
                break;
            }
        }
    }

    if params.id == 1 {
        eprintln!("DEBUG ELDENS after loop: ane={:.6e}, an={:.6e}", ane, an);
    }

    // 确保电子密度为正
    if ane <= 0.0 {
        ane = 1e-6 * an;
    }

    // 计算最终值
    anerel = ane / an;
    let ahtot = ah;
    let ahmol = anh * anh * q2;

    // 计算质子密度 (从 Saha 方程)
    // Saha: n_p * n_e / n_H = qh
    // 所以: n_p = n_H * qh / n_e
    let anp_raw = anh / ane * qh;

    // 确保物理一致性:
    // 1. 质子密度不能超过总氢密度
    // 2. 对于氢主导的气体, ane ≈ anp (电荷守恒)
    // 3. 如果 anp > ahtot, 说明计算有问题, 使用 ane 作为估计
    let anp = if anp_raw > ahtot || anp_raw < 0.0 {
        // Newton-Raphson 可能未收敛, 使用电荷守恒估计
        // 对于氢主导气体: ane ≈ anp (电子来自氢电离)
        ane.min(ahtot).max(0.0)
    } else {
        anp_raw
    };

    let anhm = anh * ane * qmi;
    let rhoter = params.wmy * ah * HMASS;

    // 调用 ENTENE 计算内能和熵
    // f2r_depends: ENTENE, MOLEQ
    let rr_dummy = [[0.0f64; 2]; 30];
    let enev_dummy = [[0.0f64; 2]; 30];
    let amas_dummy = [0.0f64; 30];
    let entene_params = EnteneParams {
        t,
        ah,
        anh,
        anpr: anp,
        ane,
        rr: &rr_dummy,
        enev: &enev_dummy,
        amas: &amas_dummy,
        natoms: 1,
        bolk: BOLK,
        un: UN,
    };
    let entene_output = entene(&entene_params);
    let mut energ = entene_output.energ;
    let mut entt = entene_output.entrop;

    // H2 的能量和熵修正
    if t < 9000.0 && ahmol > 0.0 && uh2 > 0.0 {
        energ = energ + (duh2 - 51951.8 / t) * tk * ahmol;
        entt = entt + ahmol * (tkln15 - ahmol.ln() + uh2.ln() + 1.0487 + 103.973) * BOLK;
    }

    let wm = rhoter / an / HMASS;

    // ANATO 计算 (Fortran lines 245-256)
    // if(ifmol.le.0.or.t.ge.tmolim) then
    let (anato_h, anato_he) = if params.config.ifmol <= 0 || t >= params.config.tmolim {
        if let Some(ref ad) = params.anato_data {
            // anato(1,id)
            let ah = if ad.n0hn > 0 {
                ad.popul_h
            } else {
                ad.lte_estimate
            };
            // anato(2,id)
            let ahe = if ad.iathe > 0 {
                ad.popul_he
            } else {
                ad.lte_estimate * ad.abndd_he
            };
            (Some(ah), Some(ahe))
        } else {
            (None, None)
        }
    } else {
        (None, None)
    };

    // WMM 更新值: wmm(id) = dens(id)/(an-ane) = rhoter/(an-ane)
    let wmm_update = if an > ane { rhoter / (an - ane) } else { 0.0 };

    if params.id == 1 {
        eprintln!("DEBUG ELDENS return: id={}, ane={:.6e}, anp={:.6e}, ahtot={:.6e}, anerel={:.6e}",
            params.id, ane, anp, ahtot, anerel);
    }

    EldensOutput {
        ane,
        anp,
        ahtot,
        ahmol,
        anhm,
        energ,
        entt,
        wm,
        rhoter,
        anerel,
        anato_h,
        anato_he,
        wmm_update,
    }
}

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

    #[test]
    fn test_eldens_basic() {
        let config = EldensConfig::default();
        let params = EldensParams {
            id: 1,
            t: 10000.0,
            an: 1e15,
            ytot: 1.0,
            qref: 0.0,
            dqnr: 0.0,
            wmy: 1.0,
            config,
            state_params: None,
            molecule_data: None,
            anato_data: None,
        };

        let output = eldens_pure(&params, 0);

        // 验证输出
        assert!(output.ane > 0.0);
        assert!(output.anerel.is_finite());
    }

    #[test]
    fn test_eldens_low_temp() {
        let config = EldensConfig {
            ifmol: 0,
            tmolim: 5000.0,
            ..Default::default()
        };
        let params = EldensParams {
            id: 1,
            t: 6000.0,
            an: 1e15,
            ytot: 1.0,
            qref: 0.0,
            dqnr: 0.0,
            wmy: 1.0,
            config,
            state_params: None,
            molecule_data: None,
            anato_data: None,
        };

        let output = eldens_pure(&params, 0);

        // 低温时电子密度应该较低
        assert!(output.anerel < 0.5);
    }
}