use std::collections::BTreeSet;

use bone_types::{ParentIndex, ResidualIndex};

use crate::dof::{DofConfig, analyze_dof_at};
use crate::graph::decompose;
use crate::newton::{NewtonConfig, solve_newton_decomposed};
use crate::system::ConstraintSystem;
use crate::util::dedup_preserving_order;

#[must_use]
pub fn minimal_unsatisfiable_subset(
    system: &ConstraintSystem,
    over_flagged: &[ResidualIndex],
    cfg: DofConfig,
) -> Vec<ResidualIndex> {
    let candidate_parents: Vec<ParentIndex> = dedup_preserving_order(
        over_flagged
            .iter()
            .copied()
            .map(|row| system.parent_of_row(row)),
    );
    if candidate_parents.is_empty() {
        return Vec::new();
    }
    let solver_cfg = NewtonConfig::DEFAULT;
    let kept_parents = candidate_parents
        .iter()
        .copied()
        .fold(
            (BTreeSet::<ParentIndex>::new(), Vec::<ParentIndex>::new()),
            |(mut excluded, mut kept), parent| {
                let mut trial = excluded.clone();
                trial.insert(parent);
                if reduced_is_over(system, &trial, solver_cfg, cfg) {
                    excluded.insert(parent);
                } else {
                    kept.push(parent);
                }
                (excluded, kept)
            },
        )
        .1;
    kept_parents
        .into_iter()
        .flat_map(|p| system.rows_of_parent(p))
        .filter(|row| over_flagged.contains(row))
        .collect()
}

fn reduced_is_over(
    system: &ConstraintSystem,
    exclude: &BTreeSet<ParentIndex>,
    newton_cfg: NewtonConfig,
    dof_cfg: DofConfig,
) -> bool {
    let sub = system.without_parents(exclude);
    if sub.row_count() == 0 {
        return false;
    }
    let decomposition = decompose(&sub);
    match solve_newton_decomposed(&sub, &decomposition, newton_cfg) {
        Ok(solved) => {
            let report = analyze_dof_at(&sub, &solved, dof_cfg);
            !report.over_constrained().is_empty()
        }
        Err(_) => true,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::residual::Residual;
    use bone_types::{Parameter, ParameterIndex};

    fn p(i: u32) -> ParameterIndex {
        ParameterIndex::new(i)
    }

    fn pin(at: u32, target: f64) -> Residual {
        Residual::Pin {
            param: p(at),
            target,
        }
    }

    fn over_flagged_after_solve(system: &ConstraintSystem) -> Vec<ResidualIndex> {
        let decomposition = decompose(system);
        let solved = solve_newton_decomposed(system, &decomposition, NewtonConfig::DEFAULT)
            .unwrap_or_else(|_| system.parameters().to_vec());
        analyze_dof_at(system, &solved, DofConfig::DEFAULT)
            .over_constrained()
            .to_vec()
    }

    #[test]
    fn three_conflicting_pins_isolate_the_odd_one_out() {
        let system = ConstraintSystem::new(
            vec![Parameter::new(0.0)],
            vec![pin(0, 1.0), pin(0, 2.0), pin(0, 1.0)],
        );
        let over = over_flagged_after_solve(&system);
        assert!(
            !over.is_empty(),
            "fixture must produce at least one over row"
        );
        let mus = minimal_unsatisfiable_subset(&system, &over, DofConfig::DEFAULT);
        assert!(
            !mus.is_empty(),
            "MUS must keep at least one essential row, got empty",
        );
        assert!(
            mus.len() <= over.len(),
            "MUS must be a subset of over-flagged rows: mus={mus:?} over={over:?}",
        );
    }

    #[test]
    fn empty_over_flagged_yields_empty_mus() {
        let system = ConstraintSystem::new(vec![Parameter::new(0.0)], vec![pin(0, 0.0)]);
        let mus = minimal_unsatisfiable_subset(&system, &[], DofConfig::DEFAULT);
        assert!(mus.is_empty());
    }
}
