use std::collections::BTreeMap;
use std::time::{Duration, Instant};

use bone_types::{
    BudgetCeiling, NewtonDamping, NewtonStepTolerance, Parameter, ParameterIndex, ResidualIndex,
    SolverResidual,
};
use faer::Mat;
use faer::linalg::solvers::Solve;
use faer::sparse::{SparseColMat, Triplet as FaerTriplet};
use nalgebra::DVector;

use crate::error::{Result, SolverError};
use crate::graph::Decomposition;
use crate::jacobian::{evaluate_residuals, jacobian_triplets};
use crate::residual::Triplet;
use crate::system::ConstraintSystem;

#[derive(Copy, Clone, Debug, PartialEq)]
pub struct NewtonConfig {
    pub max_iterations: u32,
    pub residual_tolerance: SolverResidual,
    pub step_tolerance: NewtonStepTolerance,
    pub line_search_shrinks: u32,
    pub damping: NewtonDamping,
    pub budget: Option<BudgetCeiling>,
}

impl NewtonConfig {
    pub const DEFAULT: Self = Self {
        max_iterations: 64,
        residual_tolerance: SolverResidual::new(1e-10),
        step_tolerance: NewtonStepTolerance::DEFAULT,
        line_search_shrinks: 20,
        damping: NewtonDamping::DEFAULT,
        budget: None,
    };
}

#[must_use]
pub fn residual_norm(residuals: &[f64]) -> SolverResidual {
    let sum: f64 = residuals.iter().map(|r| r * r).sum();
    SolverResidual::new(sum.sqrt())
}

#[derive(Copy, Clone, Debug)]
struct Stopwatch {
    started_at: Instant,
    budget: Option<BudgetCeiling>,
}

impl Stopwatch {
    fn start(budget: Option<BudgetCeiling>) -> Self {
        Self {
            started_at: Instant::now(),
            budget,
        }
    }

    fn check(self) -> Result<(), SolverError> {
        let Some(budget) = self.budget else {
            return Ok(());
        };
        let elapsed = self.started_at.elapsed();
        if elapsed >= budget.duration() {
            Err(SolverError::Budget { elapsed })
        } else {
            Ok(())
        }
    }
}

#[must_use]
fn remaining_budget(outer: BudgetCeiling, elapsed: Duration) -> BudgetCeiling {
    BudgetCeiling::new(outer.duration().saturating_sub(elapsed))
}

struct Iterate {
    params: Vec<f64>,
    residuals: Vec<f64>,
    norm: SolverResidual,
}

impl Iterate {
    fn evaluate(system: &ConstraintSystem, params: Vec<f64>) -> Self {
        let residuals = evaluate_residuals(system, &params);
        let norm = residual_norm(&residuals);
        Self {
            params,
            residuals,
            norm,
        }
    }
}

pub fn solve_newton(system: &ConstraintSystem, cfg: NewtonConfig) -> Result<Vec<Parameter>> {
    let stopwatch = Stopwatch::start(cfg.budget);
    let seed: Vec<f64> = system.parameters().iter().map(|p| p.value()).collect();
    converge(
        system,
        cfg,
        &Iterate::evaluate(system, seed),
        cfg.max_iterations,
        stopwatch,
    )
}

pub fn solve_newton_decomposed(
    system: &ConstraintSystem,
    decomposition: &Decomposition,
    cfg: NewtonConfig,
) -> Result<Vec<Parameter>> {
    let stopwatch = Stopwatch::start(cfg.budget);
    let seed: Vec<Parameter> = system.parameters().to_vec();
    decomposition
        .components()
        .iter()
        .try_fold(seed, |acc, component| {
            stopwatch.check()?;
            let sub = system.subsystem(component);
            if sub.system().row_count() == 0 {
                return Ok(acc);
            }
            let sub_cfg = NewtonConfig {
                budget: cfg
                    .budget
                    .map(|b| remaining_budget(b, stopwatch.started_at.elapsed())),
                ..cfg
            };
            let solved = solve_newton(sub.system(), sub_cfg)
                .map_err(|err| lift_subsystem_error(err, sub.param_map()))?;
            Ok(splice(acc, sub.param_map(), &solved))
        })
}

fn lift_subsystem_error(err: SolverError, param_map: &[ParameterIndex]) -> SolverError {
    match err {
        SolverError::InvalidSolutionFound { at } => SolverError::InvalidSolutionFound {
            at: param_map[at.as_usize()],
        },
        other => other,
    }
}

fn splice(
    mut whole: Vec<Parameter>,
    param_map: &[bone_types::ParameterIndex],
    sub: &[Parameter],
) -> Vec<Parameter> {
    param_map
        .iter()
        .zip(sub.iter().copied())
        .for_each(|(orig, value)| {
            whole[orig.as_usize()] = value;
        });
    whole
}

fn converge(
    system: &ConstraintSystem,
    cfg: NewtonConfig,
    state: &Iterate,
    remaining: u32,
    stopwatch: Stopwatch,
) -> Result<Vec<Parameter>> {
    if state.norm.value() < cfg.residual_tolerance.value() {
        return Ok(to_parameters(&state.params));
    }
    if remaining == 0 {
        return Err(SolverError::NoSolutionFound { last: state.norm });
    }
    stopwatch.check()?;
    let triplets = jacobian_triplets(system, &state.params);
    let step = least_squares_step(
        system.parameter_count(),
        &triplets,
        &state.residuals,
        cfg.damping,
    )?;
    if step.norm() < cfg.step_tolerance.value() {
        return Ok(to_parameters(&state.params));
    }
    match line_search(
        system,
        &state.params,
        &step,
        state.norm,
        cfg.line_search_shrinks,
    ) {
        Some(next) => converge(system, cfg, &next, remaining - 1, stopwatch),
        None => Err(SolverError::NoSolutionFound { last: state.norm }),
    }
}

fn to_parameters(values: &[f64]) -> Vec<Parameter> {
    values.iter().copied().map(Parameter::new).collect()
}

fn least_squares_step(
    params_len: usize,
    triplets: &[Triplet],
    residuals: &[f64],
    damping: NewtonDamping,
) -> Result<DVector<f64>> {
    let (jtj_triplets, jtr) =
        assemble_normal_equations(params_len, triplets, residuals, damping.value());
    let Ok(jtj) =
        SparseColMat::<usize, f64>::try_new_from_triplets(params_len, params_len, &jtj_triplets)
    else {
        unreachable!("normal-equations triplets are built with in-bounds indices")
    };
    let lu = jtj.sp_lu().map_err(|err| match err {
        faer::sparse::linalg::LuError::SymbolicSingular { index } => {
            let Ok(at) = u32::try_from(index) else {
                unreachable!("J^T J column count fits in u32 via ParameterIndex")
            };
            SolverError::InvalidSolutionFound {
                at: ParameterIndex::new(at),
            }
        }
        faer::sparse::linalg::LuError::Generic(_) => {
            unreachable!("faer sp_lu returns Generic only on malformed input")
        }
    })?;
    let mut rhs = Mat::<f64>::from_fn(params_len, 1, |i, _| -jtr[i]);
    lu.solve_in_place(rhs.as_mut());
    Ok(DVector::from_iterator(
        params_len,
        (0..params_len).map(|i| rhs[(i, 0)]),
    ))
}

fn assemble_normal_equations(
    params_len: usize,
    triplets: &[Triplet],
    residuals: &[f64],
    damping: f64,
) -> (Vec<FaerTriplet<usize, usize, f64>>, Vec<f64>) {
    let rows: BTreeMap<ResidualIndex, Vec<(usize, f64)>> =
        triplets
            .iter()
            .fold(BTreeMap::new(), |mut acc, (row, col, val)| {
                acc.entry(*row).or_default().push((col.as_usize(), *val));
                acc
            });
    let jtj: BTreeMap<(usize, usize), f64> = rows
        .values()
        .flat_map(|entries| {
            entries
                .iter()
                .flat_map(move |(i, vi)| entries.iter().map(move |(j, vj)| ((*i, *j), vi * vj)))
        })
        .fold(BTreeMap::new(), |mut acc, (key, val)| {
            *acc.entry(key).or_insert(0.0) += val;
            acc
        });
    let damped = (0..params_len).fold(jtj, |mut acc, i| {
        *acc.entry((i, i)).or_insert(0.0) += damping;
        acc
    });
    let jtj_triplets: Vec<FaerTriplet<usize, usize, f64>> = damped
        .into_iter()
        .map(|((i, j), v)| FaerTriplet::new(i, j, v))
        .collect();
    let jtr = triplets
        .iter()
        .fold(vec![0.0_f64; params_len], |mut acc, (row, col, val)| {
            acc[col.as_usize()] += val * residuals[row.as_usize()];
            acc
        });
    (jtj_triplets, jtr)
}

fn line_search(
    system: &ConstraintSystem,
    params: &[f64],
    step: &DVector<f64>,
    baseline: SolverResidual,
    max_shrinks: u32,
) -> Option<Iterate> {
    let count = usize::try_from(max_shrinks)
        .unwrap_or(usize::MAX)
        .saturating_add(1);
    std::iter::successors(Some(1.0_f64), |a| Some(a * 0.5))
        .take(count)
        .find_map(|alpha| {
            let trial: Vec<f64> = params
                .iter()
                .zip(step.iter())
                .map(|(p, s)| p + alpha * s)
                .collect();
            let candidate = Iterate::evaluate(system, trial);
            (candidate.norm.value() < baseline.value()).then_some(candidate)
        })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn lift_subsystem_error_translates_singular_index() {
        let param_map = vec![
            ParameterIndex::new(4),
            ParameterIndex::new(9),
            ParameterIndex::new(12),
        ];
        let local = SolverError::InvalidSolutionFound {
            at: ParameterIndex::new(1),
        };
        match lift_subsystem_error(local, &param_map) {
            SolverError::InvalidSolutionFound { at } => assert_eq!(at, ParameterIndex::new(9)),
            other => panic!("expected InvalidSolutionFound, got {other:?}"),
        }
    }

    #[test]
    fn lift_subsystem_error_passes_through_other_variants() {
        let param_map = [ParameterIndex::new(2), ParameterIndex::new(5)];
        let no_solution = SolverError::NoSolutionFound {
            last: SolverResidual::new(3.5),
        };
        match lift_subsystem_error(no_solution, &param_map) {
            SolverError::NoSolutionFound { last } => {
                assert!((last.value() - 3.5).abs() < f64::EPSILON);
            }
            other => panic!("expected NoSolutionFound, got {other:?}"),
        }
    }

    #[test]
    fn remaining_budget_subtracts_elapsed() {
        let outer = BudgetCeiling::new(Duration::from_millis(10));
        let r = remaining_budget(outer, Duration::from_millis(3));
        assert_eq!(r.duration(), Duration::from_millis(7));
    }

    #[test]
    fn remaining_budget_saturates_at_zero_when_overshot() {
        let outer = BudgetCeiling::new(Duration::from_millis(10));
        let r = remaining_budget(outer, Duration::from_millis(25));
        assert_eq!(r.duration(), Duration::ZERO);
    }

    #[test]
    fn decomposed_zero_budget_fires_across_disjoint_components() {
        use crate::graph::decompose;
        use crate::residual::Residual;
        use crate::system::ConstraintSystem;
        let system = ConstraintSystem::new(
            vec![Parameter::new(0.0), Parameter::new(0.0)],
            vec![
                Residual::Pin {
                    param: ParameterIndex::new(0),
                    target: 1.0,
                },
                Residual::Pin {
                    param: ParameterIndex::new(1),
                    target: 2.0,
                },
            ],
        );
        let decomp = decompose(&system);
        assert!(
            decomp.components().len() >= 2,
            "fixture must decompose into at least two components, got {}",
            decomp.components().len(),
        );
        let cfg = NewtonConfig {
            budget: Some(BudgetCeiling::new(Duration::ZERO)),
            ..NewtonConfig::DEFAULT
        };
        match solve_newton_decomposed(&system, &decomp, cfg) {
            Err(SolverError::Budget { .. }) => {}
            other => panic!("expected Budget, got {other:?}"),
        }
    }
}
