//! 输出计算好的大气模型。
//!
//! 重构自 TLUSTY `output.f`
//! 将模型写入文件，可作为后续运行的输入模型。

use crate::tlusty::io::{FortranWriter, Result, format_exp_fortran};
use crate::tlusty::state::config::TlustyConfig;
use crate::tlusty::state::model::ModelState;
use std::io::Write;

/// OUTPUT 参数结构体
pub struct OutputParams<'a> {
    /// 配置
    pub config: &'a TlustyConfig,
    /// 模型状态
    pub model: &'a ModelState,
}

/// 输出模型到文件。
///
/// # 参数
/// * `writer7` - 主输出写入器 (fort.7)
/// * `params` - 输出参数
/// * `writer17` - 诊断输出写入器 (fort.17)，IPRIND>0 时使用
/// * `writer20` - BFAC 数据写入器 (fort.20)，NLTE 且 IPRIND>0 时使用
///
/// # 返回值
/// 成功返回 Ok(())
pub fn output<W: Write>(
    writer7: &mut FortranWriter<W>,
    params: &OutputParams,
    writer17: Option<&mut FortranWriter<W>>,
    mut writer20: Option<&mut FortranWriter<W>>,
) -> Result<()> {
    let config = params.config;
    let model = params.model;

    let nd = config.basnum.nd as usize;
    let nlevel = config.basnum.nlevel as usize;
    let idisk = config.basnum.idisk;
    let ifmol = config.basnum.ifmol;
    let lte = config.inppar.lte;
    let iprinp = config.prints.iprinp;
    let iprind = config.prints.iprind;

    // 计算 NUMLT 和 NUMPAR
    let mut numlt: i32 = 3;
    if idisk == 1 {
        numlt = 4;
    }
    if ifmol == 1 {
        numlt += 1;
    }

    let mut numpar: i32 = nlevel as i32 + numlt;
    if lte && iprinp == 0 {
        numpar = numlt;
    }
    if ifmol > 0 {
        numpar = -numpar;
    }

    // 写入头部: ND, NUMPAR (unit 7)
    writer7.write_raw(&format!("{:5}{:5}", nd, numpar))?;
    writer7.write_newline()?;

    // 写入 DM 数组 (6 个一组, FORMAT 1P6E13.6)
    write_dm_array(writer7, &model.modpar.dm, nd)?;

    // 写入每个深度点的数据 (unit 7)
    write_depth_data(writer7, model, nd, nlevel, idisk, ifmol, lte, iprinp)?;

    // IPRIND > 0: 写诊断输出到 unit 17 和 unit 20
    if iprind > 0 {
        if let Some(w17) = writer17 {
            w17.write_raw(&format!("{:5}{:5}", nd, numpar))?;
            w17.write_newline()?;
            write_dm_array(w17, &model.modpar.dm, nd)?;

            if idisk == 0 {
                if lte {
                    for id in 0..nd {
                        write_depth_line(w17, &[
                            model.modpar.temp[id],
                            model.modpar.elec[id],
                            model.modpar.dens[id],
                        ])?;
                    }
                } else {
                    // NLTE: 写 unit 20 头部和 DM
                    if let Some(w20) = writer20.as_mut() {
                        w20.write_raw(&format!("{:5}{:5}", nd, numpar))?;
                        w20.write_newline()?;
                        write_dm_array(w20, &model.modpar.dm, nd)?;
                    }
                    for id in 0..nd {
                        write_depth_line_with_popul(
                            w17, model.modpar.temp[id], model.modpar.elec[id],
                            model.modpar.dens[id], None,
                            &model.levpop.popul, id, nlevel,
                        )?;
                        if let Some(w20) = writer20.as_mut() {
                            write_depth_line_with_popul(
                                w20, model.modpar.temp[id], model.modpar.elec[id],
                                model.modpar.dens[id], None,
                                &model.levpop.bfac, id, nlevel,
                            )?;
                        }
                    }
                }
            } else {
                // 圆盘模型
                if lte {
                    for id in 0..nd {
                        write_depth_line(w17, &[
                            model.modpar.temp[id],
                            model.modpar.elec[id],
                            model.modpar.dens[id],
                            model.modpar.zd[id],
                        ])?;
                    }
                } else {
                    if let Some(w20) = writer20.as_mut() {
                        w20.write_raw(&format!("{:5}{:5}", nd, numpar))?;
                        w20.write_newline()?;
                        write_dm_array(w20, &model.modpar.dm, nd)?;
                    }
                    for id in 0..nd {
                        write_depth_line_with_popul_disk(
                            w17, model.modpar.temp[id], model.modpar.elec[id],
                            model.modpar.dens[id], None, model.modpar.zd[id],
                            &model.levpop.popul, id, nlevel,
                        )?;
                        if let Some(w20) = writer20.as_mut() {
                            write_depth_line_with_popul_disk(
                                w20, model.modpar.temp[id], model.modpar.elec[id],
                                model.modpar.dens[id], None, model.modpar.zd[id],
                                &model.levpop.bfac, id, nlevel,
                            )?;
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

/// 写入 DM 数组 (6 个一组, FORMAT 1P6E13.6)
fn write_dm_array<W: Write>(writer: &mut FortranWriter<W>, dm: &[f64], nd: usize) -> Result<()> {
    let mut dm_line = String::new();
    for (i, &dm_val) in dm.iter().take(nd).enumerate() {
        dm_line.push_str(&format_exp_fortran(dm_val, 13, 6, false));
        if (i + 1) % 6 == 0 || i == nd - 1 {
            writer.write_raw(&dm_line)?;
            writer.write_newline()?;
            dm_line.clear();
        }
    }
    Ok(())
}

/// 写入每个深度点的数据（unit 7 主输出）
fn write_depth_data<W: Write>(
    writer: &mut FortranWriter<W>,
    model: &ModelState,
    nd: usize,
    nlevel: usize,
    idisk: i32,
    ifmol: i32,
    lte: bool,
    iprinp: i32,
) -> Result<()> {
    for id in 0..nd {
        let temp = model.modpar.temp[id];
        let elec = model.modpar.elec[id];
        let dens = model.modpar.dens[id];

        if idisk == 0 {
            // 平面平行模型
            if lte && iprinp == 0 {
                if ifmol == 0 {
                    write_depth_line(writer, &[temp, elec, dens])?;
                } else {
                    let totn = model.modpar.totn[id];
                    write_depth_line(writer, &[temp, elec, dens, totn])?;
                }
            } else {
                if ifmol == 0 {
                    write_depth_line_with_popul(writer, temp, elec, dens, None, &model.levpop.popul, id, nlevel)?;
                } else {
                    let totn = model.modpar.totn[id];
                    write_depth_line_with_popul(writer, temp, elec, dens, Some(totn), &model.levpop.popul, id, nlevel)?;
                }
            }
        } else {
            // 圆盘模型
            let zd = model.modpar.zd[id];
            if lte && iprinp == 0 {
                if ifmol == 0 {
                    write_depth_line(writer, &[temp, elec, dens, zd])?;
                } else {
                    let totn = model.modpar.totn[id];
                    write_depth_line(writer, &[temp, elec, dens, totn, zd])?;
                }
            } else {
                if ifmol == 0 {
                    write_depth_line_with_popul_disk(writer, temp, elec, dens, None, zd, &model.levpop.popul, id, nlevel)?;
                } else {
                    let totn = model.modpar.totn[id];
                    write_depth_line_with_popul_disk(writer, temp, elec, dens, Some(totn), zd, &model.levpop.popul, id, nlevel)?;
                }
            }
        }
    }
    Ok(())
}

/// 写入深度点数据（无占据数）
/// 使用 Fortran 格式 1P5E15.6
fn write_depth_line<W: Write>(writer: &mut FortranWriter<W>, values: &[f64]) -> Result<()> {
    let mut line = String::new();
    for &val in values {
        line.push_str(&format_exp_fortran(val, 15, 6, false));
    }
    writer.write_raw(&line)?;
    writer.write_newline()?;
    Ok(())
}

/// 写入深度点数据（带占据数，平面平行）
/// 使用 Fortran 格式 1P5E15.6
fn write_depth_line_with_popul<W: Write>(
    writer: &mut FortranWriter<W>,
    temp: f64,
    elec: f64,
    dens: f64,
    totn: Option<f64>,
    popul: &[Vec<f64>],
    id: usize,
    nlevel: usize,
) -> Result<()> {
    // Fortran: WRITE(7,503) TEMP(ID),ELEC(ID),DENS(ID),(POPUL(J,ID),J=1,NLEVEL)
    // FORMAT(1P5E15.6) - 5 values per line, flowing continuously
    // So first line has temp, elec, dens, pop[0], pop[1]
    // then subsequent lines have 5 pop values each

    let mut line = String::new();
    let mut col = 0;

    // Write temp, elec, dens
    for &val in &[temp, elec, dens] {
        line.push_str(&format_exp_fortran(val, 15, 6, false));
        col += 1;
    }
    if let Some(t) = totn {
        line.push_str(&format_exp_fortran(t, 15, 6, false));
        col += 1;
    }

    // Write populations, continuing from where temp/elec/dens left off
    for j in 0..nlevel {
        line.push_str(&format_exp_fortran(popul[j][id], 15, 6, false));
        col += 1;
        if col % 5 == 0 || j == nlevel - 1 {
            writer.write_raw(&line)?;
            writer.write_newline()?;
            line.clear();
        }
    }

    Ok(())
}

/// 写入深度点数据（带占据数，圆盘）
/// 使用 Fortran 格式 1P5E15.6
fn write_depth_line_with_popul_disk<W: Write>(
    writer: &mut FortranWriter<W>,
    temp: f64,
    elec: f64,
    dens: f64,
    totn: Option<f64>,
    zd: f64,
    popul: &[Vec<f64>],
    id: usize,
    nlevel: usize,
) -> Result<()> {
    // Same flow as planar: values flow continuously, 5 per line
    let mut line = String::new();
    let mut col = 0;

    for &val in &[temp, elec, dens] {
        line.push_str(&format_exp_fortran(val, 15, 6, false));
        col += 1;
    }
    if let Some(t) = totn {
        line.push_str(&format_exp_fortran(t, 15, 6, false));
        col += 1;
    }
    line.push_str(&format_exp_fortran(zd, 15, 6, false));
    col += 1;

    for j in 0..nlevel {
        line.push_str(&format_exp_fortran(popul[j][id], 15, 6, false));
        col += 1;
        if col % 5 == 0 || j == nlevel - 1 {
            writer.write_raw(&line)?;
            writer.write_newline()?;
            line.clear();
        }
    }

    Ok(())
}

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

    fn create_test_state() -> (TlustyConfig, ModelState) {
        let mut config = TlustyConfig::default();
        config.basnum.nd = 5;
        config.basnum.nlevel = 3;
        config.basnum.idisk = 0;
        config.basnum.ifmol = 0;
        config.inppar.lte = true;
        config.prints.iprinp = 0;

        let mut model = ModelState::new();
        for i in 0..5 {
            model.modpar.dm[i] = (i + 1) as f64 * 1e2;
            model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
            model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
            model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
        }

        (config, model)
    }

    #[test]
    fn test_output_lte_plane_parallel() {
        let (config, model) = create_test_state();

        let mut writer = FortranWriter::to_memory();

        let params = OutputParams {
            config: &config,
            model: &model,
        };

        let result = output(&mut writer, &params, None, None);
        assert!(result.is_ok());

        let output_str = writer.into_string().unwrap();

        // 验证头部
        let lines: Vec<&str> = output_str.lines().collect();
        assert_eq!(lines[0], "    5    3"); // ND=5, NUMPAR=3 (LTE, plane-parallel)

        // 验证 DM 数组 (6 个一组)
        assert!(lines[1].contains("1.000000E+02"));

        // 验证深度数据
        assert!(lines[2].contains("1.000000E+04")); // temp
        assert!(lines[2].contains("1.000000E+12")); // elec
        assert!(lines[2].contains("1.000000E-07")); // dens
    }

    #[test]
    fn test_output_nlte_with_populations() {
        let mut config = TlustyConfig::default();
        config.basnum.nd = 3;
        config.basnum.nlevel = 7; // 测试多行占据数
        config.basnum.idisk = 0;
        config.basnum.ifmol = 0;
        config.inppar.lte = false; // NLTE
        config.prints.iprinp = 1; // 需要输出占据数

        let mut model = ModelState::new();
        for i in 0..3 {
            model.modpar.dm[i] = (i + 1) as f64 * 1e2;
            model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
            model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
            model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
        }
        // 设置占据数
        for j in 0..7 {
            for i in 0..3 {
                model.levpop.popul[j][i] = (j + 1) as f64 * 1e10 + i as f64;
            }
        }

        let mut writer = FortranWriter::to_memory();

        let params = OutputParams {
            config: &config,
            model: &model,
        };

        let result = output(&mut writer, &params, None, None);
        assert!(result.is_ok());

        let output_str = writer.into_string().unwrap();
        let lines: Vec<&str> = output_str.lines().collect();

        // 验证头部: NUMPAR = NLEVEL + NUMLT = 7 + 3 = 10
        assert_eq!(lines[0], "    3   10");

        // 验证 DM 数组
        assert!(lines[1].contains("1.000000E+02"));

        // 验证第一深度点的数据
        // Line 2: temp, elec, dens (3 个值)
        assert!(lines[2].contains("1.000000E+04")); // temp
        assert!(lines[2].contains("1.000000E+12")); // elec
        assert!(lines[2].contains("1.000000E-07")); // dens

        // Line 3-4: 占据数 (7 个，每行最多 5 个，但第一行只有 2 个空位)
        // 实际上，由于第一行已经写了 3 个值，占据数从新行开始
        assert!(lines[3].contains("1.000000E+10")); // popul[0][0]
        assert!(lines[3].contains("2.000000E+10")); // popul[1][0]
    }

    #[test]
    fn test_output_disk_model() {
        let mut config = TlustyConfig::default();
        config.basnum.nd = 3;
        config.basnum.nlevel = 2;
        config.basnum.idisk = 1; // 圆盘模型
        config.basnum.ifmol = 0;
        config.inppar.lte = true;
        config.prints.iprinp = 0;

        let mut model = ModelState::new();
        for i in 0..3 {
            model.modpar.dm[i] = (i + 1) as f64 * 1e2;
            model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
            model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
            model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
            model.modpar.zd[i] = (i + 1) as f64 * 1e5;
        }

        let mut writer = FortranWriter::to_memory();

        let params = OutputParams {
            config: &config,
            model: &model,
        };

        let result = output(&mut writer, &params, None, None);
        assert!(result.is_ok());

        let output_str = writer.into_string().unwrap();
        let lines: Vec<&str> = output_str.lines().collect();

        // 验证头部: NUMPAR = NUMLT = 4 (LTE, disk)
        assert_eq!(lines[0], "    3    4");

        // 验证深度数据包含 ZD
        assert!(lines[2].contains("1.000000E+05")); // zd
    }

    #[test]
    fn test_output_with_molecules() {
        let mut config = TlustyConfig::default();
        config.basnum.nd = 3;
        config.basnum.nlevel = 2;
        config.basnum.idisk = 0;
        config.basnum.ifmol = 1; // 有分子
        config.inppar.lte = true;
        config.prints.iprinp = 0;

        let mut model = ModelState::new();
        for i in 0..3 {
            model.modpar.dm[i] = (i + 1) as f64 * 1e2;
            model.modpar.temp[i] = 10000.0 - i as f64 * 1000.0;
            model.modpar.elec[i] = 1e12 - i as f64 * 1e11;
            model.modpar.dens[i] = 1e-7 + i as f64 * 1e-8;
            model.modpar.totn[i] = (i + 1) as f64 * 1e14;
        }

        let mut writer = FortranWriter::to_memory();

        let params = OutputParams {
            config: &config,
            model: &model,
        };

        let result = output(&mut writer, &params, None, None);
        assert!(result.is_ok());

        let output_str = writer.into_string().unwrap();
        let lines: Vec<&str> = output_str.lines().collect();

        // 验证头部: NUMPAR = -4 (负数表示有分子)
        assert_eq!(lines[0], "    3   -4");

        // 验证深度数据包含 TOTN
        assert!(lines[2].contains("1.000000E+14")); // totn
    }
}