use bone_types::ParameterIndex;
use nalgebra::DMatrix;

use crate::residual::Triplet;
use crate::system::ConstraintSystem;

#[derive(Clone, Debug, PartialEq)]
pub struct SparsityPattern {
    rows: usize,
    cols: usize,
    entries: Vec<(usize, ParameterIndex)>,
}

impl SparsityPattern {
    #[must_use]
    pub fn rows(&self) -> usize {
        self.rows
    }

    #[must_use]
    pub fn cols(&self) -> usize {
        self.cols
    }

    #[must_use]
    pub fn entries(&self) -> &[(usize, ParameterIndex)] {
        &self.entries
    }

    #[must_use]
    pub fn display(&self) -> String {
        let mut grid = vec![vec![false; self.cols]; self.rows];
        self.entries.iter().for_each(|(r, c)| {
            grid[*r][c.as_usize()] = true;
        });
        grid.into_iter()
            .map(|row| {
                row.into_iter()
                    .map(|filled| if filled { 'X' } else { '.' })
                    .collect::<String>()
            })
            .collect::<Vec<_>>()
            .join("\n")
    }
}

#[must_use]
pub fn evaluate_residuals(system: &ConstraintSystem, params: &[f64]) -> Vec<f64> {
    let mut out = vec![0.0; system.row_count()];
    let offsets = system.row_offsets();
    system
        .residuals()
        .iter()
        .zip(offsets.iter().copied())
        .for_each(|(r, start)| {
            let len = r.rows();
            let begin = start.as_usize();
            r.evaluate(params, &mut out[begin..begin + len]);
        });
    out
}

#[must_use]
pub fn assemble_jacobian(system: &ConstraintSystem, params: &[f64]) -> DMatrix<f64> {
    let rows = system.row_count();
    let cols = system.parameter_count();
    let triplets = collect_triplets(system, params);
    triplets.into_iter().fold(
        DMatrix::<f64>::zeros(rows, cols),
        |mut mat, (row, col, val)| {
            mat[(row.as_usize(), col.as_usize())] += val;
            mat
        },
    )
}

#[must_use]
pub fn sparsity_pattern(system: &ConstraintSystem) -> SparsityPattern {
    let dummy: Vec<f64> = std::iter::successors(Some(1.0_f64), |v| Some(v + 0.25))
        .take(system.parameter_count())
        .collect();
    let triplets = collect_triplets(system, &dummy);
    let mut entries: Vec<(usize, ParameterIndex)> = triplets
        .into_iter()
        .map(|(r, c, _)| (r.as_usize(), c))
        .collect();
    entries.sort_unstable_by(|a, b| a.0.cmp(&b.0).then(a.1.value().cmp(&b.1.value())));
    entries.dedup();
    SparsityPattern {
        rows: system.row_count(),
        cols: system.parameter_count(),
        entries,
    }
}

#[must_use]
pub fn jacobian_triplets(system: &ConstraintSystem, params: &[f64]) -> Vec<Triplet> {
    collect_triplets(system, params)
}

fn collect_triplets(system: &ConstraintSystem, params: &[f64]) -> Vec<Triplet> {
    let offsets = system.row_offsets();
    system
        .residuals()
        .iter()
        .zip(offsets)
        .fold(Vec::new(), |mut triplets, (r, start)| {
            r.jacobian(params, start, &mut triplets);
            triplets
        })
}
