uff-relax 1.0.7

use uff_relax::{System, Atom, Bond, UnitCell, UffOptimizer};
use glam::DVec3;
use std::time::Instant;

fn create_polyethylene(n_carbons: usize, with_charges: bool) -> System {
    let mut atoms = Vec::new();
    let mut bonds = Vec::new();
    let bond_len = 1.53;
    let angle = 109.5f64.to_radians();
    
    for i in 0..n_carbons {
        let x = i as f64 * bond_len * (angle/2.0).sin();
        let y = if i % 2 == 0 { 0.0 } else { bond_len * (angle/2.0).cos() };
        let pos = DVec3::new(x, y, 0.0);
        let mut atom = Atom::new(6, pos);
        if with_charges { atom.charge = -0.06; }
        atoms.push(atom);
        if i > 0 { bonds.push(Bond { atom_indices: (i - 1, i), order: 1.0 }); }
    }
    for i in 0..n_carbons {
        let c_pos = atoms[i].position;
        let h_offsets = [DVec3::new(0.0, 0.0, 1.0), DVec3::new(0.0, 0.0, -1.0)];
        for &offset in &h_offsets {
            let mut h_atom = Atom::new(1, c_pos + offset);
            if with_charges { h_atom.charge = 0.03; }
            atoms.push(h_atom);
            bonds.push(Bond { atom_indices: (i, atoms.len() - 1), order: 1.0 });
        }
    }
    System::new(atoms, bonds, UnitCell::new_none())
}

fn run_benchmark(cutoff: f64, threshold: f64, with_charges: bool) {
    let mut system = create_polyethylene(50, true);
    if !with_charges {
        for atom in &mut system.atoms { atom.charge = 0.0; }
    }
    let optimizer = UffOptimizer::new(3000, threshold)
        .with_cutoff(cutoff)
        .with_verbose(false);

    let steps = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
    let steps_clone = steps.clone();
    let optimizer = optimizer.with_step_hook(move |iter, _, _| {
        steps_clone.store(iter + 1, std::sync::atomic::Ordering::SeqCst);
    });

    let start = Instant::now();
    optimizer.optimize(&mut system);
    let duration = start.elapsed();

    println!("Charges: {:>5} | Cutoff: {:>4.1} Å | Steps: {:>4} | Time: {:>8.3?}", 
        with_charges, cutoff, steps.load(std::sync::atomic::Ordering::SeqCst), duration);
}

fn main() {
    let threshold = 0.1;
    println!("Benchmarking Polyethylene (50C) (Threshold: {} kcal/mol/Å)", threshold);
    println!("------------------------------------------------------------------");
    run_benchmark(6.0, threshold, true);
    run_benchmark(6.0, threshold, false);
    run_benchmark(12.0, threshold, true);
    run_benchmark(12.0, threshold, false);
    println!("------------------------------------------------------------------");
}