//! LTE-Grey 盘模型初始计算。
//!
//! 重构自 TLUSTY `ltegrd.f`。
//!
//! # 功能
//!
//! 计算初始的 LTE-Grey 盘模型，作为后续非 LTE 迭代的起点。
//! 这是盘模型（disk model）的计算，与大气模型（ltegr）不同。

use crate::math::zmrho;
use crate::state::constants::{HALF, MDEPTH, TWO, UN, SIG4P, SIGE, BOLK};

// ============================================================================
// 常量
// ============================================================================

/// 收敛容差
const ERRT: f64 = 1e-3;
/// 1/3
const THIRD: f64 = 1.0 / 3.0;
/// 4.0
const FOUR: f64 = 4.0;

// ============================================================================
// 配置结构体
// ============================================================================

/// LTEGRD 配置参数。
#[derive(Debug, Clone)]
pub struct LtegrdConfig {
    /// Grey 模型深度点数 (NDGREY)
    /// 0 = 使用 ND
    pub ndgrey: i32,
    /// 深度标尺模式 (IDGREY)
    pub idgrey: i32,
    /// 最大全局迭代次数 (ITGMX0)
    pub itgmx0: i32,
    /// 深度标尺重算次数 (NNEWD)
    pub nnewd: i32,
    /// 对流内部迭代次数 (NCONIT)
    pub nconit: i32,
    /// 诊断输出级别 (IPRING)
    pub ipring: i32,
    /// 混合长度参数 (HMIX0)
    pub hmix0: f64,
    /// 初始电离度估计 (DION0)
    pub dion0: f64,
    /// 初始 Rosseland 不透明度 (ABROS0)
    pub abros0: f64,
    /// 初始 Planck 平均不透明度 (ABPLA0)
    pub abpla0: f64,
    /// 第一深度点质量 (DM1)
    pub dm1: f64,
    /// DM 固定标志 (IDMFIX)
    pub idmfix: i32,
    /// 粘性参数 α (ALPHAV)
    pub alphav: f64,
    /// 粘性分数 (FRACTV)
    pub fractv: f64,
    /// ZETA0 参数
    pub zeta0: f64,
    /// ZETA1 参数
    pub zeta1: f64,
    /// 粘性质量比 (DMVISC)
    pub dmvisc: f64,
    /// 通量平均不透明度乘数 (ABFLXM)
    pub abflxm: f64,
}

impl Default for LtegrdConfig {
    fn default() -> Self {
        Self {
            ndgrey: 0,
            idgrey: 0,
            itgmx0: 5,
            nnewd: 0,
            nconit: 0,
            ipring: 0,
            hmix0: 0.0,
            dion0: 0.5,
            abros0: 0.4,
            abpla0: 0.4,
            dm1: 0.0,
            idmfix: 0,
            alphav: 1.0,
            fractv: 0.5,
            zeta0: 0.0,
            zeta1: 0.0,
            dmvisc: 0.1,
            abflxm: 0.4,
        }
    }
}

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

/// LTEGRD 输入参数。
pub struct LtegrdParams<'a> {
    /// 配置
    pub config: LtegrdConfig,
    /// 深度点数 (ND)
    pub nd: usize,
    /// 能级数 (NLEVEL)
    pub nlevel: usize,
    /// 有效温度 (TEFF)
    pub teff: f64,
    /// 表面重力加速度 (QGRAV)
    pub qgrav: f64,
    /// 平均分子量 [深度] (WMM)
    pub wmm: &'a [f64],
    /// 初始温度 [深度] (TEMP)
    pub temp: &'a mut [f64],
    /// 初始电子密度 [深度] (ELEC)
    pub elec: &'a mut [f64],
    /// 初始粒子密度 [深度] (DENS)
    pub dens: &'a mut [f64],
    /// 初始柱质量密度 [深度] (DM)
    pub dm: &'a mut [f64],
    /// 几何深度 [深度] (ZD)
    pub zd: &'a mut [f64],
    /// 总压力 [深度] (PTOTAL)
    pub ptotal: &'a mut [f64],
    /// 气体压力 [深度] (PGS)
    pub pgs: &'a mut [f64],
    /// Rosseland 光学深度 [深度] (TAUROS)
    pub tauros: &'a mut [f64],
    /// Rosseland 平均不透明度 [深度] (ABROSD)
    pub abrosd: &'a mut [f64],
    /// Planck 平均不透明度 [深度] (ABPLAD)
    pub abplad: &'a mut [f64],
    /// 湍流速度 [深度] (VTURB)
    pub vturb: &'a [f64],
    /// TAUTHE [深度]
    pub tauthe: &'a mut [f64],
    /// TAUFLX [深度]
    pub tauflx: &'a mut [f64],
    /// THETA [深度]
    pub theta: &'a mut [f64],
    /// VISCD [深度]
    pub viscd: &'a mut [f64],
    /// GAMJ [深度]
    pub gamj: &'a mut [f64],
    /// TOTJ [深度]
    pub totj: &'a mut [f64],
    /// TOTH [深度]
    pub toth: &'a mut [f64],
    /// TOTK [深度]
    pub totk: &'a mut [f64],
    /// RDOPAC [深度]
    pub rdopac: &'a mut [f64],
    /// FLOPAC [深度]
    pub flopac: &'a mut [f64],
}

/// LTEGRD 原子数据（简化版）。
pub struct LtegrdAtomicData<'a> {
    /// Saha-Boltzmann 因子 [能级][深度]
    pub sbf: &'a mut [Vec<f64>],
    /// 占据概率 [能级][深度]
    pub wop: &'a mut [Vec<f64>],
}

/// LTEGRD 输出。
#[derive(Debug, Clone)]
pub struct LtegrdOutput {
    /// 深度点数 (ND)
    pub nd: usize,
    /// 柱质量密度 [深度] (DM)
    pub dm: Vec<f64>,
    /// 温度 [深度] (TEMP)
    pub temp: Vec<f64>,
    /// 电子密度 [深度] (ELEC)
    pub elec: Vec<f64>,
    /// 总粒子密度 [深度] (DENS)
    pub dens: Vec<f64>,
    /// 几何深度 [深度] (ZD)
    pub zd: Vec<f64>,
    /// 总压力 [深度] (PTOTAL)
    pub ptotal: Vec<f64>,
    /// 气体压力 [深度] (PGS)
    pub pgs: Vec<f64>,
    /// Rosseland 光学深度 [深度] (TAUROS)
    pub tauros: Vec<f64>,
    /// Rosseland 平均不透明度 [深度]
    pub abrosd: Vec<f64>,
    /// Planck 平均不透明度 [深度]
    pub abplad: Vec<f64>,
    /// Eddington 因子 GAMH
    pub gamh: f64,
    /// 迭代计数 ITGREY
    pub itgrey: i32,
    /// 盘总质量 DMTOT
    pub dmtot: f64,
    /// 盘耗散 EDISC
    pub edisc: f64,
}

// ============================================================================
// 工作数组
// ============================================================================

/// LTEGRD 内部工作数组。
struct LtegrdWork {
    /// 备份 TEMP0
    temp0: Vec<f64>,
    /// 备份 ELEC0
    elec0: Vec<f64>,
    /// 备份 DENS0
    dens0: Vec<f64>,
    /// 备份 ZD0
    zd0: Vec<f64>,
    /// 备份 DM0
    dm0: Vec<f64>,
}

impl LtegrdWork {
    fn new() -> Self {
        Self {
            temp0: vec![0.0; MDEPTH],
            elec0: vec![0.0; MDEPTH],
            dens0: vec![0.0; MDEPTH],
            zd0: vec![0.0; MDEPTH],
            dm0: vec![0.0; MDEPTH],
        }
    }
}

// ============================================================================
// 核心计算函数
// ============================================================================

/// 执行 LTE-Grey 盘模型计算（纯计算，无 I/O）。
///
/// # 参数
/// - `params`: 输入参数
///
/// # 返回值
/// 计算结果
pub fn ltegrd_pure(params: &mut LtegrdParams) -> LtegrdOutput {
    let config = &params.config;
    let mut work = LtegrdWork::new();

    // 1. 确定深度点数
    let mut ndepth = if config.ndgrey == 0 {
        params.nd
    } else {
        config.ndgrey as usize
    };

    if ndepth > MDEPTH {
        panic!("NDEPTH too large in LTEGRD: {} > {}", ndepth, MDEPTH);
    }

    let idepth = config.idgrey;
    let mut itgmax = config.itgmx0;
    let mut nconit = config.nconit;

    if config.hmix0 > 0.0 && nconit == 0 {
        nconit = 10;
    }

    // 处理 DION0
    let mut dion0 = config.dion0;
    let abpmin = if dion0 < 0.0 {
        let abpmin_val = -dion0;
        dion0 = 1.0;
        abpmin_val
    } else {
        1e-10
    };
    let _ = abpmin;

    // 2. 计算基本量
    let t4 = params.teff.powi(4);
    let totf = SIG4P * t4;
    let abfl0 = SIGE / params.wmm[0];

    let (t0, dmtot, edisc) = if config.idmfix == 1 {
        let t0 = params.teff;
        let dmtot = totf / 0.1;
        (t0, dmtot, totf / dmtot)
    } else {
        let t0 = params.teff;
        let dmtot = totf / (SIGE / params.wmm[0] * 2.0);
        let edisc = totf / dmtot;
        (t0, dmtot, edisc)
    };

    // 3. 计算标高
    let vtb = params.vturb[0];
    let vsnd20: f64 = 2.76e-16 * t0 / params.wmm[0] * dion0 + vtb * vtb;
    let hscalg: f64 = (TWO * vsnd20 / params.qgrav).sqrt();
    let hscalr: f64 = 4.19168946e-10 * totf * abfl0 / params.qgrav;
    let r: f64 = hscalr / hscalg;

    // 诊断输出被简化（无 writer）
    if config.ipring >= 2 {
        eprintln!(" GAS PRESSURE SCALE HEIGHT  = {:+.3E}", hscalg);
        eprintln!(" RAD.PRESSURE SCALE HEIGHT  = {:+.3E}", hscalr);
        eprintln!(" RATIO                      = {:+.3E}", r);
    }

    // 4. 初始化 Eddington 因子
    let mut gamh = UN;
    let _fak0 = THIRD;

    let mut anerel = (dion0 - HALF) / dion0;
    if anerel < ERRT {
        anerel = ERRT;
    }
    let _ = anerel;

    if ndepth == 0 {
        ndepth = params.nd;
    }

    let nd0 = params.nd;
    let mut nd = ndepth;

    // 保存原始 DM
    for id in 0..nd0 {
        work.dm0[id] = params.dm[id];
    }

    // 5. 调用 ZMRHO 计算质量-深度-密度-几何深度
    nd = zmrho(
        r,
        hscalg,
        config.dm1,
        dmtot,
        nd,
        params.dm,
        params.dens,
        params.zd,
    );

    // 6. 初始化迭代
    let mut itgrey = -1;
    let amuv0 = config.dmvisc.powf(config.zeta0 + UN);
    let amuv1 = UN - amuv0;

    // 初始化各种数组
    for id in 0..nd {
        params.pgs[id] = params.dens[id] * vsnd20;

        // 计算粘性相关量
        if params.dm[id] <= config.dmvisc * params.dm[nd - 1] {
            params.viscd[id] = (UN - config.fractv) * (config.zeta1 + UN)
                / config.dmvisc.powf(config.zeta1 + UN)
                * (params.dm[id] / params.dm[nd - 1]).powf(config.zeta1);
            params.theta[id] = (UN - config.fractv)
                * (params.dm[id] / config.dmvisc / params.dm[nd - 1]).powf(config.zeta1 + UN);
        } else {
            params.viscd[id] = config.fractv * (config.zeta0 + UN) / amuv1
                * (params.dm[id] / params.dm[nd - 1]).powf(config.zeta0);
            params.theta[id] = (UN - config.fractv)
                + config.fractv
                    * ((params.dm[id] / params.dm[nd - 1]).powf(config.zeta0 + UN) - amuv0)
                    / amuv1;
        }

        params.gamj[id] = UN;

        // 初始 Rosseland 不透明度和 TAUTHE
        if id == 0 {
            let taur = params.dm[id] * config.abros0;
            params.tauthe[id] = taur * params.theta[id] / (config.zeta1 + TWO);
            params.abrosd[id] = config.abros0;
            params.abplad[id] = config.abpla0;
            params.tauros[id] = taur;
        } else {
            let ddm = params.dm[id] - params.dm[id - 1];
            params.tauros[id] = params.tauros[id - 1] + ddm * params.abrosd[id - 1];
            params.tauthe[id] =
                params.tauthe[id - 1] + ddm * params.abrosd[id - 1] * params.theta[id];
            params.abrosd[id] = params.abrosd[id - 1];
            params.abplad[id] = params.abplad[id - 1];
        }

        // 计算灰大气温度
        let taur = params.tauros[id];
        params.temp[id] = compute_grey_temperature(taur, params.teff);
    }

    // 7. 主迭代循环
    loop {
        itgrey += 1;

        // 更新温度
        for id in 0..nd {
            let taur = if itgrey > 1 {
                params.tauflx[id]
            } else {
                params.tauros[id]
            };
            params.temp[id] = compute_grey_temperature(taur, params.teff);
        }

        // 对流处理
        if config.hmix0 > 0.0 {
            break;
        }

        // 检查迭代结束
        if itgmax == 0 {
            break;
        }

        if itgrey == 0 {
            itgrey = 1;
        }

        // 简化的 RADTOT 计算
        for id in 0..nd {
            params.totj[id] = SIG4P * params.temp[id].powi(4);
            params.toth[id] = totf * (UN - params.theta[id]);
            params.totk[id] = params.totj[id] / 3.0;
            params.rdopac[id] = params.abrosd[id] * params.dens[id];
            params.flopac[id] = params.abrosd[id] * params.toth[id];
        }

        // 插值 TOTH 和 FLOPAC
        for id in 1..nd - 1 {
            let a1 = params.dm[id + 1] - params.dm[id - 1];
            if a1.abs() > 1e-30 {
                let a0 = (params.dm[id] - params.dm[id - 1]) / a1;
                let a1_frac = (params.dm[id + 1] - params.dm[id]) / a1;
                params.toth[id] = a0 * params.toth[id + 1] + a1_frac * params.toth[id];
                params.flopac[id] = a0 * params.flopac[id + 1] + a1_frac * params.flopac[id];
            }
        }
        params.toth[nd - 1] = 0.0;
        params.flopac[nd - 1] = params.flopac[nd - 2];

        // Unsöld-Lucy 温度修正
        let mut dfint = 0.0;
        let mut db0 = 0.0;
        let mut abflxm = config.abflxm;

        for id in 0..nd {
            let hmech = totf * (UN - params.theta[id]);
            let dflux = params.toth[id] - hmech;
            let fkk = if params.totj[id] > 0.0 {
                params.totk[id] / params.totj[id]
            } else {
                THIRD
            };
            let abrad = if params.totj[id] > 0.0 {
                params.rdopac[id] / params.dens[id] / params.totj[id]
            } else {
                params.abrosd[id]
            };

            params.gamj[id] = abrad / params.abplad[id] / fkk * THIRD;

            let abflx = if id != nd - 1 {
                if params.toth[id] > 0.0 {
                    params.flopac[id] / params.toth[id]
                } else {
                    params.abrosd[id]
                }
            } else {
                abflxm
            };

            if id == 0 {
                let fhh = if params.totj[id] > 0.0 {
                    params.toth[id] / params.totj[id]
                } else {
                    1.0
                };
                gamh = fkk / fhh / 0.57753;
                params.tauflx[id] = abflx * params.dm[id];
                params.tauthe[id] = params.tauflx[id] * params.theta[id] / (config.zeta1 + TWO);
                dfint = params.tauflx[id] * dflux;
                db0 = fkk / fhh * dflux;
            } else {
                let zetad = if params.dm[id] <= config.dmvisc * params.dm[nd - 1] {
                    config.zeta1
                } else {
                    config.zeta0
                };

                let ddm = params.dm[id] - params.dm[id - 1];
                if ddm.abs() > 1e-30 {
                    let a0 = (abflxm * params.dm[id] - abflx * params.dm[id - 1])
                        / ddm
                        / (zetad + TWO);
                    let a1 = (abflx - abflxm) / ddm / (zetad + 3.0);

                    params.tauflx[id] = params.tauflx[id - 1] + ddm * HALF * (abflxm + abflx);
                    params.tauthe[id] = params.tauthe[id - 1]
                        + a0 * (params.theta[id] * params.dm[id]
                            - params.theta[id - 1] * params.dm[id - 1])
                        + a1 * (params.theta[id] * params.dm[id].powi(2)
                            - params.theta[id - 1] * params.dm[id - 1].powi(2));
                    dfint = dfint + ddm * HALF * (abflxm * dflux + abflx * dflux);
                }
            }

            abflxm = abflx;

            if itgmax >= 0 {
                let b0 = FOUR * SIG4P * params.temp[id].powi(4);
                let dis = totf * params.viscd[id] / params.abplad[id] / params.dm[nd - 1];
                let db1 = abrad / params.abplad[id] * params.totj[id] - b0 + dis;
                let db = db1 - 3.0 * params.gamj[id] * (db0 + dfint);
                let bnew = FOUR * SIG4P * params.temp[id].powi(4) + db;

                if bnew > 0.0 {
                    params.temp[id] = (bnew / FOUR / SIG4P).powf(0.25);
                }
            }

            if id < nd - 1 {
                db0 = params.gamj[id] * (db0 + dfint);
            }
        }

        if itgrey >= itgmax.abs() {
            break;
        }
    }

    // 8. 插值到最终深度标尺
    if idepth > 0 {
        for i in 0..nd0.min(nd) {
            work.temp0[i] = params.temp[i];
            work.elec0[i] = params.elec[i];
            work.dens0[i] = params.dens[i];
            work.zd0[i] = params.zd[i];
        }

        nd = nd0;
        for i in 0..nd {
            params.dm[i] = work.dm0[i];
            params.temp[i] = work.temp0[i];
            params.elec[i] = work.elec0[i];
            params.dens[i] = work.dens0[i];
            params.zd[i] = work.zd0[i];
        }
    }

    // 9. 重新计算粒子数
    for id in 0..nd {
        let t = params.temp[id];
        let wmm_id = if id < params.wmm.len() { params.wmm[id] } else { 1.0 };
        let an = params.dens[id] / wmm_id + params.elec[id];
        params.ptotal[id] = an * BOLK * t;
        params.pgs[id] = params.ptotal[id];
    }

    LtegrdOutput {
        nd,
        dm: params.dm.to_vec(),
        temp: params.temp.to_vec(),
        elec: params.elec.to_vec(),
        dens: params.dens.to_vec(),
        zd: params.zd.to_vec(),
        ptotal: params.ptotal.to_vec(),
        pgs: params.pgs.to_vec(),
        tauros: params.tauros.to_vec(),
        abrosd: params.abrosd.to_vec(),
        abplad: params.abplad.to_vec(),
        gamh,
        itgrey,
        dmtot,
        edisc,
    }
}

/// 计算灰大气温度分布。
fn compute_grey_temperature(tau: f64, teff: f64) -> f64 {
    // 确保 tau 非负
    let tau = tau.max(0.0);

    let q = if tau < 1e-4 {
        0.5772 // Hopf 函数表面值
    } else if tau < 1.0 {
        0.5772 + 0.4 * tau.powf(0.6)
    } else if tau < 10.0 {
        0.710 + 0.05 * (tau - 1.0)
    } else {
        0.710 + 0.05 * 9.0 + 0.02 * (tau - 10.0).min(90.0)
    };

    // T = Teff * (3/4 * (tau + q))^{1/4}
    // 当 tau = 0 时，T = Teff * (3/4 * 0.5772)^{1/4} ≈ 0.811 * Teff
    teff * (0.75 * (tau + q)).powf(0.25)
}

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

    #[test]
    fn test_compute_grey_temperature() {
        let teff = 10000.0;

        let t0 = compute_grey_temperature(0.0, teff);
        assert!(t0 > 0.0 && t0 < teff);

        let t23 = compute_grey_temperature(2.0 / 3.0, teff);
        assert!((t23 - teff).abs() / teff < 0.1);

        let t1 = compute_grey_temperature(1.0, teff);
        assert!(t1 > teff * 0.9);

        let t10 = compute_grey_temperature(10.0, teff);
        assert!(t10 > t1);
    }

    #[test]
    fn test_ltegrd_config_default() {
        let config = LtegrdConfig::default();
        assert_eq!(config.ndgrey, 0);
        assert_eq!(config.idgrey, 0);
        assert_eq!(config.itgmx0, 5);
    }

    #[test]
    fn test_ltegrd_basic() {
        let config = LtegrdConfig::default();

        let nd = 50;
        let nlevel = 100;

        let wmm = vec![1.0; MDEPTH];
        let mut temp = vec![0.0; MDEPTH];
        let mut elec = vec![0.0; MDEPTH];
        let mut dens = vec![0.0; MDEPTH];
        let mut dm = vec![0.0; MDEPTH];
        let mut zd = vec![0.0; MDEPTH];
        let mut ptotal = vec![0.0; MDEPTH];
        let mut pgs = vec![0.0; MDEPTH];
        let mut tauros = vec![0.0; MDEPTH];
        let mut abrosd = vec![0.4; MDEPTH];
        let mut abplad = vec![0.4; MDEPTH];
        let vturb = vec![5.0; MDEPTH];
        let mut tauthe = vec![0.0; MDEPTH];
        let mut tauflx = vec![0.0; MDEPTH];
        let mut theta = vec![0.0; MDEPTH];
        let mut viscd = vec![0.0; MDEPTH];
        let mut gamj = vec![1.0; MDEPTH];
        let mut totj = vec![0.0; MDEPTH];
        let mut toth = vec![0.0; MDEPTH];
        let mut totk = vec![0.0; MDEPTH];
        let mut rdopac = vec![0.0; MDEPTH];
        let mut flopac = vec![0.0; MDEPTH];

        let mut params = LtegrdParams {
            config,
            nd,
            nlevel,
            teff: 35000.0,
            qgrav: 1e4,
            wmm: &wmm,
            temp: &mut temp,
            elec: &mut elec,
            dens: &mut dens,
            dm: &mut dm,
            zd: &mut zd,
            ptotal: &mut ptotal,
            pgs: &mut pgs,
            tauros: &mut tauros,
            abrosd: &mut abrosd,
            abplad: &mut abplad,
            vturb: &vturb,
            tauthe: &mut tauthe,
            tauflx: &mut tauflx,
            theta: &mut theta,
            viscd: &mut viscd,
            gamj: &mut gamj,
            totj: &mut totj,
            toth: &mut toth,
            totk: &mut totk,
            rdopac: &mut rdopac,
            flopac: &mut flopac,
        };

        let result = ltegrd_pure(&mut params);

        assert!(result.nd > 0);
        assert!(!result.dm.is_empty());
        assert!(!result.temp.is_empty());

        for i in 0..result.nd.min(10) {
            assert!(result.temp[i] > 0.0, "Temperature at {} should be positive", i);
        }
    }
}