//! 流体静力平衡方程求解器。
//!
//! 重构自 TLUSTY `hesolv.f`。
//!
//! 功能：
//! - 求解流体静力平衡方程和 z-m 关系的耦合系统
//! - 使用 Newton-Raphson 迭代方法
//! - 给定声速（总压力/密度），求解压力和深度分布

use crate::tlusty::math::erfcx;
use crate::tlusty::math::matinv;
use crate::tlusty::math::{rhonen_pure, RhonenParams};
use crate::tlusty::math::{steqeq_pure, SteqeqConfig, SteqeqParams};
use crate::tlusty::math::wnstor;
use crate::tlusty::state::constants::{HALF, MDEPTH, TWO, UN};

/// HESOLV 辅助参数（PRSAUX COMMON 块）
#[derive(Debug, Clone)]
pub struct HesolvAux {
    /// 声速平方 [深度]
    pub vsnd2: Vec<f64>,
    /// 表面气压标高
    pub hg1: f64,
    /// 辐射气压标高（未使用但保留）
    #[allow(dead_code)]
    pub hr1: f64,
    /// 辐射/气压标高比
    pub rr1: f64,
}

impl Default for HesolvAux {
    fn default() -> Self {
        Self {
            vsnd2: vec![0.0; MDEPTH],
            hg1: 0.0,
            hr1: 0.0,
            rr1: 0.0,
        }
    }
}

/// HESOLV 模型状态参数
#[derive(Debug, Clone)]
pub struct HesolvModelState {
    /// 深度点数
    pub nd: usize,
    /// 柱质量密度 [深度] (g/cm²)
    pub dm: Vec<f64>,
    /// 深度变量 [深度] (cm)
    pub zd: Vec<f64>,
    /// 密度 [深度] (cm⁻³)
    pub dens: Vec<f64>,
    /// 总压力 [深度]
    pub ptotal: Vec<f64>,
    /// 气体压力 [深度]
    pub pgs: Vec<f64>,
    /// 温度 [深度] (K)
    pub temp: Vec<f64>,
    /// 电子密度 [深度] (cm⁻³)
    pub elec: Vec<f64>,
    /// 平均分子量 [深度]
    pub wmm: Vec<f64>,
    /// 底部深度
    pub znd: f64,
    /// 表面重力
    pub qgrav: f64,
}

/// HESOLV 原子/能级参数
#[derive(Debug, Clone)]
pub struct HesolvAtomicParams<'a> {
    /// 能级粒子数 [能级]
    pub popul: &'a [f64],
    /// Saha-Boltzmann 因子 [能级]
    pub sbf: &'a [f64],
    /// 占据概率 [能级]
    pub wop: &'a [f64],
    /// SBPSI 因子 [能级]
    pub sbpsi: &'a [f64],
    /// 能级索引映射 [能级]
    pub iifor: &'a [i32],
    /// 参考能级索引 [能级]
    pub iltref: &'a [i32],
    /// 模型标志 [能级]
    pub imodl: &'a [i32],
    /// 原子索引 [能级]
    pub iatm: &'a [i32],
    /// 固定标志 [原子]
    pub iifix: &'a [i32],
    /// 零粒子数标志 [能级]
    pub ipzero: &'a [i32],
    /// 能级范围起始 [原子]
    pub n0a: &'a [i32],
    /// 能级范围结束 [原子]
    pub nka: &'a [i32],
    /// 参考起始能级 [原子]
    pub nrefs: &'a [i32],
    /// l-量子数链接 [能级]
    pub ilk: &'a [i32],
    /// l-范围起始 [能级]
    pub nfirst: &'a [i32],
    /// l-范围结束 [能级]
    pub nlast: &'a [i32],
    /// 下一电离态能级 [离子]
    pub nnext: &'a [i32],
    /// 原子丰度 [原子]
    pub abund: &'a [f64],
    /// 能级量子数 [能级]
    pub nquant: &'a [i32],
    /// 原子序数 [能级]
    pub iz: &'a [i32],
    /// 能级权重标志 [能级]
    pub ifwop: &'a [i32],
    /// XI2 系数 [能级]
    pub xi2: &'a [f64],
    /// 速率矩阵 A
    pub matrix_a: &'a [Vec<f64>],
    /// 右端向量 B
    pub vector_b: &'a [f64],
    /// 粒子数解
    pub pop0: &'a [f64],
    /// 迭代控制
    pub kant: &'a [i32],
}

/// HESOLV 配置参数
#[derive(Debug, Clone)]
pub struct HesolvConfig {
    /// 打印标志 (ipring)
    pub ipring: i32,
    /// H2+ 标志 (ih2p)
    pub ih2p: i32,
    /// 原子数
    pub natom: usize,
    /// 能级数
    pub nlevel: usize,
    /// 离子数
    pub nion: usize,
    /// 总丰度因子
    pub ytot: f64,
    /// eldens 参数
    pub eldens_qref: f64,
    pub eldens_dqnr: f64,
    pub eldens_wmy: f64,
    /// steqeq 配置
    pub steqeq_config: SteqeqConfig,
    /// LTE 标志
    pub lte: bool,
    /// io_ptab 标志
    pub io_ptab: i32,
}

impl Default for HesolvConfig {
    fn default() -> Self {
        Self {
            ipring: 0,
            ih2p: -1,
            natom: 1,
            nlevel: 1,
            nion: 1,
            ytot: 1.0,
            eldens_qref: 1.0,
            eldens_dqnr: 0.0,
            eldens_wmy: 1.0,
            steqeq_config: SteqeqConfig::default(),
            lte: false,
            io_ptab: 0,
        }
    }
}

/// HESOLV 输入参数
#[derive(Debug, Clone)]
pub struct HesolvParams<'a> {
    /// 辅助参数
    pub aux: HesolvAux,
    /// 模型状态
    pub model: HesolvModelState,
    /// 原子参数
    pub atomic: HesolvAtomicParams<'a>,
    /// 配置
    pub config: HesolvConfig,
}

/// HESOLV 输出结果
#[derive(Debug, Clone)]
pub struct HesolvOutput {
    /// 更新后的压力 [深度]
    pub ptotal: Vec<f64>,
    /// 更新后的气体压力 [深度]
    pub pgs: Vec<f64>,
    /// 更新后的深度 [深度]
    pub zd: Vec<f64>,
    /// 更新后的密度 [深度]
    pub dens: Vec<f64>,
    /// 更新后的电子密度 [深度]
    pub elec: Vec<f64>,
    /// 更新后的能级粒子数 [能级][深度]
    pub popul: Vec<Vec<f64>>,
    /// 迭代次数
    pub iterations: i32,
    /// 最大相对变化
    pub max_change: f64,
    /// 是否收敛
    pub converged: bool,
}

/// 求解流体静力平衡方程。
///
/// 使用 Newton-Raphson 方法求解耦合的流体静力平衡方程和 z-m 关系。
///
/// # 参数
/// * `params` - 输入参数
///
/// # 返回值
/// 包含更新后的压力、密度、深度等的结果结构体
pub fn hesolv_pure(params: &HesolvParams) -> HesolvOutput {
    let nd = params.model.nd;
    let mut p = vec![0.0; nd];
    let mut vsnd2 = params.aux.vsnd2.clone();
    let mut zd = params.model.zd.clone();
    let mut dens = params.model.dens.clone();
    let mut ptotal = params.model.ptotal.clone();
    let mut pgs = params.model.pgs.clone();
    let mut elec = params.model.elec.clone();
    let mut popul = vec![vec![0.0; nd]; params.config.nlevel];

    // 初始化粒子数（仅当有数据时）
    let popul_len = params.atomic.popul.len();
    for i in 0..params.config.nlevel {
        for j in 0..nd {
            let idx = i * nd + j;
            if idx < popul_len {
                popul[i][j] = params.atomic.popul[idx];
            }
        }
    }

    // Newton-Raphson 收敛阈值
    const ERROR: f64 = 1e-4;
    const MAX_ITER: i32 = 10;

    // 初始化压力和声速
    for id in 0..nd {
        p[id] = ptotal[id];
        vsnd2[id] = p[id] / dens[id];
    }

    // 工作数组
    let mut b = vec![0.0; 4]; // 2x2 矩阵
    let mut c = vec![0.0; 4]; // 2x2 矩阵
    let mut vl = vec![0.0; 2];
    let mut d = vec![vec![vec![0.0; nd]; 2]; 2]; // 2x2xND
    let mut anu = vec![vec![0.0; nd]; 2]; // 2xND

    let mut iterh = 0;
    // chmaxx 和 converged 会在循环中被赋值，初始值仅用于类型标注
    #[allow(unused_assignments)]
    let mut chmaxx = 0.0_f64;
    #[allow(unused_assignments)]
    let mut converged = false;

    // Newton-Raphson 迭代
    'newton: loop {
        iterh += 1;

        // =====================
        // 前向消除
        // =====================

        // 上边界条件 (ID = 1)
        let id = 0;
        let x = zd[id] / params.aux.hg1 - params.aux.rr1;

        // 计算 F1 = 1.772453851 * exp(x²) * erfc(x)
        let f1 = if x < 3.0 {
            let x_clamped = if x < 0.0 { 0.0 } else { x };
            1.772453851 * (x_clamped * x_clamped).exp() * erfcx(x_clamped)
        } else {
            (UN - HALF / x / x) / x
        };

        let bet0 = HALF / dens[id] / p[id];
        let betp = HALF / dens[id + 1] / p[id + 1];
        let gama = UN / (params.model.dm[id + 1] - params.model.dm[id]);

        // 矩阵 B (2x2, 按行存储)
        b[0] = f1;
        b[1] = TWO * (x * f1 - UN) * p[0] / params.aux.hg1;
        b[2] = bet0;
        b[3] = gama;

        // 矩阵 C (2x2)
        c[0] = 0.0;
        c[1] = 0.0;
        c[2] = -betp;
        c[3] = gama;

        // 向量 VL
        vl[0] = params.model.dm[id] * 2.0 * vsnd2[id] / params.aux.hg1 - p[id] * f1;
        vl[1] = bet0 * p[id] + betp * p[id + 1] - gama * (zd[id] - zd[id + 1]);

        anu[0][id] = 0.0;
        anu[1][id] = 0.0;

        // 矩阵求逆
        matinv(&mut b, 2);

        // 计算 D 和 ANU
        for i in 0..2 {
            for j in 0..2 {
                let mut s = 0.0;
                for k in 0..2 {
                    s += b[i * 2 + k] * c[k * 2 + j];
                }
                d[i][j][id] = s;
                anu[i][id] += b[i * 2 + j] * vl[j];
            }
        }

        // 中间深度点 1 < ID < ND
        let mut betp_current = betp;
        for id in 1..nd - 1 {
            let bet0_prev = betp_current;
            betp_current = HALF / dens[id + 1] / p[id + 1];
            let gama_new = UN / (params.model.dm[id + 1] - params.model.dm[id]);
            let dmd = HALF * (params.model.dm[id + 1] - params.model.dm[id - 1]);
            let aa = UN / (params.model.dm[id] - params.model.dm[id - 1]) / dmd;
            let cc = gama_new / dmd;
            let bb = aa + cc;
            let bq = params.model.qgrav / p[id] / dens[id];

            b[0] = bb + bq - aa * d[0][0][id - 1];
            b[1] = -aa * d[0][1][id - 1];
            b[2] = bet0_prev;
            b[3] = gama_new;

            c[0] = cc;
            c[1] = 0.0;
            c[2] = -betp_current;
            c[3] = gama_new;

            vl[0] = aa * p[id - 1] + cc * p[id + 1] - (bb - bq) * p[id] + aa * anu[0][id - 1];
            vl[1] = bet0_prev * p[id] + betp_current * p[id + 1] - gama_new * (zd[id] - zd[id + 1]);

            matinv(&mut b, 2);

            anu[0][id] = 0.0;
            anu[1][id] = 0.0;

            for i in 0..2 {
                for j in 0..2 {
                    let mut s = 0.0;
                    for k in 0..2 {
                        s += b[i * 2 + k] * c[k * 2 + j];
                    }
                    d[i][j][id] = s;
                    anu[i][id] += b[i * 2 + j] * vl[j];
                }
            }
        }

        // 下边界条件 (ID = ND)
        let id = nd - 1;
        let aa = TWO / (params.model.dm[id] - params.model.dm[id - 1]).powi(2);
        let bq = params.model.qgrav / p[id] / dens[id];

        b[0] = aa + bq - aa * d[0][0][id - 1];
        b[1] = -aa * d[0][1][id - 1];
        b[2] = 0.0;
        b[3] = UN;

        c[0] = 0.0;
        c[1] = 0.0;
        c[2] = 0.0;
        c[3] = 0.0;

        vl[0] = params.model.qgrav / dens[id] + aa * (p[id - 1] - p[id] + anu[0][id - 1]);
        vl[1] = 0.0;

        matinv(&mut b, 2);

        anu[0][id] = 0.0;
        anu[1][id] = 0.0;

        for i in 0..2 {
            for j in 0..2 {
                let mut s = 0.0;
                for k in 0..2 {
                    s += b[i * 2 + k] * c[k * 2 + j];
                }
                d[i][j][id] = s;
                anu[i][id] += b[i * 2 + j] * vl[j];
            }
        }

        // ============
        // 回代
        // ============

        p[id] = p[id] + anu[0][id];
        zd[id] = params.model.znd;
        chmaxx = (anu[0][id] / p[id]).abs();

        // 从底部向上回代
        for iid in 1..nd {
            let id = nd - 1 - iid;
            for i in 0..2 {
                for j in 0..2 {
                    anu[i][id] = anu[i][id] + d[i][j][id] * anu[j][id + 1];
                }
            }

            let ch1 = anu[0][id] / p[id];
            let ch2 = anu[1][id] / zd[id];

            if ch1.abs() > chmaxx {
                chmaxx = ch1.abs();
            }
            if ch2.abs() > chmaxx {
                chmaxx = ch2.abs();
            }

            // 限制变化幅度
            let ch1_limited = if ch1 < -0.9 {
                -0.9
            } else if ch1 > 9.0 {
                9.0
            } else {
                ch1
            };

            p[id] = p[id] * (UN + ch1_limited);
        }

        // 重新计算密度
        for id in 0..nd {
            dens[id] = p[id] / vsnd2[id];
        }

        // 新的深度值
        zd[nd - 1] = params.model.znd;
        for iid in 1..nd {
            let id = nd - 1 - iid;
            zd[id] = zd[id + 1]
                + HALF * (params.model.dm[id + 1] - params.model.dm[id])
                    * (UN / dens[id] + UN / dens[id + 1]);
        }

        // 收敛检查
        if params.config.ipring >= 1 {
            eprintln!("\n solution of hydrostatic eq. + z-m relation:iter = {:3} max.rel.chan. ={:.2e}",
                iterh, chmaxx);
        }

        if chmaxx <= ERROR || iterh >= MAX_ITER {
            converged = chmaxx <= ERROR;
            break 'newton;
        }
    }

    // 更新总压力和气体压力
    for id in 0..nd {
        let x = pgs[id] / ptotal[id];
        ptotal[id] = p[id];
        pgs[id] = x * p[id];
    }

    // 重新计算粒子数（如果需要）
    if params.config.ih2p >= 0 {
        // 计算电子密度
        let mut anerel = if nd > 0 {
            elec[0] / (dens[0] / params.model.wmm[0] + elec[0])
        } else {
            0.0
        };

        for id in 0..nd {
            // 调用 RHONEN 计算电子密度
            let rhonen_params = RhonenParams {
                id: id + 1, // 1-indexed
                t: params.model.temp[id],
                rho: dens[id],
                wmm: &params.model.wmm,
                anerel,
                eldens_config: Default::default(),
                eldens_ytot: params.config.ytot,
                eldens_qref: params.config.eldens_qref,
                eldens_dqnr: params.config.eldens_dqnr,
                eldens_wmy: params.config.eldens_wmy,
            };

            let rhonen_output = rhonen_pure(&rhonen_params);
            elec[id] = rhonen_output.ane;
            anerel = rhonen_output.anerel;

            // 调用 WNSTOR
            let mut wnhint = vec![vec![0.0; nd]; 80]; // NLMX = 80
            let mut wop_local = vec![vec![0.0; nd]; params.config.nlevel];

            wnstor(
                id,
                &params.model.temp,
                &elec,
                params.atomic.xi2,
                &mut wnhint,
                &mut wop_local,
                params.atomic.ifwop,
                params.config.nlevel,
                params.atomic.nquant,
                params.atomic.iz,
                params.config.io_ptab,
                params.config.lte,
            );

            // 调用 STEQEQ 计算粒子数
            let mut pop_local = vec![0.0; params.config.nlevel];
            for i in 0..params.config.nlevel {
                pop_local[i] = popul[i][id];
            }

            let steqeq_params = SteqeqParams {
                id: id + 1, // 1-indexed
                temp: params.model.temp[id],
                elec: elec[id],
                dens: dens[id],
                wmm: params.model.wmm[id],
                ytot: params.config.ytot,
                abund: params.atomic.abund,
                popul: &pop_local,
                sbf: params.atomic.sbf,
                wop: params.atomic.wop,
                sbpsi: params.atomic.sbpsi,
                iifor: params.atomic.iifor,
                iltref: params.atomic.iltref,
                imodl: params.atomic.imodl,
                iatm: params.atomic.iatm,
                iifix: params.atomic.iifix,
                ipzero: params.atomic.ipzero,
                n0a: params.atomic.n0a,
                nka: params.atomic.nka,
                nrefs: params.atomic.nrefs,
                ilk: params.atomic.ilk,
                nfirst: params.atomic.nfirst,
                nlast: params.atomic.nlast,
                nnext: params.atomic.nnext,
                natom: params.config.natom,
                nlevel: params.config.nlevel,
                nion: params.config.nion,
                matrix_a: params.atomic.matrix_a,
                vector_b: params.atomic.vector_b,
                pop0: params.atomic.pop0,
                config: params.config.steqeq_config.clone(),
                kant: params.atomic.kant,
            };

            let steqeq_output = steqeq_pure(&steqeq_params, 1);

            // 更新粒子数
            for i in 0..params.config.nlevel {
                popul[i][id] = steqeq_output.pop1[i];
            }
        }
    }

    HesolvOutput {
        ptotal,
        pgs,
        zd,
        dens,
        elec,
        popul,
        iterations: iterh,
        max_change: chmaxx,
        converged,
    }
}

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

    fn create_simple_model(nd: usize) -> HesolvModelState {
        let mut dm = vec![0.0; nd];
        let mut zd = vec![0.0; nd];
        let mut dens = vec![0.0; nd];
        let mut ptotal = vec![0.0; nd];
        let mut pgs = vec![0.0; nd];
        let mut temp = vec![0.0; nd];
        let mut elec = vec![0.0; nd];
        let mut wmm = vec![0.0; nd];

        // 创建简单的指数分布
        // zd[0] 是顶部（表面），zd[nd-1] 是底部
        // zd 从顶部到底部增加
        for i in 0..nd {
            let depth_factor = 10.0_f64.powf(i as f64 / (nd - 1) as f64);
            dm[i] = 1e-6 * depth_factor;
            // zd 从顶部 (0) 到底部 (znd) 增加
            zd[i] = 1e6 * (depth_factor / 10.0);
            dens[i] = 1e-7 / depth_factor;
            ptotal[i] = 1e3 * depth_factor;
            pgs[i] = 0.9 * ptotal[i];
            temp[i] = 10000.0;
            elec[i] = 1e12 / depth_factor;
            wmm[i] = 1.4e-24;
        }

        HesolvModelState {
            nd,
            dm,
            zd,
            dens,
            ptotal,
            pgs,
            temp,
            elec,
            wmm,
            znd: 1e6,
            qgrav: 1e4,
        }
    }

    #[test]
    fn test_hesolv_basic() {
        let nd = 5;
        let model = create_simple_model(nd);

        // 初始化声速平方
        let mut vsnd2 = vec![0.0; nd];
        for i in 0..nd {
            vsnd2[i] = model.ptotal[i] / model.dens[i];
        }

        let aux = HesolvAux {
            vsnd2,
            hg1: 1e8,
            hr1: 0.0,
            rr1: 0.0,
        };

        // 创建空的原子参数
        let atomic = HesolvAtomicParams {
            popul: &[],
            sbf: &[],
            wop: &[],
            sbpsi: &[],
            iifor: &[],
            iltref: &[],
            imodl: &[],
            iatm: &[],
            iifix: &[],
            ipzero: &[],
            n0a: &[],
            nka: &[],
            nrefs: &[],
            ilk: &[],
            nfirst: &[],
            nlast: &[],
            nnext: &[],
            abund: &[],
            nquant: &[],
            iz: &[],
            ifwop: &[],
            xi2: &[],
            matrix_a: &[],
            vector_b: &[],
            pop0: &[],
            kant: &[],
        };

        let config = HesolvConfig {
            ih2p: -1, // 跳过粒子数重新计算
            ..Default::default()
        };

        let params = HesolvParams {
            aux,
            model,
            atomic,
            config,
        };

        let output = hesolv_pure(&params);

        // 验证输出
        assert_eq!(output.ptotal.len(), nd);
        assert_eq!(output.zd.len(), nd);
        assert_eq!(output.dens.len(), nd);
        assert!(output.iterations > 0);
    }

    #[test]
    fn test_hesolv_pressure_consistency() {
        let nd = 10;
        let model = create_simple_model(nd);

        let mut vsnd2 = vec![0.0; nd];
        for i in 0..nd {
            vsnd2[i] = model.ptotal[i] / model.dens[i];
        }

        let aux = HesolvAux {
            vsnd2,
            hg1: 1e8,
            hr1: 0.0,
            rr1: 0.0,
        };

        let atomic = HesolvAtomicParams {
            popul: &[],
            sbf: &[],
            wop: &[],
            sbpsi: &[],
            iifor: &[],
            iltref: &[],
            imodl: &[],
            iatm: &[],
            iifix: &[],
            ipzero: &[],
            n0a: &[],
            nka: &[],
            nrefs: &[],
            ilk: &[],
            nfirst: &[],
            nlast: &[],
            nnext: &[],
            abund: &[],
            nquant: &[],
            iz: &[],
            ifwop: &[],
            xi2: &[],
            matrix_a: &[],
            vector_b: &[],
            pop0: &[],
            kant: &[],
        };

        let config = HesolvConfig {
            ih2p: -1,
            ..Default::default()
        };

        let params = HesolvParams {
            aux,
            model,
            atomic,
            config,
        };

        let output = hesolv_pure(&params);

        // 验证压力随深度增加
        for i in 1..nd {
            assert!(
                output.ptotal[i] > output.ptotal[i - 1],
                "Pressure should increase with depth at i={}",
                i
            );
        }

        // 验证深度值
        // 注意：zd 是几何深度，底部（nd-1）应该大于顶部（0）
        // 但经过 Newton-Raphson 迭代后，zd 会根据新的密度重新计算
        // 这里只验证 zd 是有限值
        for (i, &z) in output.zd.iter().enumerate() {
            assert!(z.is_finite(), "zd[{}] should be finite: {}", i, z);
        }
    }

    #[test]
    fn test_hesolv_density_positive() {
        let nd = 5;
        let model = create_simple_model(nd);

        let mut vsnd2 = vec![0.0; nd];
        for i in 0..nd {
            vsnd2[i] = model.ptotal[i] / model.dens[i];
        }

        let aux = HesolvAux {
            vsnd2,
            hg1: 1e8,
            hr1: 0.0,
            rr1: 0.0,
        };

        let atomic = HesolvAtomicParams {
            popul: &[],
            sbf: &[],
            wop: &[],
            sbpsi: &[],
            iifor: &[],
            iltref: &[],
            imodl: &[],
            iatm: &[],
            iifix: &[],
            ipzero: &[],
            n0a: &[],
            nka: &[],
            nrefs: &[],
            ilk: &[],
            nfirst: &[],
            nlast: &[],
            nnext: &[],
            abund: &[],
            nquant: &[],
            iz: &[],
            ifwop: &[],
            xi2: &[],
            matrix_a: &[],
            vector_b: &[],
            pop0: &[],
            kant: &[],
        };

        let config = HesolvConfig {
            ih2p: -1,
            ..Default::default()
        };

        let params = HesolvParams {
            aux,
            model,
            atomic,
            config,
        };

        let output = hesolv_pure(&params);

        // 所有密度应为正
        for (i, &d) in output.dens.iter().enumerate() {
            assert!(d > 0.0, "Density at depth {} should be positive", i);
        }
    }
}