use std::collections::HashSet;

use rand::{rngs::StdRng, Rng};

use crate::config::{PackConfig, StructureSpec};
use crate::constraint_penalty::penalty_and_grad;
use crate::constraints::satisfies_constraints;
use crate::error::{PackError, PackResult};
use crate::geom::{Quaternion, Vec3};
use crate::atom_params::AtomParams;
use crate::pack::AtomRecord;
use crate::pack_ops::{append_bonds, append_ter_after, apply_chain, apply_resnumbers, random_center, transform_atoms, TemplateEntry};
use crate::gencan::optimize_gencan;
use crate::pbc::PbcBox;
use crate::placement::PlacementRecord;
use crate::spatial_hash::{SpatialHash, SpatialHashParamsExt};

pub(crate) fn run_movebad_pass(
    cfg: &PackConfig,
    templates: &[TemplateEntry],
    atoms: &mut Vec<AtomRecord>,
    positions: &mut Vec<Vec3>,
    params: &mut Vec<AtomParams>,
    hash: &mut SpatialHash,
    bonds: &mut Vec<(usize, usize)>,
    ter_after: &mut Vec<usize>,
    placements: &mut [PlacementRecord],
    box_size: [f32; 3],
    box_origin: Vec3,
    pbc: Option<PbcBox>,
    rng: &mut StdRng,
    global_resid_counter: &mut i32,
    mol_id: &mut i32,
    mol_spec: &[usize],
    cell_size: f32,
) -> PackResult<()> {
    let movefrac = cfg.movefrac.unwrap_or(0.05);
    if movefrac <= 0.0 {
        return Ok(());
    }
    let precision = cfg.precision.unwrap_or(1.0e-2);
    let (bounds_min, bounds_max) = compute_bounds(cfg, positions, params, box_origin, box_size);
    let per_spec = collect_molecules_per_spec(atoms, mol_spec, templates.len());
    let total_mols = mol_spec.len();
    let atom_overlap = compute_atom_overlap(positions, params, atoms, pbc, cell_size);
    let atom_constraint = compute_atom_constraint(atoms, positions, cfg, mol_spec);
    let (mol_fdist, mol_frest) = compute_mol_scores(total_mols, atoms, &atom_overlap, &atom_constraint);
    let mut moved_any = false;

    for (spec_index, mol_ids) in per_spec.iter().enumerate() {
        if mol_ids.is_empty() {
            continue;
        }
        let spec = &cfg.structures[spec_index];
        let entry = &templates[spec_index];
        let mut scores = Vec::with_capacity(mol_ids.len());
        let mut nbad = 0usize;
        for &mol_id in mol_ids {
            let mol_idx = mol_id.max(1) as usize - 1;
            let fdist = *mol_fdist.get(mol_idx).unwrap_or(&0.0);
            let frest = *mol_frest.get(mol_idx).unwrap_or(&0.0);
            let bad = fdist > precision || frest > precision;
            if bad {
                nbad += 1;
            }
            let score = if bad { fdist + frest } else { 0.0 };
            scores.push((mol_id, score));
        }
        if nbad == 0 {
            continue;
        }
        let frac = (nbad as f32 / mol_ids.len() as f32).min(movefrac);
        let good_pool_len = ((mol_ids.len() as f32) * frac).floor().max(1.0) as usize;
        let mut nmove = ((mol_ids.len() as f32) * frac).max(1.0).round() as usize;
        nmove = nmove.min(entry.maxmove.max(1)).min(mol_ids.len());
        scores.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
        let good_pool_len = good_pool_len.min(scores.len());
        let bad_start = scores.len().saturating_sub(nmove);
        let bad_list: Vec<i32> = scores[bad_start..].iter().map(|v| v.0).collect();

        for target in bad_list {
            remove_molecule(
                atoms,
                positions,
                params,
                hash,
                bonds,
                ter_after,
                target,
                cell_size,
            );
            let good_idx = rng.gen_range(0..good_pool_len);
            let good_mol = scores[good_idx].0;
            let good_index = (good_mol.max(1) - 1) as usize;
            let anchor_center = if cfg.movebadrandom {
                None
            } else {
                placements.get(good_index).map(|p| p.center)
            };
            let anchor_euler = placements.get(good_index).map(|p| p.euler);
            let bounds = if cfg.movebadrandom {
                Some((bounds_min, bounds_max))
            } else {
                None
            };
            let (record, needs_constraints) = place_one_molecule(
                cfg,
                spec,
                entry,
                atoms,
                positions,
                params,
                hash,
                bonds,
                ter_after,
                box_size,
                box_origin,
                pbc,
                rng,
                *global_resid_counter,
                target,
                anchor_center,
                anchor_euler,
                bounds,
                false,
            )?;
            if let Some(slot) = placements.get_mut((target - 1) as usize) {
                *slot = record;
            }
            if needs_constraints && !spec.constraints.is_empty() {
                let _ = optimize_gencan(
                    cfg,
                    atoms,
                    positions,
                    params,
                    placements,
                    mol_spec,
                    templates,
                    box_size,
                    box_origin.to_array(),
                    pbc,
                    false,
                    None,
                    Some((target - 1) as usize),
                )?;
                *hash = SpatialHash::new(cell_size);
                for (idx, pos) in positions.iter().copied().enumerate() {
                    hash.insert(idx, pos);
                }
            }
            *global_resid_counter += entry.template.resid_count;
            moved_any = true;
        }
    }

    if moved_any {
        *mol_id = atoms.iter().map(|a| a.mol_id).max().unwrap_or(0) + 1;
    }
    Ok(())
}

fn collect_molecules_per_spec(
    atoms: &[AtomRecord],
    mol_spec: &[usize],
    n_specs: usize,
) -> Vec<Vec<i32>> {
    let mut per_spec: Vec<Vec<i32>> = vec![Vec::new(); n_specs];
    let mut seen: HashSet<i32> = HashSet::new();
    for atom in atoms {
        let mol_id = atom.mol_id;
        if !seen.insert(mol_id) {
            continue;
        }
        let idx = (mol_id - 1) as usize;
        if let Some(&spec_index) = mol_spec.get(idx) {
            if spec_index < per_spec.len() {
                per_spec[spec_index].push(mol_id);
            }
        }
    }
    per_spec
}

fn compute_atom_overlap(
    positions: &[Vec3],
    params: &[AtomParams],
    atoms: &[AtomRecord],
    pbc: Option<PbcBox>,
    cell_size: f32,
) -> Vec<f32> {
    let mut overlap = vec![0.0f32; positions.len()];
    let mut hash = SpatialHash::new(cell_size);
    for (i, pos_i) in positions.iter().enumerate() {
        hash.for_each_neighbor(*pos_i, |j| {
            if j >= i {
                return;
            }
            let mol_i = atoms[i].mol_id;
            let mol_j = atoms[j].mol_id;
            if mol_i == mol_j {
                return;
            }
            let delta = if let Some(pbc_box) = pbc {
                pbc_box.delta(*pos_i, positions[j])
            } else {
                pos_i.sub(positions[j])
            };
            let target = params[i].radius + params[j].radius;
            let d2 = delta.dot(delta);
            let tol2 = target * target;
            if d2 < tol2 {
                let diff = tol2 - d2;
                overlap[i] = overlap[i].max(diff);
                overlap[j] = overlap[j].max(diff);
            }
        });
        hash.insert(i, *pos_i);
    }
    overlap
}

fn compute_atom_constraint(
    atoms: &[AtomRecord],
    positions: &[Vec3],
    cfg: &PackConfig,
    mol_spec: &[usize],
) -> Vec<f32> {
    let mut violations = vec![0.0f32; positions.len()];
    for (i, atom) in atoms.iter().enumerate() {
        let mol_idx = atom.mol_id.max(1) as usize - 1;
        let Some(&spec_idx) = mol_spec.get(mol_idx) else {
            continue;
        };
        let spec = &cfg.structures[spec_idx];
        if spec.constraints.is_empty() {
            continue;
        }
        let mut max_v = 0.0f32;
        for constraint in &spec.constraints {
            let pen = penalty_and_grad(positions[i], constraint);
            if pen.value > max_v {
                max_v = pen.value;
            }
        }
        violations[i] = max_v;
    }
    violations
}

fn compute_mol_scores(
    total_mols: usize,
    atoms: &[AtomRecord],
    atom_overlap: &[f32],
    atom_constraint: &[f32],
) -> (Vec<f32>, Vec<f32>) {
    let mut fdist = vec![0.0f32; total_mols];
    let mut frest = vec![0.0f32; total_mols];
    for (i, atom) in atoms.iter().enumerate() {
        let mol_idx = atom.mol_id.max(1) as usize - 1;
        if mol_idx >= total_mols {
            continue;
        }
        fdist[mol_idx] = fdist[mol_idx].max(atom_overlap[i]);
        frest[mol_idx] = frest[mol_idx].max(atom_constraint[i]);
    }
    (fdist, frest)
}

fn remove_molecule(
    atoms: &mut Vec<AtomRecord>,
    positions: &mut Vec<Vec3>,
    params: &mut Vec<AtomParams>,
    hash: &mut SpatialHash,
    bonds: &mut Vec<(usize, usize)>,
    ter_after: &mut Vec<usize>,
    target_mol: i32,
    cell_size: f32,
) {
    let mut keep_indices = Vec::new();
    for (i, atom) in atoms.iter().enumerate() {
        if atom.mol_id != target_mol {
            keep_indices.push(i);
        }
    }
    if keep_indices.len() == atoms.len() {
        return;
    }
    let mut new_atoms = Vec::with_capacity(keep_indices.len());
    let mut new_positions = Vec::with_capacity(keep_indices.len());
    let mut new_params = Vec::with_capacity(keep_indices.len());
    let mut index_map = vec![None; atoms.len()];
    for (new_idx, &old_idx) in keep_indices.iter().enumerate() {
        new_atoms.push(atoms[old_idx].clone());
        new_positions.push(positions[old_idx]);
        new_params.push(params[old_idx]);
        index_map[old_idx] = Some(new_idx);
    }
    let mut new_bonds = Vec::new();
    for (a, b) in bonds.iter().copied() {
        if let (Some(na), Some(nb)) = (index_map[a], index_map[b]) {
            new_bonds.push((na, nb));
        }
    }
    let mut new_ter = Vec::new();
    for &idx in ter_after.iter() {
        if let Some(nidx) = index_map.get(idx).and_then(|v| *v) {
            new_ter.push(nidx);
        }
    }
    *atoms = new_atoms;
    *positions = new_positions;
    *params = new_params;
    *bonds = new_bonds;
    *ter_after = new_ter;
    *hash = SpatialHash::new(cell_size);
    for (i, pos) in positions.iter().copied().enumerate() {
        hash.insert(i, pos);
    }
}

fn compute_bounds(
    cfg: &PackConfig,
    positions: &[Vec3],
    params: &[AtomParams],
    origin: Vec3,
    box_size: [f32; 3],
) -> (Vec3, Vec3) {
    if let Some(sidemax) = cfg.sidemax {
        let center = Vec3::new(
            origin.x + box_size[0] * 0.5,
            origin.y + box_size[1] * 0.5,
            origin.z + box_size[2] * 0.5,
        );
        let min = Vec3::new(center.x - sidemax, center.y - sidemax, center.z - sidemax);
        let max = Vec3::new(center.x + sidemax, center.y + sidemax, center.z + sidemax);
        return (min, max);
    }
    if positions.is_empty() {
        let max = Vec3::new(origin.x + box_size[0], origin.y + box_size[1], origin.z + box_size[2]);
        return (origin, max);
    }
    let mut min = Vec3::new(f32::INFINITY, f32::INFINITY, f32::INFINITY);
    let mut max = Vec3::new(f32::NEG_INFINITY, f32::NEG_INFINITY, f32::NEG_INFINITY);
    for pos in positions {
        min.x = min.x.min(pos.x);
        min.y = min.y.min(pos.y);
        min.z = min.z.min(pos.z);
        max.x = max.x.max(pos.x);
        max.y = max.y.max(pos.y);
        max.z = max.z.max(pos.z);
    }
    let radmax = params.iter().fold(0.0f32, |acc, p| acc.max(p.radius));
    let pad = 1.1 * radmax;
    min = Vec3::new(min.x - pad, min.y - pad, min.z - pad);
    max = Vec3::new(max.x + pad, max.y + pad, max.z + pad);
    (min, max)
}

fn place_one_molecule(
    cfg: &PackConfig,
    spec: &StructureSpec,
    entry: &TemplateEntry,
    atoms: &mut Vec<AtomRecord>,
    positions: &mut Vec<Vec3>,
    params: &mut Vec<AtomParams>,
    hash: &mut SpatialHash,
    bonds: &mut Vec<(usize, usize)>,
    ter_after: &mut Vec<usize>,
    box_size: [f32; 3],
    box_origin: Vec3,
    pbc: Option<PbcBox>,
    rng: &mut StdRng,
    global_resid_counter: i32,
    mol_id: i32,
    anchor: Option<Vec3>,
    anchor_euler: Option<[f32; 3]>,
    bounds: Option<(Vec3, Vec3)>,
    check_overlap: bool,
) -> PackResult<(PlacementRecord, bool)> {
    for _ in 0..cfg.max_attempts.unwrap_or(10000) {
        let (rotation, euler) = if spec.rotate {
            if let Some(anchor_euler) = anchor_euler {
                (
                    Quaternion::from_packmol_euler(anchor_euler[0], anchor_euler[1], anchor_euler[2]),
                    anchor_euler,
                )
            } else 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(rng);
                let (b, g, t) = rot.to_packmol_euler();
                (rot, [b, g, t])
            }
        } else {
            (Quaternion::identity(), [0.0, 0.0, 0.0])
        };
        let center = if let Some(anchor) = anchor {
            let local_anchor = anchor.sub(box_origin);
            let jitter = (0.3 * entry.template.dmax).max(1.0e-3);
            let local_center = random_center_near(rng, local_anchor, jitter)?;
            local_center.add(box_origin)
        } else if let Some((min, max)) = bounds {
            random_center_bounds(rng, min, max)
        } else {
            let mut center = random_center(rng, box_size, entry.template.radius)?;
            center = center.add(box_origin);
            center
        };
        let candidate = transform_atoms(&entry.template, rotation, center, mol_id);
        let cand_positions: Vec<Vec3> = candidate.iter().map(|a| a.position).collect();
        let constraints_ok = satisfies_constraints(&cand_positions, &spec.constraints, pbc);
        if !constraints_ok && spec.constraints.is_empty() {
            continue;
        }
        if check_overlap && hash.overlaps_params(&cand_positions, &entry.atom_params, &positions, params) {
            continue;
        }
        let mut candidate = candidate;
        apply_chain(&mut candidate, spec, mol_id);
        apply_resnumbers(
            &mut candidate,
            entry.res_mode,
            entry.template.resid_min,
            global_resid_counter,
            1,
            mol_id,
        )?;
        let base = atoms.len();
        if entry.connect {
            append_bonds(bonds, &entry.template.bonds, base);
        }
        if cfg.amber_ter_preserve {
            append_ter_after(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);
            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);
            }
        }
        return Ok((PlacementRecord::new(center, euler), !constraints_ok));
    }
    Err(PackError::Placement("movebad failed to re-place".into()))
}

fn random_center_near<R: Rng + ?Sized>(
    rng: &mut R,
    anchor: Vec3,
    jitter: f32,
) -> PackResult<Vec3> {
    let range = jitter.max(1.0e-3);
    let center = Vec3::new(
        anchor.x + rng.gen_range(-range..range),
        anchor.y + rng.gen_range(-range..range),
        anchor.z + rng.gen_range(-range..range),
    );
    Ok(center)
}

fn random_center_bounds<R: Rng + ?Sized>(rng: &mut R, min: Vec3, max: Vec3) -> Vec3 {
    Vec3::new(
        rng.gen_range(min.x..max.x),
        rng.gen_range(min.y..max.y),
        rng.gen_range(min.z..max.z),
    )
}
