use std::collections::VecDeque;

use crate::config::PackConfig;
use crate::error::{PackError, PackResult};
use crate::gencan_objective::{
    build_bounds, compute_gradient, compute_objective, compute_objective_only, grad_norm,
    grad_sup_norm, pack_variables, project_step, projected_gradient, update_positions,
    validate_molecule_counts, Bounds, MolInfo,
};
use crate::atom_params::AtomParams;
use crate::geom::Vec3;
use crate::pack::AtomRecord;
use crate::pack_ops::TemplateEntry;
use crate::pbc::PbcBox;
use crate::placement::PlacementRecord;

#[derive(Clone, Copy, Debug, Default)]
pub(crate) struct GencanResult {
    pub(crate) value: f32,
    pub(crate) max_overlap: f32,
    pub(crate) max_constraint: f32,
}

pub(crate) fn optimize_gencan(
    cfg: &PackConfig,
    atoms: &mut Vec<AtomRecord>,
    positions: &mut Vec<Vec3>,
    atom_params: &[AtomParams],
    placements: &mut [PlacementRecord],
    mol_spec: &[usize],
    templates: &[TemplateEntry],
    box_size: [f32; 3],
    box_origin: [f32; 3],
    pbc: Option<PbcBox>,
    include_overlap: bool,
    active_spec: Option<usize>,
    active_mol: Option<usize>,
) -> PackResult<GencanResult> {
    if !cfg.use_gencan {
        return Ok(GencanResult::default());
    }
    let total_mols = placements.len();
    if total_mols == 0 {
        return Ok(GencanResult::default());
    }
    let mol_info = build_mol_info(cfg, atoms, mol_spec, templates, active_spec, active_mol)?;

    let mut x = pack_variables(placements);
    let bounds = build_bounds(&mol_info, placements, box_size, box_origin);
    let maxit = cfg.gencan_maxit.unwrap_or(200);
    let gencan_step = cfg.gencan_step;
    let precision = cfg.precision.unwrap_or(1.0e-2);
    let use_short_tol = cfg.use_short_tol;
    let short_dist = cfg.short_tol_dist.unwrap_or(0.0);
    let short_scale = cfg.short_tol_scale.unwrap_or(1.0);
    let fbins = cfg.fbins.unwrap_or(3.0f32.sqrt());
    let radmax = atom_params
        .iter()
        .fold(0.0f32, |acc, p| acc.max(p.radius));

    let nonmon_m = 10usize;
    let gamma = 1.0e-4f32;
    let beta = 0.5f32;
    let theta = 1.0e-6f32;
    let eta = 0.9f32;
    let sigma1 = 0.1f32;
    let sigma2 = 0.9f32;
    let nint = 2.0f32;
    let min_interp = 4usize;
    let max_ls = 40usize;
    let epsgpsn = 1.0e-6f32;
    let lspg_min = 1.0e-10f32;
    let lspg_max = 1.0e10f32;
    let delmin = 2.0f32;

    let mut f_hist = VecDeque::with_capacity(nonmon_m);
    let mut prev_x: Option<Vec<f32>> = None;
    let mut prev_g: Option<Vec<f32>> = None;
    let mut pg_norm0: Option<f32> = None;
    let ometa2 = (1.0 - eta) * (1.0 - eta);

    if cfg.chkgrad {
        update_positions(&mut x, placements, &mol_info, atoms, positions);
            let obj = compute_objective(
                atoms,
                positions,
                atom_params,
                &mol_info,
                pbc,
                use_short_tol,
                include_overlap,
                short_dist,
                short_scale,
                fbins,
                radmax,
            );
        let g = compute_gradient(&obj, &x, &mol_info);
        let eps = 1.0e-4f32;
        let mut max_rel = 0.0f32;
        let ncheck = x.len().min(20);
        for i in 0..ncheck {
            let mut x_plus = x.clone();
            let mut x_minus = x.clone();
            x_plus[i] += eps;
            x_minus[i] -= eps;
            update_positions(&mut x_plus, placements, &mol_info, atoms, positions);
            let f_plus = compute_objective_only(
                atoms,
                positions,
                atom_params,
                &mol_info,
                pbc,
                use_short_tol,
                include_overlap,
                short_dist,
                short_scale,
                fbins,
                radmax,
            );
            update_positions(&mut x_minus, placements, &mol_info, atoms, positions);
            let f_minus = compute_objective_only(
                atoms,
                positions,
                atom_params,
                &mol_info,
                pbc,
                use_short_tol,
                include_overlap,
                short_dist,
                short_scale,
                fbins,
                radmax,
            );
            let fd = (f_plus - f_minus) / (2.0 * eps);
            let denom = fd.abs().max(g[i].abs()).max(1.0e-6);
            let rel = (fd - g[i]).abs() / denom;
            if rel > max_rel {
                max_rel = rel;
            }
        }
        if max_rel > 1.0e-2 {
            return Err(PackError::Invalid(format!(
                "chkgrad failed: max relative error {:.3e}",
                max_rel
            )));
        }
        update_positions(&mut x, placements, &mol_info, atoms, positions);
    }

    for iter in 0..maxit {
        update_positions(&mut x, placements, &mol_info, atoms, positions);
        let obj = compute_objective(
            atoms,
            positions,
            atom_params,
            &mol_info,
            pbc,
            use_short_tol,
            include_overlap,
            short_dist,
            short_scale,
            fbins,
            radmax,
        );
        if obj.max_overlap <= precision && obj.max_constraint <= precision {
            break;
        }
        let g = compute_gradient(&obj, &x, &mol_info);
        let pg = projected_gradient(&x, &g, &bounds);
        let pg_norm = grad_norm(&pg);
        let pg_sup = grad_sup_norm(&pg);
        let gpeucn2 = pg_norm * pg_norm;
        if pg_norm0.is_none() {
            pg_norm0 = Some(pg_norm.max(1.0e-6));
        }
        if pg_sup <= epsgpsn {
            break;
        }
        if pg_norm < 1.0e-6 {
            break;
        }

        let xnorm = norm2(&x).sqrt();
        let (sts, sty) = match (&prev_x, &prev_g) {
            (Some(px), Some(pg_prev)) => {
                let mut s_dot_s = 0.0f32;
                let mut s_dot_y = 0.0f32;
                for i in 0..x.len() {
                    let s = x[i] - px[i];
                    let y = g[i] - pg_prev[i];
                    s_dot_s += s * s;
                    s_dot_y += s * y;
                }
                (s_dot_s, s_dot_y)
            }
            _ => (0.0, 0.0),
        };
        let gieucn2 = internal_pg_norm2(&x, &pg, &bounds, 1.0e-6);
        let use_spg = gieucn2 <= ometa2 * gpeucn2;

        let mut lamspg = if iter == 0 || sty <= 0.0 {
            let denom = gpeucn2.max(1.0e-12).sqrt();
            (1.0f32.max(xnorm)) / denom
        } else {
            (sts / sty).abs()
        };
        if iter == 0 {
            if let Some(step) = gencan_step {
                lamspg = step.abs();
            }
        }
        let lamspg = lamspg.clamp(lspg_min, lspg_max);
        let spg_dir = spg_direction(&x, &g, &bounds, lamspg);

        let mut direction = if use_spg {
            spg_dir.clone()
        } else {
            truncated_newton_direction(
                &x,
                &g,
                &bounds,
                atoms,
                positions,
                atom_params,
                placements,
                &mol_info,
                pbc,
                use_short_tol,
                include_overlap,
                short_dist,
                short_scale,
                fbins,
                radmax,
            )
        };
        let mut gtd = dot(&g, &direction);
        if gtd >= 0.0 {
            direction = spg_dir.clone();
            gtd = dot(&g, &direction);
            if gtd >= 0.0 {
                break;
            }
        }

        if !use_spg {
            let gnorm = grad_norm(&g);
            let dnorm = grad_norm(&direction);
            if gtd > -theta * gnorm * dnorm {
                direction = spg_dir.clone();
                gtd = dot(&g, &direction);
                if gtd >= 0.0 {
                    break;
                }
            }
            let delta = if iter == 0 {
                delmin.max(0.1 * 1.0f32.max(xnorm))
            } else {
                delmin.max(10.0 * sts.max(0.0).sqrt())
            };
            let dnorm = grad_norm(&direction);
            if dnorm > delta && dnorm > 0.0 {
                let scale = delta / dnorm;
                for v in direction.iter_mut() {
                    *v *= scale;
                }
                gtd = dot(&g, &direction);
            }
        }

        let f_ref = f_hist
            .iter()
            .cloned()
            .fold(obj.value, |a, b| a.max(b));
        let (x_new, f_new) = match line_search(
            &x,
            &direction,
            gtd,
            &bounds,
            obj.value,
            f_ref,
            gamma,
            if use_spg { None } else { Some(beta) },
            sigma1,
            sigma2,
            nint,
            min_interp,
            max_ls,
            atoms,
            positions,
            atom_params,
            placements,
            &mol_info,
            pbc,
            use_short_tol,
            include_overlap,
            short_dist,
            short_scale,
            fbins,
            radmax,
        ) {
            Some(v) => v,
            None => break,
        };

        let mut x_eval = x_new.clone();
        update_positions(&mut x_eval, placements, &mol_info, atoms, positions);
        let new_obj = compute_objective(
            atoms,
            positions,
            atom_params,
            &mol_info,
            pbc,
            use_short_tol,
            include_overlap,
            short_dist,
            short_scale,
            fbins,
            radmax,
        );
        let _g_new = compute_gradient(&new_obj, &x_new, &mol_info);

        if f_hist.len() == nonmon_m {
            f_hist.pop_front();
        }
        f_hist.push_back(f_new);

        prev_x = Some(x);
        prev_g = Some(g);
        x = x_new;

        if let Some(freq) = cfg.iprint1 {
            if freq > 0 && (iter as i32) % freq == 0 {
                eprintln!(
                    "gencan iter {} f={:.6} pg={:.6} max_overlap={:.6} max_constraint={:.6}",
                    iter, new_obj.value, pg_norm, new_obj.max_overlap, new_obj.max_constraint
                );
            }
        }
    }

    update_positions(&mut x, placements, &mol_info, atoms, positions);
    update_positions(&mut x, placements, &mol_info, atoms, positions);
    let obj = compute_objective(
        atoms,
        positions,
        atom_params,
        &mol_info,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
    );
    Ok(GencanResult {
        value: obj.value,
        max_overlap: obj.max_overlap,
        max_constraint: obj.max_constraint,
    })
}

pub(crate) fn evaluate_state(
    cfg: &PackConfig,
    atoms: &[AtomRecord],
    positions: &[Vec3],
    atom_params: &[AtomParams],
    mol_spec: &[usize],
    templates: &[TemplateEntry],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    active_spec: Option<usize>,
    active_mol: Option<usize>,
) -> PackResult<GencanResult> {
    let mol_info = build_mol_info(cfg, atoms, mol_spec, templates, active_spec, active_mol)?;
    let radmax = atom_params
        .iter()
        .fold(0.0f32, |acc, p| acc.max(p.radius));
    let short_dist = if use_short_tol {
        cfg.short_tol_dist.unwrap_or(0.0)
    } else {
        0.0
    };
    let short_scale = if use_short_tol {
        cfg.short_tol_scale.unwrap_or(1.0)
    } else {
        1.0
    };
    let obj = compute_objective(
        atoms,
        positions,
        atom_params,
        &mol_info,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        cfg.fbins.unwrap_or(3.0f32.sqrt()),
        radmax,
    );
    Ok(GencanResult {
        value: obj.value,
        max_overlap: obj.max_overlap,
        max_constraint: obj.max_constraint,
    })
}

fn build_mol_info(
    cfg: &PackConfig,
    atoms: &[AtomRecord],
    mol_spec: &[usize],
    templates: &[TemplateEntry],
    active_spec: Option<usize>,
    active_mol: Option<usize>,
) -> PackResult<Vec<MolInfo>> {
    let total_mols = mol_spec.len();
    let mut mol_atoms: Vec<Vec<usize>> = vec![Vec::new(); total_mols];
    for (idx, atom) in atoms.iter().enumerate() {
        let mol_idx = atom.mol_id.max(1) as usize - 1;
        if mol_idx < total_mols {
            mol_atoms[mol_idx].push(idx);
        }
    }
    let mut mol_info = Vec::with_capacity(total_mols);
    for (mol_idx, &spec_index) in mol_spec.iter().enumerate() {
        let spec = &cfg.structures[spec_index];
        let entry = templates
            .get(spec_index)
            .ok_or_else(|| PackError::Placement("missing template entry".into()))?;
        let active_spec_ok = active_spec.map_or(true, |idx| idx == spec_index);
        let active_mol_ok = active_mol.map_or(true, |idx| idx == mol_idx);
        let active = active_spec_ok && active_mol_ok;
        let atom_indices = mol_atoms[mol_idx].clone();
        validate_molecule_counts(&atom_indices, entry.template.atoms.len())?;
        let local_positions = entry
            .template
            .atoms
            .iter()
            .map(|a| a.position)
            .collect::<Vec<_>>();
        mol_info.push(MolInfo {
            atom_indices,
            local_positions,
            radius: entry.template.radius,
            movable: !spec.fixed && active,
            rotatable: spec.rotate && active,
            constraints: if active { spec.constraints.clone() } else { Vec::new() },
            rot_bounds: spec.rot_bounds,
        });
    }
    Ok(mol_info)
}

fn spg_direction(x: &[f32], g: &[f32], bounds: &Bounds, lamspg: f32) -> Vec<f32> {
    let trial = project_step(x, g, -lamspg, bounds);
    let mut dir = Vec::with_capacity(x.len());
    for i in 0..x.len() {
        dir.push(trial[i] - x[i]);
    }
    dir
}

fn line_search(
    x: &[f32],
    direction: &[f32],
    gtd: f32,
    bounds: &Bounds,
    f0: f32,
    f_ref: f32,
    gamma: f32,
    beta: Option<f32>,
    sigma1: f32,
    sigma2: f32,
    nint: f32,
    min_interp: usize,
    max_ls: usize,
    atoms: &mut [AtomRecord],
    positions: &mut [Vec3],
    atom_params: &[AtomParams],
    placements: &mut [PlacementRecord],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> Option<(Vec<f32>, f32)> {
    if gtd >= 0.0 {
        return None;
    }
    let mut step = 1.0f32;
    let mut interp = 0usize;
    for _ in 0..max_ls {
        let x_new = project_step(x, direction, step, bounds);
        let mut x_eval = x_new.clone();
        update_positions(&mut x_eval, placements, mols, atoms, positions);
        let f_new = compute_objective_only(
            atoms,
            positions,
            atom_params,
            mols,
            pbc,
            use_short_tol,
            include_overlap,
            short_dist,
            short_scale,
            fbins,
            radmax,
        );
        let mut accept = f_new <= f_ref + gamma * step * gtd;
        if accept {
            if let Some(beta_val) = beta {
                let g_new = eval_grad(
                    &x_new,
                    atoms,
                    positions,
                    atom_params,
                    placements,
                    mols,
                    pbc,
                    use_short_tol,
                    include_overlap,
                    short_dist,
                    short_scale,
                    fbins,
                    radmax,
                );
                let gptd = dot(&g_new, direction);
                if gptd >= beta_val * gtd {
                    accept = false;
                }
            }
        }
        if accept {
            return Some((x_new, f_new));
        }
        interp += 1;
        if interp >= min_interp && step < 1.0e-12 {
            return None;
        }
        let denom = 2.0 * (f_new - f0 - step * gtd);
        let mut step_new = if denom > 0.0 {
            -gtd * step * step / denom
        } else {
            step / nint
        };
        let lower = sigma1 * step;
        let upper = sigma2 * step;
        if step_new < lower || step_new > upper {
            step_new = step / nint;
        }
        step = step_new;
    }
    None
}

fn truncated_newton_direction(
    x: &[f32],
    g: &[f32],
    bounds: &Bounds,
    atoms: &mut [AtomRecord],
    positions: &mut [Vec3],
    atom_params: &[AtomParams],
    placements: &mut [PlacementRecord],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> Vec<f32> {
    let n = x.len();
    let mask = active_mask(x, g, bounds, 1.0e-6);
    let mut d = vec![0.0f32; n];
    let mut r = vec![0.0f32; n];
    let mut p = vec![0.0f32; n];
    for i in 0..n {
        if mask[i] {
            r[i] = -g[i];
            p[i] = r[i];
        }
    }
    let mut rs_old = dot(&r, &r);
    if rs_old.sqrt() < 1.0e-6 {
        return d;
    }
    let max_cg = 30usize;
    let tol = 1.0e-3f32 * rs_old.sqrt();

    for _ in 0..max_cg {
        let ap = hess_vec(
            x,
            &p,
            bounds,
            atoms,
            positions,
            atom_params,
            placements,
            mols,
            pbc,
            use_short_tol,
            include_overlap,
            short_dist,
            short_scale,
            fbins,
            radmax,
        );
        let p_ap = dot(&p, &ap);
        if p_ap <= 1.0e-8 {
            break;
        }
        let alpha = rs_old / p_ap;
        for i in 0..n {
            if mask[i] {
                d[i] += alpha * p[i];
                r[i] -= alpha * ap[i];
            }
        }
        let rs_new = dot(&r, &r);
        if rs_new.sqrt() < tol {
            break;
        }
        let beta = rs_new / rs_old;
        for i in 0..n {
            if mask[i] {
                p[i] = r[i] + beta * p[i];
            }
        }
        rs_old = rs_new;
    }
    d
}

fn hess_vec(
    x: &[f32],
    v: &[f32],
    bounds: &Bounds,
    atoms: &mut [AtomRecord],
    positions: &mut [Vec3],
    atom_params: &[AtomParams],
    placements: &mut [PlacementRecord],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> Vec<f32> {
    let n = x.len();
    let base = eval_grad(
        x,
        atoms,
        positions,
        atom_params,
        placements,
        mols,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
    );
    let norm_v = grad_norm(v);
    let eps = if norm_v > 0.0 { 1.0e-3 / norm_v } else { 1.0e-3 };
    let mut x2 = Vec::with_capacity(n);
    for i in 0..n {
        x2.push(x[i] + eps * v[i]);
    }
    let x2 = project_step(x, &x2_sub(&x2, x), 1.0, bounds);
    let g2 = eval_grad(
        &x2,
        atoms,
        positions,
        atom_params,
        placements,
        mols,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
    );
    let mut hv = vec![0.0f32; n];
    for i in 0..n {
        hv[i] = (g2[i] - base[i]) / eps;
    }
    hv
}

fn eval_grad(
    x: &[f32],
    atoms: &mut [AtomRecord],
    positions: &mut [Vec3],
    atom_params: &[AtomParams],
    placements: &mut [PlacementRecord],
    mols: &[MolInfo],
    pbc: Option<PbcBox>,
    use_short_tol: bool,
    include_overlap: bool,
    short_dist: f32,
    short_scale: f32,
    fbins: f32,
    radmax: f32,
) -> Vec<f32> {
    let mut x_eval = x.to_vec();
    update_positions(&mut x_eval, placements, mols, atoms, positions);
    let obj = compute_objective(
        atoms,
        positions,
        atom_params,
        mols,
        pbc,
        use_short_tol,
        include_overlap,
        short_dist,
        short_scale,
        fbins,
        radmax,
    );
    compute_gradient(&obj, x, mols)
}

fn active_mask(x: &[f32], g: &[f32], bounds: &Bounds, eps: f32) -> Vec<bool> {
    let mut mask = vec![true; x.len()];
    for i in 0..x.len() {
        let lo = bounds.lower[i];
        let hi = bounds.upper[i];
        if (x[i] <= lo + eps && g[i] >= 0.0) || (x[i] >= hi - eps && g[i] <= 0.0) {
            mask[i] = false;
        }
    }
    mask
}

fn dot(a: &[f32], b: &[f32]) -> f32 {
    let mut sum = 0.0f32;
    for i in 0..a.len() {
        sum += a[i] * b[i];
    }
    sum
}

fn norm2(a: &[f32]) -> f32 {
    dot(a, a)
}

fn x2_sub(x2: &[f32], x: &[f32]) -> Vec<f32> {
    let mut out = Vec::with_capacity(x.len());
    for i in 0..x.len() {
        out.push(x2[i] - x[i]);
    }
    out
}

fn internal_pg_norm2(x: &[f32], pg: &[f32], bounds: &Bounds, eps: f32) -> f32 {
    let mut sum = 0.0f32;
    for i in 0..x.len() {
        let lo = bounds.lower[i];
        let hi = bounds.upper[i];
        if x[i] > lo + eps && x[i] < hi - eps {
            sum += pg[i] * pg[i];
        }
    }
    sum
}
