//! LAMMPS data and settings file writer.
//!
//! Generates LAMMPS-compatible input files from parameterized molecular
//! systems. Supports both periodic and non-periodic boundary conditions,
//! with appropriate electrostatics handling (PPPM for periodic, direct
//! summation for non-periodic).

use crate::io::error::Error;
use crate::model::system::Bond;
use crate::model::topology::{
    AnglePotential, BondPotential, ForgedSystem, ImproperPotential, Potentials, VdwPairPotential,
};
use std::collections::{HashMap, HashSet};
use std::io::Write;

/// Configuration for LAMMPS output generation.
///
/// Controls cutoffs, boundary conditions, and box sizing for the
/// generated simulation input files.
///
/// # Examples
///
/// ```
/// use dreid_forge::io::{LammpsConfig, SystemType};
///
/// // Periodic system with standard cutoffs
/// let config = LammpsConfig {
///     system_type: SystemType::Periodic,
///     nonbonded_cutoff: 12.0,
///     ..Default::default()
/// };
///
/// // Non-periodic (gas phase) system
/// let gas_config = LammpsConfig {
///     system_type: SystemType::NonPeriodic,
///     aabb_margin: 10.0,
///     ..Default::default()
/// };
/// ```
#[derive(Debug, Clone)]
pub struct LammpsConfig {
    /// Cutoff distance for non-bonded interactions (Å).
    pub nonbonded_cutoff: f64,
    /// Inner cutoff for hydrogen bond potential (Å).
    pub hbond_cutoff_inner: f64,
    /// Outer cutoff for hydrogen bond potential (Å).
    pub hbond_cutoff_outer: f64,
    /// Angular cutoff for hydrogen bond geometry (degrees).
    pub hbond_angle_cutoff: f64,
    /// Margin added to bounding box for non-periodic systems (Å).
    pub aabb_margin: f64,
    /// Boundary condition type (periodic or non-periodic).
    pub system_type: SystemType,
}

/// Boundary condition type for LAMMPS simulations.
///
/// Determines how electrostatics are handled and what pair styles
/// are used for long-range interactions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SystemType {
    /// Periodic boundary conditions with PPPM electrostatics.
    Periodic,
    /// Non-periodic (shrink-wrapped) boundaries with direct Coulomb.
    NonPeriodic,
}

impl Default for LammpsConfig {
    fn default() -> Self {
        Self {
            nonbonded_cutoff: 12.0,
            hbond_cutoff_inner: 10.0,
            hbond_cutoff_outer: 12.0,
            hbond_angle_cutoff: 90.0,
            aabb_margin: 5.0,
            system_type: SystemType::Periodic,
        }
    }
}

/// Writes both LAMMPS data and settings files.
///
/// This is the primary entry point for generating complete LAMMPS
/// simulation input. Writes a data file (topology + coordinates)
/// and a settings file (force field parameters) that can be used
/// together in a LAMMPS input script.
///
/// # Arguments
///
/// * `data_writer` — Output for the data file (`.data` or `.lammps`)
/// * `settings_writer` — Output for the settings file (`.in.settings`)
/// * `forged` — Parameterized system from [`forge`](crate::forge)
/// * `config` — Output configuration options
///
/// # Errors
///
/// Returns [`Error`] if the system has validation issues or if
/// I/O operations fail.
///
/// # Examples
///
/// ```no_run
/// use dreid_forge::io::{write_lammps_package, LammpsConfig};
/// use dreid_forge::ForgedSystem;
/// use std::fs::File;
///
/// let forged: ForgedSystem = todo!(); // from forge()
/// let mut data = File::create("system.data")?;
/// let mut settings = File::create("system.in.settings")?;
///
/// write_lammps_package(&mut data, &mut settings, &forged, &LammpsConfig::default())?;
/// # Ok::<(), dreid_forge::io::Error>(())
/// ```
pub fn write_package<W1: Write, W2: Write>(
    mut data_writer: W1,
    mut settings_writer: W2,
    forged: &ForgedSystem,
    config: &LammpsConfig,
) -> Result<(), Error> {
    write_data_file(&mut data_writer, forged, config)?;
    write_settings_file(&mut settings_writer, forged, config)?;
    Ok(())
}

/// Writes a LAMMPS data file.
///
/// Generates a topology file containing atom coordinates, types,
/// charges, molecular connectivity, and simulation box dimensions.
/// This file is read by LAMMPS using the `read_data` command.
///
/// # Arguments
///
/// * `writer` — Output destination for the data file
/// * `forged` — Parameterized system from [`forge`](crate::forge)
/// * `config` — Output configuration options
///
/// # Errors
///
/// Returns [`Error`] if validation fails or I/O operations fail.
pub fn write_data_file<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    config: &LammpsConfig,
) -> Result<(), Error> {
    validate_forged(forged)?;

    let mapping = TypeMapping::new(forged)?;
    let molecules = build_molecules(&forged.system.bonds, forged.system.atom_count());
    let box_bounds = compute_bounds(&forged.system, config.aabb_margin);

    writeln!(writer, "# LAMMPS data file generated by dreid-forge")?;
    writeln!(writer)?;
    writeln!(writer, "{} atoms", forged.system.atom_count())?;
    writeln!(writer, "{} bonds", forged.potentials.bonds.len())?;
    writeln!(writer, "{} angles", forged.potentials.angles.len())?;
    writeln!(writer, "{} dihedrals", forged.potentials.dihedrals.len())?;
    writeln!(writer, "{} impropers", forged.potentials.impropers.len())?;
    writeln!(writer)?;
    writeln!(writer, "{} atom types", forged.atom_types.len())?;
    writeln!(writer, "{} bond types", mapping.bond_types.len())?;
    writeln!(writer, "{} angle types", mapping.angle_types.len())?;
    writeln!(writer, "{} dihedral types", mapping.dihedral_types.len())?;
    writeln!(writer, "{} improper types", mapping.improper_types.len())?;
    writeln!(writer)?;
    writeln!(
        writer,
        "{:.6} {:.6} xlo xhi",
        box_bounds.xlo, box_bounds.xhi
    )?;
    writeln!(
        writer,
        "{:.6} {:.6} ylo yhi",
        box_bounds.ylo, box_bounds.yhi
    )?;
    writeln!(
        writer,
        "{:.6} {:.6} zlo zhi",
        box_bounds.zlo, box_bounds.zhi
    )?;
    if let Some(tilt) = box_bounds.tilt {
        writeln!(
            writer,
            "{:.6} {:.6} {:.6} xy xz yz",
            tilt.xy, tilt.xz, tilt.yz
        )?;
    }
    writeln!(writer)?;

    write_masses(writer, forged)?;
    write_bond_coeffs(writer, &mapping)?;
    write_angle_coeffs(writer, &mapping)?;
    write_dihedral_coeffs(writer, &mapping)?;
    write_improper_coeffs(writer, &mapping)?;
    write_atoms(writer, forged, &molecules)?;
    write_bonds(writer, forged, &mapping)?;
    write_angles(writer, forged, &mapping)?;
    write_dihedrals(writer, forged, &mapping)?;
    write_impropers(writer, forged, &mapping)?;

    Ok(())
}

/// Writes a LAMMPS settings file.
///
/// Generates force field parameter commands including style
/// definitions and coefficient specifications. This file is
/// typically included in a LAMMPS input script using `include`.
///
/// # Arguments
///
/// * `writer` — Output destination for the settings file
/// * `forged` — Parameterized system from [`forge`](crate::forge)
/// * `config` — Output configuration options
///
/// # Errors
///
/// Returns [`Error`] if validation fails or I/O operations fail.
pub fn write_settings_file<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    config: &LammpsConfig,
) -> Result<(), Error> {
    validate_forged(forged)?;
    let potentials = &forged.potentials;
    let has_lj = potentials
        .vdw_pairs
        .iter()
        .any(|p| matches!(p, VdwPairPotential::LennardJones { .. }));
    let has_buck = potentials
        .vdw_pairs
        .iter()
        .any(|p| matches!(p, VdwPairPotential::Exponential6 { .. }));
    let has_hbond = !potentials.h_bonds.is_empty();

    writeln!(writer, "# LAMMPS settings generated by dreid-forge")?;
    writeln!(writer, "units           real")?;
    writeln!(writer, "atom_style      full")?;
    writeln!(writer, "special_bonds   lj/coul 0.0 0.0 1.0")?;
    writeln!(writer, "pair_modify     mix geometric")?;

    let bond_styles = collect_bond_styles(potentials);
    let angle_styles = collect_angle_styles(potentials);
    let dihedral_styles = collect_dihedral_styles(potentials);
    let improper_styles = collect_improper_styles(potentials);

    write_style_line(writer, "bond_style", &bond_styles)?;
    write_style_line(writer, "angle_style", &angle_styles)?;
    write_style_line(writer, "dihedral_style", &dihedral_styles)?;
    write_style_line(writer, "improper_style", &improper_styles)?;

    if has_lj || has_buck || has_hbond {
        let (lj_keyword, buck_keyword, kspace_line) = match config.system_type {
            SystemType::Periodic => (
                "lj/cut/coul/long",
                "buck/coul/long",
                Some("kspace_style    pppm 1.0e-4"),
            ),
            SystemType::NonPeriodic => ("lj/cut/coul/cut", "buck/coul/cut", None),
        };

        let mut pair_styles = Vec::new();
        if has_lj {
            pair_styles.push(format!("{lj_keyword} {:.3}", config.nonbonded_cutoff));
        }
        if has_buck {
            pair_styles.push(format!("{buck_keyword} {:.3}", config.nonbonded_cutoff));
        }
        if has_hbond {
            pair_styles.push(format!(
                "hbond/dreiding/lj 2 {:.3} {:.3} {:.3}",
                config.hbond_cutoff_inner, config.hbond_cutoff_outer, config.hbond_angle_cutoff
            ));
        }

        if pair_styles.len() == 1 {
            writeln!(writer, "pair_style      {}", pair_styles[0])?;
        } else {
            writeln!(
                writer,
                "pair_style      hybrid/overlay {}",
                pair_styles.join(" ")
            )?;
        }
        if let Some(kspace) = kspace_line {
            writeln!(writer, "{}", kspace)?;
        }
    }

    writeln!(writer)?;
    writeln!(writer, "# Pair coefficients")?;
    let (lj_keyword, buck_keyword) = match config.system_type {
        SystemType::Periodic => ("lj/cut/coul/long", "buck/coul/long"),
        SystemType::NonPeriodic => ("lj/cut/coul/cut", "buck/coul/cut"),
    };

    for p in &potentials.vdw_pairs {
        match p {
            VdwPairPotential::LennardJones {
                type1_idx,
                type2_idx,
                sigma,
                epsilon,
            } => writeln!(
                writer,
                "pair_coeff {:>3} {:>3} {lj_keyword} {:.6} {:.6}",
                type1_idx + 1,
                type2_idx + 1,
                epsilon,
                sigma
            )?,
            VdwPairPotential::Exponential6 {
                type1_idx,
                type2_idx,
                a,
                b,
                c,
            } => writeln!(
                writer,
                "pair_coeff {:>3} {:>3} {buck_keyword} {:.6} {:.6} {:.6}",
                type1_idx + 1,
                type2_idx + 1,
                a,
                b,
                c
            )?,
        }
    }

    if has_hbond {
        writeln!(writer)?;
        writeln!(writer, "# Hydrogen bond coefficients")?;
        for hb in &potentials.h_bonds {
            let (i, j, donor_flag) = if hb.donor_type_idx <= hb.acceptor_type_idx {
                (hb.donor_type_idx, hb.acceptor_type_idx, "i")
            } else {
                (hb.acceptor_type_idx, hb.donor_type_idx, "j")
            };
            let epsilon = hb.d0;
            let sigma = hb.r0 / 1.0954_f64.sqrt();
            let n = 2; // Improved: cos^2(θ) (Original Dreiding uses cos^4(θ))
            writeln!(
                writer,
                "pair_coeff {:>3} {:>3} hbond/dreiding/lj {} {} {:.6} {:.6} {}",
                i + 1,
                j + 1,
                hb.hydrogen_type_idx + 1,
                donor_flag,
                epsilon,
                sigma,
                n
            )?;
        }
    }

    Ok(())
}

fn write_style_line<W: Write>(
    writer: &mut W,
    keyword: &str,
    styles: &[&'static str],
) -> Result<(), Error> {
    if styles.is_empty() {
        return Ok(());
    }
    if styles.len() == 1 {
        writeln!(writer, "{:14} {}", keyword, styles[0])?;
    } else {
        writeln!(writer, "{:14} hybrid {}", keyword, styles.join(" "))?;
    }
    Ok(())
}

fn collect_bond_styles(p: &Potentials) -> Vec<&'static str> {
    let mut set = HashSet::new();
    for b in &p.bonds {
        match b {
            BondPotential::Harmonic { .. } => {
                set.insert("harmonic");
            }
            BondPotential::Morse { .. } => {
                set.insert("morse");
            }
        }
    }
    collect_sorted(set)
}

fn collect_angle_styles(p: &Potentials) -> Vec<&'static str> {
    let mut set = HashSet::new();
    for a in &p.angles {
        match a {
            AnglePotential::CosineHarmonic { .. } => {
                set.insert("cosine/squared");
            }
            AnglePotential::ThetaHarmonic { .. } => {
                set.insert("harmonic");
            }
        }
    }
    collect_sorted(set)
}

fn collect_dihedral_styles(p: &Potentials) -> Vec<&'static str> {
    if p.dihedrals.is_empty() {
        return Vec::new();
    }
    vec!["charmm"]
}

fn collect_improper_styles(p: &Potentials) -> Vec<&'static str> {
    let mut set = HashSet::new();
    for imp in &p.impropers {
        match imp {
            ImproperPotential::Planar { .. } | ImproperPotential::Umbrella { .. } => {
                set.insert("umbrella");
            }
        }
    }
    collect_sorted(set)
}

fn collect_sorted(mut set: HashSet<&'static str>) -> Vec<&'static str> {
    let mut styles: Vec<_> = set.drain().collect();
    styles.sort_unstable();
    styles
}

fn write_masses<W: Write>(writer: &mut W, forged: &ForgedSystem) -> Result<(), Error> {
    writeln!(writer, "Masses")?;
    writeln!(writer)?;

    let mut mass_by_type: HashMap<usize, f64> = HashMap::new();
    for atom in &forged.atom_properties {
        mass_by_type.entry(atom.type_idx).or_insert(atom.mass);
    }

    for (idx, name) in forged.atom_types.iter().enumerate() {
        let mass = mass_by_type
            .get(&idx)
            .copied()
            .ok_or_else(|| Error::Conversion(format!("missing mass for atom type {idx}")))?;
        writeln!(writer, "{:>4} {:>12.6}  # {}", idx + 1, mass, name)?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_bond_coeffs<W: Write>(writer: &mut W, mapping: &TypeMapping) -> Result<(), Error> {
    if mapping.bond_types.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Bond Coeffs")?;
    writeln!(writer)?;
    for (id, key) in &mapping.bond_types {
        match key {
            BondKey::Harmonic { k_force, r0 } => {
                writeln!(writer, "{:>4} harmonic {:.6} {:.6}", id, k_force.0, r0.0)?;
            }
            BondKey::Morse { r0, d0, alpha } => {
                writeln!(
                    writer,
                    "{:>4} morse {:.6} {:.6} {:.6}",
                    id, d0.0, alpha.0, r0.0
                )?;
            }
        }
    }
    writeln!(writer)?;
    Ok(())
}

fn write_angle_coeffs<W: Write>(writer: &mut W, mapping: &TypeMapping) -> Result<(), Error> {
    if mapping.angle_types.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Angle Coeffs")?;
    writeln!(writer)?;
    for (id, key) in &mapping.angle_types {
        match key {
            AngleKey::CosineHarmonic { k_force, theta0 } => {
                writeln!(
                    writer,
                    "{:>4} cosine/squared {:.6} {:.6}",
                    id, k_force.0, theta0.0
                )?;
            }
            AngleKey::ThetaHarmonic { k_force, theta0 } => {
                writeln!(
                    writer,
                    "{:>4} harmonic {:.6} {:.6}",
                    id, k_force.0, theta0.0
                )?;
            }
        }
    }
    writeln!(writer)?;
    Ok(())
}

fn write_dihedral_coeffs<W: Write>(writer: &mut W, mapping: &TypeMapping) -> Result<(), Error> {
    if mapping.dihedral_types.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Dihedral Coeffs")?;
    writeln!(writer)?;
    for (id, key) in &mapping.dihedral_types {
        writeln!(
            writer,
            "{:>4} charmm {:.6} {} {:.6} 0.0",
            id, key.v_barrier.0, key.periodicity, key.phase_offset.0
        )?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_improper_coeffs<W: Write>(writer: &mut W, mapping: &TypeMapping) -> Result<(), Error> {
    if mapping.improper_types.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Improper Coeffs")?;
    writeln!(writer)?;
    for (id, key) in &mapping.improper_types {
        match key {
            ImproperKey::Planar { k_force, chi0 } => {
                writeln!(writer, "{:>4} umbrella {:.6} {:.6}", id, k_force.0, chi0.0)?;
            }
            ImproperKey::Umbrella { k_force, psi0 } => {
                writeln!(writer, "{:>4} umbrella {:.6} {:.6}", id, k_force.0, psi0.0)?;
            }
        }
    }
    writeln!(writer)?;
    Ok(())
}

fn write_atoms<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    molecules: &[usize],
) -> Result<(), Error> {
    writeln!(writer, "Atoms # id mol-id type-id charge x y z")?;
    writeln!(writer)?;
    for (idx, atom) in forged.system.atoms.iter().enumerate() {
        let param = &forged.atom_properties[idx];
        writeln!(
            writer,
            "{:>6} {:>6} {:>4} {:>12.6} {:>12.6} {:>12.6} {:>12.6}",
            idx + 1,
            molecules[idx],
            param.type_idx + 1,
            param.charge,
            atom.position[0],
            atom.position[1],
            atom.position[2]
        )?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_bonds<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    mapping: &TypeMapping,
) -> Result<(), Error> {
    if forged.potentials.bonds.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Bonds")?;
    writeln!(writer)?;
    for (idx, bond) in forged.potentials.bonds.iter().enumerate() {
        let (i, j, key) = match bond {
            BondPotential::Harmonic { i, j, k_force, r0 } => (
                *i,
                *j,
                BondKey::Harmonic {
                    k_force: F64Key(*k_force),
                    r0: F64Key(*r0),
                },
            ),
            BondPotential::Morse {
                i,
                j,
                r0,
                d0,
                alpha,
            } => (
                *i,
                *j,
                BondKey::Morse {
                    r0: F64Key(*r0),
                    d0: F64Key(*d0),
                    alpha: F64Key(*alpha),
                },
            ),
        };
        let t = mapping
            .bond_type_lookup
            .get(&key)
            .ok_or_else(|| Error::Conversion("bond type not mapped".into()))?;
        writeln!(writer, "{:>6} {:>4} {:>6} {:>6}", idx + 1, t, i + 1, j + 1)?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_angles<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    mapping: &TypeMapping,
) -> Result<(), Error> {
    if forged.potentials.angles.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Angles")?;
    writeln!(writer)?;
    for (idx, angle) in forged.potentials.angles.iter().enumerate() {
        let (i, j, k, key) = match angle {
            AnglePotential::CosineHarmonic {
                i,
                j,
                k,
                k_force,
                theta0,
            } => (
                *i,
                *j,
                *k,
                AngleKey::CosineHarmonic {
                    k_force: F64Key(*k_force),
                    theta0: F64Key(*theta0),
                },
            ),
            AnglePotential::ThetaHarmonic {
                i,
                j,
                k,
                k_force,
                theta0,
            } => (
                *i,
                *j,
                *k,
                AngleKey::ThetaHarmonic {
                    k_force: F64Key(*k_force),
                    theta0: F64Key(*theta0),
                },
            ),
        };
        let t = mapping
            .angle_type_lookup
            .get(&key)
            .ok_or_else(|| Error::Conversion("angle type not mapped".into()))?;
        writeln!(
            writer,
            "{:>6} {:>4} {:>6} {:>6} {:>6}",
            idx + 1,
            t,
            i + 1,
            j + 1,
            k + 1
        )?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_dihedrals<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    mapping: &TypeMapping,
) -> Result<(), Error> {
    if forged.potentials.dihedrals.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Dihedrals")?;
    writeln!(writer)?;
    for (idx, dih) in forged.potentials.dihedrals.iter().enumerate() {
        let key = DihedralKey {
            v_barrier: F64Key(dih.v_barrier),
            periodicity: dih.periodicity,
            phase_offset: F64Key(dih.phase_offset),
        };
        let t = mapping
            .dihedral_type_lookup
            .get(&key)
            .ok_or_else(|| Error::Conversion("dihedral type not mapped".into()))?;
        writeln!(
            writer,
            "{:>6} {:>4} {:>6} {:>6} {:>6} {:>6}",
            idx + 1,
            t,
            dih.i + 1,
            dih.j + 1,
            dih.k + 1,
            dih.l + 1
        )?;
    }
    writeln!(writer)?;
    Ok(())
}

fn write_impropers<W: Write>(
    writer: &mut W,
    forged: &ForgedSystem,
    mapping: &TypeMapping,
) -> Result<(), Error> {
    if forged.potentials.impropers.is_empty() {
        return Ok(());
    }
    writeln!(writer, "Impropers")?;
    writeln!(writer)?;
    for (idx, imp) in forged.potentials.impropers.iter().enumerate() {
        let key = match imp {
            ImproperPotential::Planar {
                i,
                j,
                k,
                l,
                k_force,
                chi0,
            } => (
                vec![i, j, k, l],
                ImproperKey::Planar {
                    k_force: F64Key(*k_force),
                    chi0: F64Key(*chi0),
                },
            ),
            ImproperPotential::Umbrella {
                center,
                p1,
                p2,
                p3,
                k_force,
                psi0,
            } => (
                vec![center, p1, p2, p3],
                ImproperKey::Umbrella {
                    k_force: F64Key(*k_force),
                    psi0: F64Key(*psi0),
                },
            ),
        };
        let t = mapping
            .improper_type_lookup
            .get(&key.1)
            .ok_or_else(|| Error::Conversion("improper type not mapped".into()))?;
        let atoms = key.0;
        writeln!(
            writer,
            "{:>6} {:>4} {:>6} {:>6} {:>6} {:>6}",
            idx + 1,
            t,
            atoms[0] + 1,
            atoms[1] + 1,
            atoms[2] + 1,
            atoms[3] + 1
        )?;
    }
    writeln!(writer)?;
    Ok(())
}

fn validate_forged(forged: &ForgedSystem) -> Result<(), Error> {
    if forged.atom_properties.len() != forged.system.atom_count() {
        return Err(Error::Conversion(
            "atom_properties length must match atoms".into(),
        ));
    }
    if forged.system.atoms.is_empty() {
        return Err(Error::Conversion("system has no atoms".into()));
    }
    Ok(())
}

fn build_molecules(bonds: &[Bond], atom_count: usize) -> Vec<usize> {
    let mut parent: Vec<usize> = (0..atom_count).collect();
    let mut size = vec![1usize; atom_count];
    fn find(p: &mut [usize], x: usize) -> usize {
        if p[x] != x {
            p[x] = find(p, p[x]);
        }
        p[x]
    }
    fn unite(p: &mut [usize], sz: &mut [usize], a: usize, b: usize) {
        let mut ra = find(p, a);
        let mut rb = find(p, b);
        if ra == rb {
            return;
        }
        if sz[ra] < sz[rb] {
            std::mem::swap(&mut ra, &mut rb);
        }
        p[rb] = ra;
        sz[ra] += sz[rb];
    }
    for b in bonds {
        unite(&mut parent, &mut size, b.i, b.j);
    }
    let mut root_to_id: HashMap<usize, usize> = HashMap::new();
    let mut next_id = 1usize;
    let mut mol_ids = vec![0usize; atom_count];
    for (i, slot) in mol_ids.iter_mut().enumerate() {
        let r = find(&mut parent, i);
        let id = *root_to_id.entry(r).or_insert_with(|| {
            let v = next_id;
            next_id += 1;
            v
        });
        *slot = id;
    }
    mol_ids
}

struct Bounds {
    xlo: f64,
    xhi: f64,
    ylo: f64,
    yhi: f64,
    zlo: f64,
    zhi: f64,
    tilt: Option<Tilt>,
}

struct Tilt {
    xy: f64,
    xz: f64,
    yz: f64,
}

fn compute_bounds(system: &crate::model::system::System, margin: f64) -> Bounds {
    if let Some(v) = system.box_vectors {
        let a = v[0];
        let b = v[1];
        let c = v[2];
        let xlo = 0.0;
        let xhi = a[0];
        let ylo = 0.0;
        let yhi = b[1];
        let zlo = 0.0;
        let zhi = c[2];
        let tilt = Tilt {
            xy: b[0],
            xz: c[0],
            yz: c[1],
        };
        Bounds {
            xlo,
            xhi,
            ylo,
            yhi,
            zlo,
            zhi,
            tilt: Some(tilt),
        }
    } else {
        let mut xmin = f64::INFINITY;
        let mut ymin = f64::INFINITY;
        let mut zmin = f64::INFINITY;
        let mut xmax = f64::NEG_INFINITY;
        let mut ymax = f64::NEG_INFINITY;
        let mut zmax = f64::NEG_INFINITY;
        for atom in &system.atoms {
            xmin = xmin.min(atom.position[0]);
            ymin = ymin.min(atom.position[1]);
            zmin = zmin.min(atom.position[2]);
            xmax = xmax.max(atom.position[0]);
            ymax = ymax.max(atom.position[1]);
            zmax = zmax.max(atom.position[2]);
        }
        Bounds {
            xlo: xmin - margin,
            xhi: xmax + margin,
            ylo: ymin - margin,
            yhi: ymax + margin,
            zlo: zmin - margin,
            zhi: zmax + margin,
            tilt: None,
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct F64Key(f64);

impl PartialEq for F64Key {
    fn eq(&self, other: &Self) -> bool {
        self.0.to_bits() == other.0.to_bits()
    }
}
impl Eq for F64Key {}
impl std::hash::Hash for F64Key {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.to_bits().hash(state);
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum BondKey {
    Harmonic {
        k_force: F64Key,
        r0: F64Key,
    },
    Morse {
        r0: F64Key,
        d0: F64Key,
        alpha: F64Key,
    },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum AngleKey {
    CosineHarmonic { k_force: F64Key, theta0: F64Key },
    ThetaHarmonic { k_force: F64Key, theta0: F64Key },
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct DihedralKey {
    v_barrier: F64Key,
    periodicity: i32,
    phase_offset: F64Key,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum ImproperKey {
    Planar { k_force: F64Key, chi0: F64Key },
    Umbrella { k_force: F64Key, psi0: F64Key },
}

fn bond_key_from(p: &BondPotential) -> BondKey {
    match p {
        BondPotential::Harmonic { k_force, r0, .. } => BondKey::Harmonic {
            k_force: F64Key(*k_force),
            r0: F64Key(*r0),
        },
        BondPotential::Morse { r0, d0, alpha, .. } => BondKey::Morse {
            r0: F64Key(*r0),
            d0: F64Key(*d0),
            alpha: F64Key(*alpha),
        },
    }
}

fn angle_key_from(p: &AnglePotential) -> AngleKey {
    match p {
        AnglePotential::CosineHarmonic {
            k_force, theta0, ..
        } => AngleKey::CosineHarmonic {
            k_force: F64Key(*k_force),
            theta0: F64Key(*theta0),
        },
        AnglePotential::ThetaHarmonic {
            k_force, theta0, ..
        } => AngleKey::ThetaHarmonic {
            k_force: F64Key(*k_force),
            theta0: F64Key(*theta0),
        },
    }
}

fn improper_key_from(p: &ImproperPotential) -> ImproperKey {
    match p {
        ImproperPotential::Planar { k_force, chi0, .. } => ImproperKey::Planar {
            k_force: F64Key(*k_force),
            chi0: F64Key(*chi0),
        },
        ImproperPotential::Umbrella { k_force, psi0, .. } => ImproperKey::Umbrella {
            k_force: F64Key(*k_force),
            psi0: F64Key(*psi0),
        },
    }
}

struct TypeMapping {
    bond_types: Vec<(usize, BondKey)>,
    angle_types: Vec<(usize, AngleKey)>,
    dihedral_types: Vec<(usize, DihedralKey)>,
    improper_types: Vec<(usize, ImproperKey)>,
    bond_type_lookup: HashMap<BondKey, usize>,
    angle_type_lookup: HashMap<AngleKey, usize>,
    dihedral_type_lookup: HashMap<DihedralKey, usize>,
    improper_type_lookup: HashMap<ImproperKey, usize>,
}

impl TypeMapping {
    fn new(forged: &ForgedSystem) -> Result<Self, Error> {
        let mut bond_type_lookup: HashMap<BondKey, usize> = HashMap::new();
        let mut angle_type_lookup: HashMap<AngleKey, usize> = HashMap::new();
        let mut dihedral_type_lookup: HashMap<DihedralKey, usize> = HashMap::new();
        let mut improper_type_lookup: HashMap<ImproperKey, usize> = HashMap::new();

        let mut bond_types = Vec::new();
        for p in &forged.potentials.bonds {
            let key = bond_key_from(p);
            let next = bond_type_lookup.len() + 1;
            let id = *bond_type_lookup.entry(key).or_insert(next);
            if id == next {
                bond_types.push((id, key));
            }
        }

        let mut angle_types = Vec::new();
        for p in &forged.potentials.angles {
            let key = angle_key_from(p);
            let next = angle_type_lookup.len() + 1;
            let id = *angle_type_lookup.entry(key).or_insert(next);
            if id == next {
                angle_types.push((id, key));
            }
        }

        let mut dihedral_types = Vec::new();
        for p in &forged.potentials.dihedrals {
            let key = DihedralKey {
                v_barrier: F64Key(p.v_barrier),
                periodicity: p.periodicity,
                phase_offset: F64Key(p.phase_offset),
            };
            let next = dihedral_type_lookup.len() + 1;
            let id = *dihedral_type_lookup.entry(key).or_insert(next);
            if id == next {
                dihedral_types.push((id, key));
            }
        }

        let mut improper_types = Vec::new();
        for p in &forged.potentials.impropers {
            let key = improper_key_from(p);
            let next = improper_type_lookup.len() + 1;
            let id = *improper_type_lookup.entry(key).or_insert(next);
            if id == next {
                improper_types.push((id, key));
            }
        }

        bond_types.sort_by_key(|(id, _)| *id);
        angle_types.sort_by_key(|(id, _)| *id);
        dihedral_types.sort_by_key(|(id, _)| *id);
        improper_types.sort_by_key(|(id, _)| *id);

        Ok(Self {
            bond_types,
            angle_types,
            dihedral_types,
            improper_types,
            bond_type_lookup,
            angle_type_lookup,
            dihedral_type_lookup,
            improper_type_lookup,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::{
        atom::Atom,
        topology::{
            AnglePotential, AtomParam, BondPotential, DihedralPotential, HBondPotential,
            ImproperPotential, Potentials,
        },
        types::{BondOrder, Element},
    };

    fn complex_forged() -> ForgedSystem {
        let atoms = vec![
            Atom::new(Element::C, [0.0, 0.0, 0.0]),
            Atom::new(Element::O, [1.2, 0.0, 0.1]),
            Atom::new(Element::H, [2.1, -0.2, 0.0]),
            Atom::new(Element::N, [0.5, 1.0, -0.3]),
            Atom::new(Element::H, [4.5, 4.5, 4.5]),
        ];

        let bonds = vec![
            Bond::new(0, 1, BondOrder::Single),
            Bond::new(1, 2, BondOrder::Single),
            Bond::new(0, 3, BondOrder::Single),
        ];

        let atom_properties = vec![
            AtomParam {
                charge: -0.1,
                mass: 12.0,
                type_idx: 0,
            },
            AtomParam {
                charge: -0.2,
                mass: 16.0,
                type_idx: 1,
            },
            AtomParam {
                charge: 0.3,
                mass: 1.0,
                type_idx: 2,
            },
            AtomParam {
                charge: -0.05,
                mass: 14.0,
                type_idx: 3,
            },
            AtomParam {
                charge: 0.1,
                mass: 1.0,
                type_idx: 2,
            },
        ];

        let potentials = Potentials {
            bonds: vec![
                BondPotential::Harmonic {
                    i: 0,
                    j: 1,
                    k_force: 700.0,
                    r0: 1.23,
                },
                BondPotential::Morse {
                    i: 1,
                    j: 2,
                    r0: 1.5,
                    d0: 4.0,
                    alpha: 2.0,
                },
                BondPotential::Harmonic {
                    i: 0,
                    j: 3,
                    k_force: 700.0,
                    r0: 1.23,
                },
            ],
            angles: vec![
                AnglePotential::CosineHarmonic {
                    i: 0,
                    j: 1,
                    k: 2,
                    k_force: 50.0,
                    theta0: 109.5,
                },
                AnglePotential::ThetaHarmonic {
                    i: 3,
                    j: 0,
                    k: 1,
                    k_force: 40.0,
                    theta0: 120.0,
                },
            ],
            dihedrals: vec![DihedralPotential {
                i: 2,
                j: 1,
                k: 0,
                l: 3,
                v_barrier: 2.5,
                periodicity: 3,
                phase_offset: 180.0,
            }],
            impropers: vec![
                ImproperPotential::Planar {
                    i: 0,
                    j: 1,
                    k: 2,
                    l: 3,
                    k_force: 10.0,
                    chi0: 0.0,
                },
                ImproperPotential::Umbrella {
                    center: 1,
                    p1: 0,
                    p2: 2,
                    p3: 3,
                    k_force: 5.0,
                    psi0: 180.0,
                },
            ],
            vdw_pairs: vec![
                VdwPairPotential::LennardJones {
                    type1_idx: 0,
                    type2_idx: 1,
                    sigma: 3.5,
                    epsilon: 0.2,
                },
                VdwPairPotential::Exponential6 {
                    type1_idx: 1,
                    type2_idx: 3,
                    a: 1000.0,
                    b: 50.0,
                    c: 2.0,
                },
            ],
            h_bonds: vec![HBondPotential {
                donor_type_idx: 1,
                hydrogen_type_idx: 2,
                acceptor_type_idx: 3,
                d0: 9.5,
                r0: 2.75,
            }],
        };

        ForgedSystem {
            system: crate::model::system::System {
                atoms,
                bonds,
                box_vectors: Some([[5.0, 0.0, 0.0], [1.0, 6.0, 0.0], [0.5, 1.5, 7.0]]),
                bio_metadata: None,
            },
            atom_types: vec!["C_3".into(), "O_2".into(), "H_".into(), "N_3".into()],
            atom_properties,
            potentials,
        }
    }

    fn write_outputs(forged: &ForgedSystem, cfg: &LammpsConfig) -> (String, String) {
        let mut data_buf = Vec::new();
        let mut settings_buf = Vec::new();
        write_data_file(&mut data_buf, forged, cfg).expect("write data");
        write_settings_file(&mut settings_buf, forged, cfg).expect("write settings");
        (
            String::from_utf8(data_buf).unwrap(),
            String::from_utf8(settings_buf).unwrap(),
        )
    }

    #[test]
    fn data_file_covers_counts_coeffs_atoms_and_tilt() {
        let forged = complex_forged();
        let cfg = LammpsConfig::default();
        let (data_out, _) = write_outputs(&forged, &cfg);

        for expected in [
            "5 atoms",
            "3 bonds",
            "2 angles",
            "1 dihedrals",
            "2 impropers",
            "4 atom types",
            "2 bond types",
            "2 angle types",
            "1 dihedral types",
            "2 improper types",
        ] {
            assert!(data_out.contains(expected), "missing line: {expected}");
        }

        assert!(data_out.contains("0.000000 5.000000 xlo xhi"));
        assert!(data_out.contains("0.000000 6.000000 ylo yhi"));
        assert!(data_out.contains("0.000000 7.000000 zlo zhi"));
        assert!(data_out.contains("1.000000 0.500000 1.500000 xy xz yz"));

        assert!(data_out.contains("Bond Coeffs"));
        assert!(data_out.contains("1 harmonic 700.000000 1.230000"));
        assert!(data_out.contains("2 morse 4.000000 2.000000 1.500000"));
        assert!(data_out.contains("Angle Coeffs"));
        assert!(data_out.contains("1 cosine/squared 50.000000 109.500000"));
        assert!(data_out.contains("2 harmonic 40.000000 120.000000"));
        assert!(data_out.contains("Dihedral Coeffs"));
        assert!(data_out.contains("1 charmm 2.500000 3 180.000000 0.0"));
        assert!(data_out.contains("Improper Coeffs"));
        assert!(data_out.contains("1 umbrella 10.000000 0.000000"));
        assert!(data_out.contains("2 umbrella 5.000000 180.000000"));

        let lines: Vec<&str> = data_out.lines().collect();
        let atoms_start = lines
            .iter()
            .position(|l| l.starts_with("Atoms #"))
            .expect("Atoms section header present");
        let atom_lines: Vec<&str> = lines
            .iter()
            .skip(atoms_start + 1)
            .skip_while(|l| l.trim().is_empty())
            .take_while(|l| !l.trim().is_empty())
            .copied()
            .collect();
        let atom5 = atom_lines
            .iter()
            .find(|l| l.trim_start().starts_with("5 "))
            .copied()
            .expect("atom 5 line present");
        let parts: Vec<&str> = atom5.split_whitespace().collect();
        assert_eq!(parts[0], "5");
        assert_eq!(parts[1], "2", "atom 5 molecule id");
        assert_eq!(parts[2], "3", "atom 5 type id");

        assert!(data_out.contains("1    1      1      2"));
        assert!(data_out.contains("2    2      2      3"));
        assert!(data_out.contains("3    1      1      4"));
        assert!(data_out.contains("Angles"));
        assert!(data_out.contains("1    1      1      2      3"));
        assert!(data_out.contains("2    2      4      1      2"));
    }

    #[test]
    fn settings_file_covers_styles_pair_and_hbond() {
        let forged = complex_forged();
        let cfg = LammpsConfig::default();
        let (_, settings_out) = write_outputs(&forged, &cfg);

        assert!(settings_out.contains("bond_style     hybrid harmonic morse"));
        assert!(settings_out.contains("angle_style    hybrid cosine/squared harmonic"));
        assert!(settings_out.contains("dihedral_style charmm"));
        assert!(settings_out.contains("improper_style umbrella"));

        assert!(settings_out.contains("pair_style      hybrid/overlay lj/cut/coul/long 12.000 buck/coul/long 12.000 hbond/dreiding/lj 2 10.000 12.000 90.000"));
        assert!(settings_out.contains("kspace_style    pppm 1.0e-4"));

        assert!(settings_out.contains("pair_coeff   1   2 lj/cut/coul/long 0.200000 3.500000"));
        assert!(
            settings_out
                .contains("pair_coeff   2   4 buck/coul/long 1000.000000 50.000000 2.000000")
        );
        assert!(
            settings_out.contains("pair_coeff   2   4 hbond/dreiding/lj 3 i 9.500000 2.627522 2")
        );
    }

    #[test]
    fn nonperiodic_uses_cutoff_only_and_no_kspace() {
        let forged = complex_forged();
        let cfg = LammpsConfig {
            system_type: SystemType::NonPeriodic,
            ..Default::default()
        };
        let (_, settings_out) = write_outputs(&forged, &cfg);

        assert!(settings_out.contains("pair_style      hybrid/overlay lj/cut/coul/cut 12.000 buck/coul/cut 12.000 hbond/dreiding/lj 2 10.000 12.000 90.000"));
        assert!(!settings_out.contains("kspace_style"));
        assert!(settings_out.contains("pair_coeff   1   2 lj/cut/coul/cut 0.200000 3.500000"));
        assert!(
            settings_out
                .contains("pair_coeff   2   4 buck/coul/cut 1000.000000 50.000000 2.000000")
        );
        assert!(
            settings_out.contains("pair_coeff   2   4 hbond/dreiding/lj 3 i 9.500000 2.627522 2")
        );
    }
}
