use std::collections::{BTreeMap, BTreeSet};

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

use crate::graph::Component;
use crate::residual::{CurveRadius, LineHandle, PointHandle, Residual};

#[derive(Clone, Debug, PartialEq)]
pub struct ConstraintSystem {
    parameters: Vec<Parameter>,
    residuals: Vec<Residual>,
}

impl ConstraintSystem {
    #[must_use]
    pub fn new(parameters: Vec<Parameter>, residuals: Vec<Residual>) -> Self {
        Self {
            parameters,
            residuals,
        }
    }

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

    #[must_use]
    pub fn residuals(&self) -> &[Residual] {
        &self.residuals
    }

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

    #[must_use]
    pub fn row_count(&self) -> usize {
        self.residuals.iter().map(Residual::rows).sum()
    }

    #[must_use]
    pub fn row_offsets(&self) -> Vec<ResidualIndex> {
        self.residuals
            .iter()
            .scan(0u32, |acc, r| {
                let start = ResidualIndex::new(*acc);
                let Ok(n) = u32::try_from(r.rows()) else {
                    unreachable!("Residual::rows is bounded to a small constant")
                };
                let Some(next) = acc.checked_add(n) else {
                    unreachable!("row count fits in u32 by construction of ResidualIndex")
                };
                *acc = next;
                Some(start)
            })
            .collect()
    }

    #[must_use]
    pub fn parameter_at(&self, idx: ParameterIndex) -> Parameter {
        self.parameters[idx.as_usize()]
    }

    #[must_use]
    pub fn residual_at_parent(&self, parent: ParentIndex) -> &Residual {
        &self.residuals[parent.as_usize()]
    }

    #[must_use]
    pub fn with_extra_residuals(&self, extra: Vec<Residual>) -> Self {
        let mut residuals = self.residuals.clone();
        residuals.extend(extra);
        Self {
            parameters: self.parameters.clone(),
            residuals,
        }
    }

    #[must_use]
    pub fn parent_of_row(&self, row: ResidualIndex) -> ParentIndex {
        let target = row.as_usize();
        let found = self
            .residuals
            .iter()
            .scan(0usize, |acc, r| {
                let start = *acc;
                *acc += r.rows();
                Some((start, *acc))
            })
            .position(|(begin, end)| target >= begin && target < end);
        let Some(idx) = found else {
            unreachable!("ResidualIndex {row} is past the residual row count")
        };
        let Ok(iv) = u32::try_from(idx) else {
            unreachable!("parent count fits in u32 via ResidualIndex")
        };
        ParentIndex::new(iv)
    }

    #[must_use]
    pub fn rows_of_parent(&self, parent: ParentIndex) -> Vec<ResidualIndex> {
        let idx = parent.as_usize();
        let Some(residual) = self.residuals.get(idx) else {
            unreachable!("ParentIndex {parent} out of range for this ConstraintSystem")
        };
        let start_usize: usize = self.residuals[..idx].iter().map(Residual::rows).sum();
        let Ok(start) = u32::try_from(start_usize) else {
            unreachable!("row count fits in u32 via ResidualIndex")
        };
        (0..residual.rows())
            .map(|o| {
                let Ok(ov) = u32::try_from(o) else {
                    unreachable!("residual row count is a small constant")
                };
                ResidualIndex::new(start + ov)
            })
            .collect()
    }

    #[must_use]
    pub fn without_parents(&self, exclude: &BTreeSet<ParentIndex>) -> Self {
        let mask: BTreeSet<usize> = exclude.iter().copied().map(ParentIndex::as_usize).collect();
        let residuals = self
            .residuals
            .iter()
            .enumerate()
            .filter(|(i, _)| !mask.contains(i))
            .map(|(_, r)| r.clone())
            .collect();
        Self {
            parameters: self.parameters.clone(),
            residuals,
        }
    }

    #[must_use]
    pub fn subsystem(&self, component: &Component) -> Subsystem {
        let orig_params: &[ParameterIndex] = component.parameters();
        let remap: BTreeMap<ParameterIndex, ParameterIndex> = orig_params
            .iter()
            .copied()
            .enumerate()
            .map(|(i, orig)| {
                let Ok(iv) = u32::try_from(i) else {
                    unreachable!("component parameter count fits in u32 by construction")
                };
                (orig, ParameterIndex::new(iv))
            })
            .collect();
        let parameters: Vec<Parameter> = orig_params
            .iter()
            .map(|p| self.parameters[p.as_usize()])
            .collect();
        let residuals: Vec<Residual> = component
            .residual_parents()
            .iter()
            .map(|&parent| remap_residual(self.residual_at_parent(parent), &remap))
            .collect();
        Subsystem {
            system: ConstraintSystem::new(parameters, residuals),
            param_map: orig_params.to_vec(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct Subsystem {
    system: ConstraintSystem,
    param_map: Vec<ParameterIndex>,
}

impl Subsystem {
    #[must_use]
    pub fn system(&self) -> &ConstraintSystem {
        &self.system
    }

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

fn remap_residual(r: &Residual, m: &BTreeMap<ParameterIndex, ParameterIndex>) -> Residual {
    let idx = |p: ParameterIndex| -> ParameterIndex {
        let Some(&mapped) = m.get(&p) else {
            unreachable!("residual parameter missing from component remap (graph bug)")
        };
        mapped
    };
    let point = |p: PointHandle| PointHandle {
        x: idx(p.x),
        y: idx(p.y),
    };
    let line = |l: LineHandle| LineHandle {
        a: point(l.a),
        b: point(l.b),
    };
    let curve = |c: CurveRadius| match c {
        CurveRadius::Explicit { center, radius } => CurveRadius::Explicit {
            center: point(center),
            radius: idx(radius),
        },
        CurveRadius::FromSpoke { center, spoke } => CurveRadius::FromSpoke {
            center: point(center),
            spoke: point(spoke),
        },
    };
    match *r {
        Residual::Pin { param, target } => Residual::Pin {
            param: idx(param),
            target,
        },
        Residual::Horizontal(l) => Residual::Horizontal(line(l)),
        Residual::Vertical(l) => Residual::Vertical(line(l)),
        Residual::Parallel(a, b) => Residual::Parallel(line(a), line(b)),
        Residual::Perpendicular(a, b) => Residual::Perpendicular(line(a), line(b)),
        Residual::TangentLineCurve { line: l, curve: c } => Residual::TangentLineCurve {
            line: line(l),
            curve: curve(c),
        },
        Residual::TangentCurveCurve { a, b } => Residual::TangentCurveCurve {
            a: curve(a),
            b: curve(b),
        },
        Residual::CoincidentPointPoint(p, q) => Residual::CoincidentPointPoint(point(p), point(q)),
        Residual::CoincidentPointLine { point: p, line: l } => Residual::CoincidentPointLine {
            point: point(p),
            line: line(l),
        },
        Residual::CoincidentPointCurve { point: p, curve: c } => Residual::CoincidentPointCurve {
            point: point(p),
            curve: curve(c),
        },
        Residual::MidpointPointLine { point: p, line: l } => Residual::MidpointPointLine {
            point: point(p),
            line: line(l),
        },
        Residual::EqualLength(a, b) => Residual::EqualLength(line(a), line(b)),
        Residual::EqualRadius(a, b) => Residual::EqualRadius(curve(a), curve(b)),
        Residual::LinearDistance { a, b, value_mm } => Residual::LinearDistance {
            a: point(a),
            b: point(b),
            value_mm,
        },
        Residual::AngularBetweenLines { a, b, angle_rad } => Residual::AngularBetweenLines {
            a: line(a),
            b: line(b),
            angle_rad,
        },
        Residual::RadiusCurve { curve: c, value_mm } => Residual::RadiusCurve {
            curve: curve(c),
            value_mm,
        },
    }
}
