use std::collections::BTreeMap;

use bone_types::{ParameterIndex, ParentIndex, ResidualIndex};

use crate::system::ConstraintSystem;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Component {
    parameters: Vec<ParameterIndex>,
    residual_parents: Vec<ParentIndex>,
}

impl Component {
    #[must_use]
    pub fn parameters(&self) -> &[ParameterIndex] {
        &self.parameters
    }

    #[must_use]
    pub fn residual_parents(&self) -> &[ParentIndex] {
        &self.residual_parents
    }

    #[must_use]
    pub fn residual_rows(&self, system: &ConstraintSystem) -> Vec<ResidualIndex> {
        let offsets = system.row_offsets();
        self.residual_parents
            .iter()
            .flat_map(|&parent| {
                let start = offsets[parent.as_usize()].value();
                (0..system.residual_at_parent(parent).rows()).map(move |o| {
                    let Ok(ov) = u32::try_from(o) else {
                        unreachable!("residual row count is a small constant")
                    };
                    ResidualIndex::new(start + ov)
                })
            })
            .collect()
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Decomposition {
    components: Vec<Component>,
}

impl Decomposition {
    #[must_use]
    pub fn components(&self) -> &[Component] {
        &self.components
    }
}

#[must_use]
pub fn decompose(system: &ConstraintSystem) -> Decomposition {
    let n = system.parameter_count();
    let dsu_final = system.residuals().iter().fold(Dsu::new(n), |dsu, r| {
        let ps = r.parameters();
        match ps.split_first() {
            None => dsu,
            Some((&first, rest)) => rest
                .iter()
                .fold(dsu, |acc, &p| acc.union(first.as_usize(), p.as_usize())),
        }
    });
    let roots: Vec<usize> = (0..n).map(|p| dsu_final.find(p)).collect();
    let param_groups: BTreeMap<usize, Vec<ParameterIndex>> =
        (0..n).fold(BTreeMap::new(), |mut acc, p| {
            let Ok(pv) = u32::try_from(p) else {
                unreachable!("parameter count stays within ParameterIndex range")
            };
            acc.entry(roots[p])
                .or_default()
                .push(ParameterIndex::new(pv));
            acc
        });
    let parent_groups: BTreeMap<usize, Vec<ParentIndex>> = system
        .residuals()
        .iter()
        .enumerate()
        .fold(BTreeMap::new(), |mut acc, (pi, r)| {
            if let Some(&first) = r.parameters().first() {
                let Ok(piv) = u32::try_from(pi) else {
                    unreachable!("parent count fits in u32 via ResidualIndex")
                };
                acc.entry(roots[first.as_usize()])
                    .or_default()
                    .push(ParentIndex::new(piv));
            }
            acc
        });
    let mut keyed: Vec<(usize, Component)> = param_groups
        .into_iter()
        .map(|(root, parameters)| {
            let residual_parents = parent_groups.get(&root).cloned().unwrap_or_default();
            let min_param = parameters
                .first()
                .copied()
                .map_or(usize::MAX, ParameterIndex::as_usize);
            (
                min_param,
                Component {
                    parameters,
                    residual_parents,
                },
            )
        })
        .collect();
    keyed.sort_by_key(|(min_param, _)| *min_param);
    Decomposition {
        components: keyed.into_iter().map(|(_, c)| c).collect(),
    }
}

#[derive(Clone, Debug)]
struct Dsu {
    parent: Vec<usize>,
}

impl Dsu {
    fn new(n: usize) -> Self {
        Self {
            parent: (0..n).collect(),
        }
    }

    fn find(&self, x: usize) -> usize {
        std::iter::successors(Some(x), |&i| {
            let p = self.parent[i];
            (p != i).then_some(p)
        })
        .fold(x, |_, i| i)
    }

    fn union(self, a: usize, b: usize) -> Self {
        let ra = self.find(a);
        let rb = self.find(b);
        if ra == rb {
            return self;
        }
        let (lo, hi) = if ra < rb { (ra, rb) } else { (rb, ra) };
        let Self { mut parent } = self;
        parent[hi] = lo;
        Self { parent }
    }
}
