use std::path::Path;
use std::time::Instant;

use rand::{rngs::StdRng, seq::SliceRandom, SeedableRng, Rng};

use crate::config::PackConfig;
use crate::constraints::satisfies_constraints;
use crate::error::{PackError, PackResult};
use crate::geom::{Quaternion, Vec3};
use crate::io::write_output;
use crate::movebad::run_movebad_pass;
use crate::atom_params::{AtomParams, build_atom_params, max_radius, scale_atom_params};
use crate::pack_ops::{
    append_bonds, append_ter_after, apply_chain, apply_resnumbers, build_order, load_template,
    maybe_write_snapshot, place_fixed, random_center, transform_atoms, validate_min_distance,
    write_bad_snapshot, PlacementSeed, TemplateEntry,
};
use crate::pbc::PbcBox;
use crate::placement::PlacementRecord;
use crate::relax::relax_overlaps;
use crate::gencan::{evaluate_state, optimize_gencan};
use crate::restart::{read_restart, write_restart, RestartEntry};
use crate::spatial_hash::{SpatialHash, SpatialHashParamsExt};

#[derive(Clone, Debug)]
pub struct AtomRecord {
    pub record_kind: AtomRecordKind,
    pub name: String,
    pub element: String,
    pub resname: String,
    pub resid: i32,
    pub chain: char,
    pub segid: String,
    pub charge: f32,
    pub position: Vec3,
    pub mol_id: i32,
}

#[derive(Clone, Debug)]
pub struct PackOutput {
    pub atoms: Vec<AtomRecord>,
    pub bonds: Vec<(usize, usize)>,
    pub box_size: [f32; 3],
    pub ter_after: Vec<usize>,
}

fn rebuild_hash_from_positions(hash: &mut SpatialHash, positions: &[Vec3], cell_size: f32) {
    *hash = SpatialHash::new(cell_size);
    for (idx, pos) in positions.iter().copied().enumerate() {
        hash.insert(idx, pos);
    }
}

fn maybe_write_gencan_snapshot(
    cfg: &PackConfig,
    atoms: &[AtomRecord],
    bonds: &[(usize, usize)],
    ter_after: &[usize],
    box_size: [f32; 3],
    last_write: &mut Instant,
    current_fx: f32,
    best_fx: &mut f32,
) -> PackResult<()> {
    let Some(interval) = cfg.writeout else {
        return Ok(());
    };
    if interval <= 0.0 {
        return Ok(());
    }
    if last_write.elapsed().as_secs_f32() < interval {
        return Ok(());
    }
    let Some(spec) = cfg.output.as_ref() else {
        return Ok(());
    };
    let improved = current_fx < *best_fx;
    if improved || cfg.writebad {
        let out = PackOutput {
            atoms: atoms.to_vec(),
            bonds: bonds.to_vec(),
            box_size,
            ter_after: ter_after.to_vec(),
        };
        let add_box_sides = cfg.add_box_sides || cfg.pbc;
        let box_fix = cfg.add_box_sides_fix.unwrap_or(0.0);
        let write_conect = !cfg.ignore_conect;
        write_output(
            &out,
            spec,
            add_box_sides,
            box_fix,
            write_conect,
            cfg.hexadecimal_indices,
        )?;
        if improved {
            *best_fx = current_fx;
        }
    }
    *last_write = Instant::now();
    Ok(())
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum AtomRecordKind {
    Atom,
    HetAtom,
}

pub fn run(config: &PackConfig) -> PackResult<PackOutput> {
    let cfg = config.normalized()?;
    let seed = cfg.seed.unwrap_or(0);
    let out = run_once(&cfg, seed, 0)?;
    if cfg.check {
        validate_min_distance(&out.atoms, cfg.min_distance.unwrap_or(2.0))?;
    }
    Ok(out)
}

fn run_once(cfg: &PackConfig, seed: u64, attempt: usize) -> PackResult<PackOutput> {
    let mut rng = StdRng::seed_from_u64(seed);
    let box_size = cfg.box_.size;
    let box_origin = Vec3::from_array(cfg.pbc_min.unwrap_or([0.0, 0.0, 0.0]));
    let pbc = if cfg.pbc {
        if let (Some(min), Some(max)) = (cfg.pbc_min, cfg.pbc_max) {
            Some(PbcBox::from_bounds(min, max)?)
        } else {
            Some(PbcBox::from_size(box_size)?)
        }
    } else {
        None
    };
    let dist_scale = if attempt == 0 {
        cfg.discale.unwrap_or(1.0)
    } else {
        1.0
    };
    let total_mols: usize = cfg.structures.iter().map(|s| s.count).sum();
    let restart_all = if let Some(path) = &cfg.restart_from {
        Some(read_restart(Path::new(path))?)
    } else {
        None
    };
    let mut restart_per_spec = Vec::with_capacity(cfg.structures.len());
    for spec in &cfg.structures {
        if let Some(path) = &spec.restart_from {
            restart_per_spec.push(Some(read_restart(Path::new(path))?));
        } else {
            restart_per_spec.push(None);
        }
    }
    if let Some(entries) = &restart_all {
        if entries.len() != total_mols {
            return Err(PackError::Invalid(
                "restart_from entry count does not match total molecule count".into(),
            ));
        }
    }
    for (idx, spec) in cfg.structures.iter().enumerate() {
        if let Some(entries) = &restart_per_spec[idx] {
            if entries.len() != spec.count {
                return Err(PackError::Invalid(format!(
                    "restart_from entry count does not match count for structure {}",
                    spec.path
                )));
            }
        }
    }
    let mut restart_offsets = vec![0usize; cfg.structures.len()];
    if restart_all.is_some() {
        let mut offset = 0usize;
        for (idx, spec) in cfg.structures.iter().enumerate() {
            restart_offsets[idx] = offset;
            offset += spec.count;
        }
    }
    let restart_entry = |spec_index: usize, local_index: usize| -> Option<RestartEntry> {
        if let Some(entries) = &restart_per_spec[spec_index] {
            return entries.get(local_index).copied();
        }
        if let Some(entries) = &restart_all {
            let offset = restart_offsets[spec_index];
            return entries.get(offset + local_index).copied();
        }
        None
    };
    let mut global_resid_counter = 1i32;
    let mut atoms = Vec::new();
    let mut bonds = Vec::new();
    let mut ter_after = Vec::new();
    let mut positions = Vec::new();
    let mut atom_params: Vec<AtomParams> = Vec::new();
    let fbins = cfg.fbins.unwrap_or(3.0_f32.sqrt());
    let mut mol_id = 1i32;
    let mut last_write = Instant::now();
    let mut mol_spec: Vec<usize> = Vec::with_capacity(total_mols);
    let mut placements = vec![PlacementRecord::default(); total_mols];
    let mut placement_filled = vec![false; total_mols];
    let short_tol_dist = cfg
        .short_tol_dist
        .unwrap_or(cfg.min_distance.unwrap_or(2.0) * 0.5)
        * dist_scale;
    let short_tol_scale = cfg.short_tol_scale.unwrap_or(3.0);
    let mut use_short_tol = cfg.use_short_tol;

    let mut templates = Vec::with_capacity(cfg.structures.len());
    let mut max_atom_radius = 0.0f32;
    for (idx, spec) in cfg.structures.iter().enumerate() {
        let format = spec.format.as_deref().or(cfg.filetype.as_deref());
        let template = load_template(
            spec,
            mol_id + idx as i32,
            format,
            cfg.ignore_conect,
            cfg.non_standard_conect,
        )?;
        let atom_params = build_atom_params(cfg, spec, template.atoms.len(), dist_scale)?;
        max_atom_radius = max_atom_radius.max(max_radius(&atom_params));
        let res_mode = spec.resnumbers.unwrap_or(1);
        let connect = spec.connect && !cfg.ignore_conect;
        let maxmove = spec
            .maxmove
            .or(cfg.maxmove)
            .unwrap_or(spec.count.max(1));
        templates.push(TemplateEntry {
            template,
            res_mode,
            connect,
            maxmove,
            atom_params,
        });
    }
    if !use_short_tol {
        use_short_tol = templates
            .iter()
            .any(|entry| entry.atom_params.iter().any(|p| p.use_short));
    }
    let short_tol = if use_short_tol {
        Some((short_tol_dist, short_tol_scale))
    } else {
        None
    };
    let cell_size = ((2.0 * max_atom_radius) * fbins * dist_scale).max(1.0e-6);
    let mut hash = SpatialHash::new(cell_size);

    for (idx, spec) in cfg.structures.iter().enumerate() {
        if !spec.fixed {
            continue;
        }
        let entry = &templates[idx];
        let mut seeds = Vec::with_capacity(spec.count);
        let mut seed_eulers = Vec::with_capacity(spec.count);
        for i in 0..spec.count {
            let (center, rotation, euler) = if let Some(positions) = &spec.positions {
                let pos = positions
                    .get(i)
                    .copied()
                    .unwrap_or([0.0, 0.0, 0.0]);
                let euler = spec
                    .fixed_eulers
                    .as_ref()
                    .and_then(|vals| vals.get(i).copied())
                    .unwrap_or([0.0, 0.0, 0.0]);
                let rotation = Quaternion::from_packmol_euler(euler[0], euler[1], euler[2]);
                (Vec3::from_array(pos), rotation, euler)
            } else if let Some(entry) = restart_entry(idx, i) {
                let rotation = Quaternion::from_packmol_euler(
                    entry.euler[0],
                    entry.euler[1],
                    entry.euler[2],
                );
                (entry.center, rotation, entry.euler)
            } else {
                return Err(PackError::Invalid(
                    "fixed structure missing positions and restart_from".into(),
                ));
            };
            seeds.push(PlacementSeed { center, rotation });
            seed_eulers.push(euler);
        }
        place_fixed(
            spec,
            &entry.template,
            &entry.atom_params,
            &seeds,
            &mut atoms,
            &mut positions,
            &mut atom_params,
            &mut hash,
            mol_id,
            pbc,
            entry.res_mode,
            &mut global_resid_counter,
            &mut bonds,
            &mut ter_after,
            entry.connect,
            cfg.add_amber_ter,
            cfg.amber_ter_preserve,
            cfg.avoid_overlap,
        )?;
        let mol_id_start = mol_id;
        for i in 0..spec.count {
            let mol_index = (mol_id_start + i as i32 - 1) as usize;
            placements[mol_index] = PlacementRecord::new(seeds[i].center, seed_eulers[i]);
            placement_filled[mol_index] = true;
            mol_spec.push(idx);
        }
        mol_id += spec.count as i32;
        maybe_write_snapshot(
            &cfg,
            &atoms,
            &bonds,
            &ter_after,
            box_size,
            &mut last_write,
        )?;
    }

    let mut order = build_order(cfg);
    if cfg.packall {
        order.shuffle(&mut rng);
    }

    if cfg.randominitialpoint {
        order.shuffle(&mut rng);
    }

    struct Candidate {
        atoms: Vec<AtomRecord>,
        penalty: f32,
        placement: PlacementRecord,
    }

    for placement in order {
        let spec = &cfg.structures[placement.spec_index];
        let entry = &templates[placement.spec_index];
        let mut best: Option<Candidate> = None;
        if let Some(restart) = restart_entry(
            placement.spec_index,
            (placement.local_index.max(1) - 1) as usize,
        ) {
            let rotation = Quaternion::from_packmol_euler(
                restart.euler[0],
                restart.euler[1],
                restart.euler[2],
            );
            let candidate = transform_atoms(&entry.template, rotation, restart.center, mol_id);
            let cand_positions: Vec<Vec3> = candidate.iter().map(|a| a.position).collect();
            if satisfies_constraints(&cand_positions, &spec.constraints, pbc) {
                if short_tol.is_some() {
                    if let Some(penalty) = hash.overlaps_short_tol_params(
                        &cand_positions,
                        &entry.atom_params,
                        &positions,
                        &atom_params,
                    ) {
                        best = Some(Candidate {
                            atoms: candidate,
                            penalty,
                            placement: PlacementRecord::new(restart.center, restart.euler),
                        });
                    }
                } else if !hash.overlaps_params(
                    &cand_positions,
                    &entry.atom_params,
                    &positions,
                    &atom_params,
                ) {
                    best = Some(Candidate {
                        atoms: candidate,
                        penalty: 0.0,
                        placement: PlacementRecord::new(restart.center, restart.euler),
                    });
                }
            }
        }
        if !best.as_ref().map(|cand| cand.penalty == 0.0).unwrap_or(false) {
            for _ in 0..cfg.max_attempts.unwrap_or(10000) {
                let (rotation, euler) = if spec.rotate {
                    if let Some(bounds) = spec.rot_bounds {
                        let bx = rng.gen_range(bounds[0][0]..=bounds[0][1]);
                        let gy = rng.gen_range(bounds[1][0]..=bounds[1][1]);
                        let tz = rng.gen_range(bounds[2][0]..=bounds[2][1]);
                        (
                            Quaternion::from_packmol_euler(bx, gy, tz),
                            [bx, gy, tz],
                        )
                    } else {
                        let rot = Quaternion::random(&mut rng);
                        let (b, g, t) = rot.to_packmol_euler();
                        (rot, [b, g, t])
                    }
                } else {
                    (Quaternion::identity(), [0.0, 0.0, 0.0])
                };
                let mut center = random_center(&mut rng, box_size, entry.template.radius)?;
                center = center.add(box_origin);
                let candidate = transform_atoms(&entry.template, rotation, center, mol_id);
                let cand_positions: Vec<Vec3> = candidate.iter().map(|a| a.position).collect();
                if !satisfies_constraints(&cand_positions, &spec.constraints, pbc) {
                    continue;
                }
                if short_tol.is_some() {
                    if let Some(penalty) = hash.overlaps_short_tol_params(
                        &cand_positions,
                        &entry.atom_params,
                        &positions,
                        &atom_params,
                    ) {
                        if best
                            .as_ref()
                            .map(|best| penalty < best.penalty)
                            .unwrap_or(true)
                        {
                            best = Some(Candidate {
                                atoms: candidate,
                                penalty,
                                placement: PlacementRecord::new(center, euler),
                            });
                        }
                        if penalty == 0.0 {
                            break;
                        }
                    }
                } else if !hash.overlaps_params(
                    &cand_positions,
                    &entry.atom_params,
                    &positions,
                    &atom_params,
                ) {
                    best = Some(Candidate {
                        atoms: candidate,
                        penalty: 0.0,
                        placement: PlacementRecord::new(center, euler),
                    });
                    break;
                }
            }
        }
        let Some(best) = best else {
            if cfg.writebad {
                write_bad_snapshot(cfg, &atoms, &bonds, &ter_after, box_size)?;
            }
            return Err(PackError::Placement(format!(
                "failed to place structure {}",
                spec.path
            )));
        };
        let mut candidate = best.atoms;
        apply_chain(&mut candidate, spec, mol_id);
        apply_resnumbers(
            &mut candidate,
            entry.res_mode,
            entry.template.resid_min,
            global_resid_counter,
            placement.local_index,
            mol_id,
        )?;
        global_resid_counter += entry.template.resid_count;
        let base = atoms.len();
        if entry.connect {
            append_bonds(&mut bonds, &entry.template.bonds, base);
        }
        if cfg.amber_ter_preserve {
            append_ter_after(&mut ter_after, &entry.template.ter_after, base);
        }
        for (offset, atom) in candidate.into_iter().enumerate() {
            let idx = positions.len();
            positions.push(atom.position);
            hash.insert(idx, atom.position);
            atom_params.push(entry.atom_params[offset]);
            atoms.push(atom);
        }
        if cfg.add_amber_ter {
            if let Some(last) = atoms.len().checked_sub(1) {
                ter_after.push(last);
            }
        }
        let mol_index = (mol_id - 1) as usize;
        placements[mol_index] = best.placement;
        placement_filled[mol_index] = true;
        mol_spec.push(placement.spec_index);
        mol_id += 1;
        maybe_write_snapshot(cfg, &atoms, &bonds, &ter_after, box_size, &mut last_write)?;

        if !cfg.disable_movebad {
            run_movebad_pass(
                cfg,
                &templates,
                &mut atoms,
                &mut positions,
                &mut atom_params,
                &mut hash,
                &mut bonds,
                &mut ter_after,
                &mut placements,
                box_size,
                box_origin,
                pbc,
                &mut rng,
                &mut global_resid_counter,
                &mut mol_id,
                &mol_spec,
                cell_size,
            )?;
        }
    }

    if placement_filled.iter().any(|filled| !filled) {
        return Err(PackError::Placement(
            "internal error: missing placement records".into(),
        ));
    }
    if mol_spec.len() != placements.len() {
        return Err(PackError::Placement(
            "internal error: molecule mapping mismatch".into(),
        ));
    }

    if cfg.use_gencan {
        let precision = cfg.precision.unwrap_or(1.0e-2);
        let nloop0_default = cfg.nloop0.unwrap_or(0);
        let nloop_default = cfg.nloop.unwrap_or(1);
        let use_short_tol = cfg.use_short_tol;
        let radscale_start = if dist_scale > 1.0 {
            1.0
        } else {
            cfg.discale.unwrap_or(1.0).max(1.0)
        };
        let mut gencan_last_write = Instant::now();
        let mut fprint = evaluate_state(
            cfg,
            &atoms,
            &positions,
            &atom_params,
            &mol_spec,
            &templates,
            pbc,
            use_short_tol,
            true,
            None,
            None,
        )
        .map(|res| res.value)
        .unwrap_or(f32::INFINITY);

        for (spec_index, spec) in cfg.structures.iter().enumerate() {
            let loops0 = spec.nloop0.unwrap_or(nloop0_default);
            let mut radscale = radscale_start;
            let mut fimp = f32::INFINITY;
            let mut base_stats = evaluate_state(
                cfg,
                &atoms,
                &positions,
                &atom_params,
                &mol_spec,
                &templates,
                pbc,
                use_short_tol,
                true,
                Some(spec_index),
                None,
            )?;
            let mut bestf = base_stats.value;
            let mut flast = base_stats.value;

            for _ in 0..loops0 {
                let scaled = scale_atom_params(&atom_params, radscale);
                let res = optimize_gencan(
                    cfg,
                    &mut atoms,
                    &mut positions,
                    &scaled,
                    &mut placements,
                    &mol_spec,
                    &templates,
                    box_size,
                    box_origin.to_array(),
                    pbc,
                    false,
                    Some(spec_index),
                    None,
                )?;
                rebuild_hash_from_positions(&mut hash, &positions, cell_size);
                base_stats = evaluate_state(
                    cfg,
                    &atoms,
                    &positions,
                    &atom_params,
                    &mol_spec,
                    &templates,
                    pbc,
                    use_short_tol,
                    true,
                    Some(spec_index),
                    None,
                )?;
                flast = base_stats.value;
                if base_stats.max_constraint <= precision {
                    break;
                }
                let _ = res;
            }

            let loops = spec.nloop.unwrap_or(nloop_default);
            for _ in 0..loops {
                if !cfg.disable_movebad && radscale <= 1.0 + 1.0e-6 && fimp <= 10.0 {
                    rebuild_hash_from_positions(&mut hash, &positions, cell_size);
                    run_movebad_pass(
                        cfg,
                        &templates,
                        &mut atoms,
                        &mut positions,
                        &mut atom_params,
                        &mut hash,
                        &mut bonds,
                        &mut ter_after,
                        &mut placements,
                        box_size,
                        box_origin,
                        pbc,
                        &mut rng,
                        &mut global_resid_counter,
                        &mut mol_id,
                        &mol_spec,
                        cell_size,
                    )?;
                }
                let scaled = scale_atom_params(&atom_params, radscale);
                let res = optimize_gencan(
                    cfg,
                    &mut atoms,
                    &mut positions,
                    &scaled,
                    &mut placements,
                    &mol_spec,
                    &templates,
                    box_size,
                    box_origin.to_array(),
                    pbc,
                    true,
                    Some(spec_index),
                    None,
                )?;
                base_stats = evaluate_state(
                    cfg,
                    &atoms,
                    &positions,
                    &atom_params,
                    &mol_spec,
                    &templates,
                    pbc,
                    use_short_tol,
                    true,
                    Some(spec_index),
                    None,
                )?;
                if flast > 0.0 {
                    let last_imp = -100.0 * (base_stats.value - flast) / flast;
                    fimp = last_imp.max(-99.99).min(99.99);
                } else {
                    fimp = 100.0;
                }
                flast = base_stats.value;
                if base_stats.value < bestf {
                    bestf = base_stats.value;
                }
                let all_stats = evaluate_state(
                    cfg,
                    &atoms,
                    &positions,
                    &atom_params,
                    &mol_spec,
                    &templates,
                    pbc,
                    use_short_tol,
                    true,
                    None,
                    None,
                )?;
                maybe_write_gencan_snapshot(
                    cfg,
                    &atoms,
                    &bonds,
                    &ter_after,
                    box_size,
                    &mut gencan_last_write,
                    all_stats.value,
                    &mut fprint,
                )?;
                if base_stats.max_overlap <= precision && base_stats.max_constraint <= precision {
                    break;
                }
                if radscale > 1.0 {
                    if (base_stats.max_overlap <= precision && fimp < 10.0) || fimp < 2.0 {
                        radscale = (0.9 * radscale).max(1.0);
                    }
                }
                let _ = res;
            }
        }

        let mut radscale = radscale_start;
        let mut fimp = f32::INFINITY;
        let mut all_stats = evaluate_state(
            cfg,
            &atoms,
            &positions,
            &atom_params,
            &mol_spec,
            &templates,
            pbc,
            use_short_tol,
            true,
            None,
            None,
        )?;
        let mut bestf = all_stats.value;
        let mut flast = all_stats.value;

        for _ in 0..nloop_default {
            if !cfg.disable_movebad && radscale <= 1.0 + 1.0e-6 && fimp <= 10.0 {
                rebuild_hash_from_positions(&mut hash, &positions, cell_size);
                run_movebad_pass(
                    cfg,
                    &templates,
                    &mut atoms,
                    &mut positions,
                    &mut atom_params,
                    &mut hash,
                    &mut bonds,
                    &mut ter_after,
                    &mut placements,
                    box_size,
                    box_origin,
                    pbc,
                    &mut rng,
                    &mut global_resid_counter,
                    &mut mol_id,
                    &mol_spec,
                    cell_size,
                )?;
            }
            let scaled = scale_atom_params(&atom_params, radscale);
            let res = optimize_gencan(
                cfg,
                &mut atoms,
                &mut positions,
                &scaled,
                &mut placements,
                &mol_spec,
                &templates,
                box_size,
                box_origin.to_array(),
                pbc,
                true,
                None,
                None,
            )?;
            all_stats = evaluate_state(
                cfg,
                &atoms,
                &positions,
                &atom_params,
                &mol_spec,
                &templates,
                pbc,
                use_short_tol,
                true,
                None,
                None,
            )?;
            if flast > 0.0 {
                let last_imp = -100.0 * (all_stats.value - flast) / flast;
                fimp = last_imp.max(-99.99).min(99.99);
            } else {
                fimp = 100.0;
            }
            flast = all_stats.value;
            if all_stats.value < bestf {
                bestf = all_stats.value;
            }
            maybe_write_gencan_snapshot(
                cfg,
                &atoms,
                &bonds,
                &ter_after,
                box_size,
                &mut gencan_last_write,
                all_stats.value,
                &mut fprint,
            )?;
            if all_stats.max_overlap <= precision && all_stats.max_constraint <= precision {
                break;
            }
            if radscale > 1.0 {
                if (all_stats.max_overlap <= precision && fimp < 10.0) || fimp < 2.0 {
                    radscale = (0.9 * radscale).max(1.0);
                }
            }
            let _ = res;
        }
    }

    if cfg.relax_steps.unwrap_or(0) > 0 {
        relax_overlaps(
            cfg,
            &mut atoms,
            &mut positions,
            &mut atom_params,
            &mut placements,
            &mol_spec,
            &templates,
            box_size,
            pbc,
            dist_scale,
        )?;
    }

    if cfg.restart_to.is_some()
        || cfg
            .structures
            .iter()
            .any(|spec| spec.restart_to.is_some())
    {
        write_restart_outputs(cfg, &placements, &mol_spec)?;
    }

    ter_after.sort_unstable();
    ter_after.dedup();

    Ok(PackOutput {
        atoms,
        bonds,
        box_size,
        ter_after,
    })
}

fn write_restart_outputs(
    cfg: &PackConfig,
    placements: &[PlacementRecord],
    mol_spec: &[usize],
) -> PackResult<()> {
    if let Some(path) = &cfg.restart_to {
        let entries: Vec<RestartEntry> = placements
            .iter()
            .map(|p| RestartEntry {
                center: p.center,
                euler: p.euler,
            })
            .collect();
        write_restart(Path::new(path), &entries)?;
    }
    for (spec_index, spec) in cfg.structures.iter().enumerate() {
        let Some(path) = &spec.restart_to else {
            continue;
        };
        let mut entries = Vec::new();
        for (mol_idx, &spec_idx) in mol_spec.iter().enumerate() {
            if spec_idx == spec_index {
                let placement = placements
                    .get(mol_idx)
                    .copied()
                    .unwrap_or_default();
                entries.push(RestartEntry {
                    center: placement.center,
                    euler: placement.euler,
                });
            }
        }
        write_restart(Path::new(path), &entries)?;
    }
    Ok(())
}
