use std::f32::consts::PI;

use crate::constraint_penalty::penalty_and_grad;
use crate::constraints::ConstraintSpec;
use crate::error::{PackError, PackResult};
use crate::gencan_math::{mat_vec, rotation_matrix, rotation_with_derivatives};
use crate::geom::Vec3;
use crate::atom_params::AtomParams;
use crate::pack::AtomRecord;
use crate::pbc::PbcBox;
use crate::placement::PlacementRecord;
use crate::spatial_hash::SpatialHash;

pub(crate) struct MolInfo {
    pub(crate) atom_indices: Vec<usize>,
    pub(crate) local_positions: Vec<Vec3>,
    pub(crate) radius: f32,
    pub(crate) movable: bool,
    pub(crate) rotatable: bool,
    pub(crate) constraints: Vec<ConstraintSpec>,
    pub(crate) rot_bounds: Option<[[f32; 2]; 3]>,
}

pub(crate) struct Bounds {
    pub(crate) lower: Vec<f32>,
    pub(crate) upper: Vec<f32>,
}

pub(crate) struct Objective {
    pub(crate) value: f32,
    pub(crate) max_overlap: f32,
    pub(crate) max_constraint: f32,
    pub(crate) grad_pos: Vec<Vec3>,
}

pub(crate) fn pack_variables(placements: &[PlacementRecord]) -> Vec<f32> {
    let mut x = Vec::with_capacity(placements.len() * 6);
    for placement in placements {
        x.push(placement.center.x);
        x.push(placement.center.y);
        x.push(placement.center.z);
        x.extend_from_slice(&placement.euler);
    }
    x
}

pub(crate) fn build_bounds(
    mols: &[MolInfo],
    placements: &[PlacementRecord],
    box_size: [f32; 3],
    box_origin: [f32; 3],
) -> Bounds {
    let mut lower = Vec::with_capacity(mols.len() * 6);
    let mut upper = Vec::with_capacity(mols.len() * 6);
    for (idx, mol) in mols.iter().enumerate() {
        let placement = placements[idx];
        let mut bounds = [
            (placement.center.x, placement.center.x),
            (placement.center.y, placement.center.y),
            (placement.center.z, placement.center.z),
        ];
        if mol.movable {
            bounds[0] = center_bounds(placement.center.x, mol.radius, box_size[0], box_origin[0]);
            bounds[1] = center_bounds(placement.center.y, mol.radius, box_size[1], box_origin[1]);
            bounds[2] = center_bounds(placement.center.z, mol.radius, box_size[2], box_origin[2]);
        }
        for (lo, hi) in bounds {
            lower.push(lo);
            upper.push(hi);
        }
        if mol.rotatable {
            if let Some(bounds) = mol.rot_bounds {
                for axis in 0..3 {
                    lower.push(bounds[axis][0]);
                    upper.push(bounds[axis][1]);
                }
            } else {
                lower.extend_from_slice(&[-PI, -PI, -PI]);
                upper.extend_from_slice(&[PI, PI, PI]);
            }
        } else {
            lower.extend_from_slice(&placement.euler);
            upper.extend_from_slice(&placement.euler);
        }
    }
    Bounds { lower, upper }
}

fn center_bounds(value: f32, radius: f32, size: f32, origin: f32) -> (f32, f32) {
    if size > 2.0 * radius {
        (origin + radius, origin + size - radius)
    } else {
        (value, value)
    }
}

pub(crate) fn update_positions(
    x: &mut [f32],
    placements: &mut [PlacementRecord],
    mols: &[MolInfo],
    atoms: &mut [AtomRecord],
    positions: &mut [Vec3],
) {
    let mut offset = 0;
    for (mol_idx, mol) in mols.iter().enumerate() {
        let center = Vec3::new(x[offset], x[offset + 1], x[offset + 2]);
        let beta = x[offset + 3];
        let gamma = x[offset + 4];
        let teta = x[offset + 5];
        placements[mol_idx].center = center;
        placements[mol_idx].euler = [beta, gamma, teta];
        let rot = rotation_matrix(beta, gamma, teta);
        for (local_idx, &atom_idx) in mol.atom_indices.iter().enumerate() {
            let local = mol.local_positions[local_idx];
            let rotated = mat_vec(rot, local);
            let pos = rotated.add(center);
            positions[atom_idx] = pos;
            atoms[atom_idx].position = pos;
        }
        offset += 6;
    }
}

pub(crate) fn compute_objective(
    atoms: &[AtomRecord],
    positions: &[Vec3],
    atom_params: &[AtomParams],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> Objective {
    let mut grad_pos = vec![Vec3::default(); atoms.len()];
    let (value, max_overlap, max_constraint) = compute_penalty(
        atoms,
        positions,
        atom_params,
        mols,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
        Some(&mut grad_pos),
    );
    Objective {
        value,
        max_overlap,
        max_constraint,
        grad_pos,
    }
}

pub(crate) fn compute_objective_only(
    atoms: &[AtomRecord],
    positions: &[Vec3],
    atom_params: &[AtomParams],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> f32 {
    let (value, _, _) = compute_penalty(
        atoms,
        positions,
        atom_params,
        mols,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
        None,
    );
    value
}

pub(crate) fn compute_penalty(
    atoms: &[AtomRecord],
    positions: &[Vec3],
    atom_params: &[AtomParams],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    _short_dist: f32,
    _short_scale: f32,
    fbins: f32,
    radmax: f32,
    mut grad_pos: Option<&mut Vec<Vec3>>,
) -> (f32, f32, f32) {
    let cell_size = (2.0 * radmax * fbins).max(1.0e-6);
    let mut hash = SpatialHash::new(cell_size);
    let mut value = 0.0f32;
    let mut max_overlap = 0.0f32;
    let mut max_constraint = 0.0f32;

    for (i, pos_i) in positions.iter().enumerate() {
        if include_overlap {
            hash.for_each_neighbor(*pos_i, |j| {
                if j >= i {
                    return;
                }
                let mol_i = atoms[i].mol_id.max(1) as usize - 1;
                let mol_j = atoms[j].mol_id.max(1) as usize - 1;
                if mol_i == mol_j || mol_i >= mols.len() || mol_j >= mols.len() {
                    return;
                }
                if !mols[mol_i].movable && !mols[mol_j].movable {
                    return;
                }
                let delta = if let Some(pbc_box) = pbc {
                    pbc_box.delta(*pos_i, positions[j])
                } else {
                    pos_i.sub(positions[j])
                };
                let d2 = delta.dot(delta);
                let pi = atom_params[i];
                let pj = atom_params[j];
                let tol = pi.radius + pj.radius;
                let tol2 = tol * tol;
                if d2 < tol2 {
                    let diff = d2 - tol2;
                    let weight = pi.fscale * pj.fscale;
                    value += weight * diff * diff;
                    let overlap = tol - d2.sqrt();
                    if overlap > max_overlap {
                        max_overlap = overlap;
                    }
                    if let Some(ref mut grads) = grad_pos {
                        let coeff = weight * 4.0 * diff;
                        let g = delta.scale(coeff);
                        grads[i] = grads[i].add(g);
                        grads[j] = grads[j].sub(g);
                    }
                } else if use_short_tol && (pi.use_short || pj.use_short) {
                    let short_r = pi.short_radius + pj.short_radius;
                    let short2 = short_r * short_r;
                    if d2 < short2 {
                        let diff2 = d2 - short2;
                        let scale = (pi.short_scale * pj.short_scale).sqrt();
                        let weight = pi.fscale * pj.fscale * scale;
                        value += weight * diff2 * diff2;
                        if let Some(ref mut grads) = grad_pos {
                            let coeff = weight * 4.0 * diff2;
                            let g = delta.scale(coeff);
                            grads[i] = grads[i].add(g);
                            grads[j] = grads[j].sub(g);
                        }
                    }
                }
            });
        }
        hash.insert(i, *pos_i);
    }

    for mol in mols {
        for &atom_idx in mol.atom_indices.iter() {
            let pos = positions[atom_idx];
            for constraint in &mol.constraints {
                let pen = penalty_and_grad(pos, constraint);
                if pen.value > 0.0 {
                    value += pen.value;
                    max_constraint = max_constraint.max(pen.violation.max(pen.value));
                    if let Some(ref mut grads) = grad_pos {
                        grads[atom_idx] = grads[atom_idx].add(pen.grad);
                    }
                }
            }
        }
    }

    (value, max_overlap, max_constraint)
}

pub(crate) fn compute_gradient(objective: &Objective, x: &[f32], mols: &[MolInfo]) -> Vec<f32> {
    let mut grad_x = vec![0.0f32; x.len()];
    let mut offset = 0usize;
    for mol in mols.iter() {
        let beta = x[offset + 3];
        let gamma = x[offset + 4];
        let teta = x[offset + 5];
        let (_mat, db, dg, dt) = rotation_with_derivatives(beta, gamma, teta);
        let mut grad_center = Vec3::default();
        let mut grad_beta = 0.0f32;
        let mut grad_gamma = 0.0f32;
        let mut grad_teta = 0.0f32;
        for (local_idx, &atom_idx) in mol.atom_indices.iter().enumerate() {
            let gpos = objective.grad_pos[atom_idx];
            grad_center = grad_center.add(gpos);
            if mol.rotatable {
                let local = mol.local_positions[local_idx];
                let dpos_db = mat_vec(db, local);
                let dpos_dg = mat_vec(dg, local);
                let dpos_dt = mat_vec(dt, local);
                grad_beta += gpos.dot(dpos_db);
                grad_gamma += gpos.dot(dpos_dg);
                grad_teta += gpos.dot(dpos_dt);
            }
        }
        if mol.movable {
            grad_x[offset] = grad_center.x;
            grad_x[offset + 1] = grad_center.y;
            grad_x[offset + 2] = grad_center.z;
        }
        if mol.rotatable {
            grad_x[offset + 3] = grad_beta;
            grad_x[offset + 4] = grad_gamma;
            grad_x[offset + 5] = grad_teta;
        }
        offset += 6;
    }
    grad_x
}

pub(crate) fn grad_norm(grad: &[f32]) -> f32 {
    let mut sum = 0.0f32;
    for g in grad {
        sum += g * g;
    }
    sum.sqrt()
}

pub(crate) fn grad_sup_norm(grad: &[f32]) -> f32 {
    let mut max_val = 0.0f32;
    for g in grad {
        let v = g.abs();
        if v > max_val {
            max_val = v;
        }
    }
    max_val
}

pub(crate) fn project_step(x: &[f32], direction: &[f32], step: f32, bounds: &Bounds) -> Vec<f32> {
    let mut out = Vec::with_capacity(x.len());
    for i in 0..x.len() {
        let mut v = x[i] + step * direction[i];
        let lo = bounds.lower[i];
        let hi = bounds.upper[i];
        if v < lo {
            v = lo;
        }
        if v > hi {
            v = hi;
        }
        out.push(v);
    }
    out
}

pub(crate) fn projected_gradient(x: &[f32], grad: &[f32], bounds: &Bounds) -> Vec<f32> {
    let mut pg = Vec::with_capacity(x.len());
    for i in 0..x.len() {
        let mut v = x[i] - grad[i];
        let lo = bounds.lower[i];
        let hi = bounds.upper[i];
        if v < lo {
            v = lo;
        }
        if v > hi {
            v = hi;
        }
        pg.push(x[i] - v);
    }
    pg
}

pub(crate) fn validate_molecule_counts(
    atom_indices: &[usize],
    template_len: usize,
) -> PackResult<()> {
    if atom_indices.len() != template_len {
        return Err(PackError::Placement(
            "molecule atom count does not match template".into(),
        ));
    }
    Ok(())
}
