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

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

use crate::graph::Component;
use crate::residual::{CurveRadius3, Ellipse3Handles, Line3Handle, Point3Handle, 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_param(m: &BTreeMap<ParameterIndex, ParameterIndex>, p: ParameterIndex) -> ParameterIndex {
    let Some(&mapped) = m.get(&p) else {
        unreachable!("residual parameter missing from component remap (graph bug)")
    };
    mapped
}

fn remap_point3(m: &BTreeMap<ParameterIndex, ParameterIndex>, p: Point3Handle) -> Point3Handle {
    Point3Handle {
        x: remap_param(m, p.x),
        y: remap_param(m, p.y),
        z: remap_param(m, p.z),
    }
}

fn remap_line3(m: &BTreeMap<ParameterIndex, ParameterIndex>, l: Line3Handle) -> Line3Handle {
    Line3Handle {
        a: remap_point3(m, l.a),
        b: remap_point3(m, l.b),
    }
}

fn remap_curve3(m: &BTreeMap<ParameterIndex, ParameterIndex>, c: CurveRadius3) -> CurveRadius3 {
    match c {
        CurveRadius3::Explicit { center, radius } => CurveRadius3::Explicit {
            center: remap_point3(m, center),
            radius: remap_param(m, radius),
        },
        CurveRadius3::FromSpoke { center, spoke } => CurveRadius3::FromSpoke {
            center: remap_point3(m, center),
            spoke: remap_point3(m, spoke),
        },
    }
}

fn remap_residual(r: &Residual, m: &BTreeMap<ParameterIndex, ParameterIndex>) -> Residual {
    match *r {
        Residual::Pin { param, target } => Residual::Pin {
            param: remap_param(m, param),
            target,
        },
        _ => remap_residual3(r, m),
    }
}

fn remap_residual3(r: &Residual, m: &BTreeMap<ParameterIndex, ParameterIndex>) -> Residual {
    match *r {
        Residual::Coincident3(p, q) => {
            Residual::Coincident3(remap_point3(m, p), remap_point3(m, q))
        }
        Residual::CoincidentPointCurve3 { point: p, curve: c } => Residual::CoincidentPointCurve3 {
            point: remap_point3(m, p),
            curve: remap_curve3(m, c),
        },
        Residual::CoincidentPointEllipse3(e) => {
            Residual::CoincidentPointEllipse3(Ellipse3Handles {
                point: remap_point3(m, e.point),
                center: remap_point3(m, e.center),
                major: remap_point3(m, e.major),
                minor: remap_point3(m, e.minor),
            })
        }
        Residual::InPlane3 {
            point: pt,
            origin,
            normal,
        } => Residual::InPlane3 {
            point: remap_point3(m, pt),
            origin: remap_point3(m, origin),
            normal,
        },
        Residual::Distance3 { a, b, value_mm } => Residual::Distance3 {
            a: remap_point3(m, a),
            b: remap_point3(m, b),
            value_mm,
        },
        Residual::OnPlane {
            point: pt,
            normal,
            offset_mm,
        } => Residual::OnPlane {
            point: remap_point3(m, pt),
            normal,
            offset_mm,
        },
        Residual::AlongAxis { line: l, basis } => Residual::AlongAxis {
            line: remap_line3(m, l),
            basis,
        },
        Residual::Perpendicular3(a, b) => {
            Residual::Perpendicular3(remap_line3(m, a), remap_line3(m, b))
        }
        Residual::Parallel3(a, b) => Residual::Parallel3(remap_line3(m, a), remap_line3(m, b)),
        Residual::Midpoint3 { point: p, line: l } => Residual::Midpoint3 {
            point: remap_point3(m, p),
            line: remap_line3(m, l),
        },
        Residual::PointOnLine3 { point: p, line: l } => Residual::PointOnLine3 {
            point: remap_point3(m, p),
            line: remap_line3(m, l),
        },
        Residual::AngularBetweenLines3 { a, b, angle_rad } => Residual::AngularBetweenLines3 {
            a: remap_line3(m, a),
            b: remap_line3(m, b),
            angle_rad,
        },
        Residual::EqualDistance3 { center, a, b } => Residual::EqualDistance3 {
            center: remap_point3(m, center),
            a: remap_point3(m, a),
            b: remap_point3(m, b),
        },
        Residual::OrthogonalSpokes3 { center, a, b } => Residual::OrthogonalSpokes3 {
            center: remap_point3(m, center),
            a: remap_point3(m, a),
            b: remap_point3(m, b),
        },
        Residual::RadiusCurve3 { curve, value_mm } => Residual::RadiusCurve3 {
            curve: remap_curve3(m, curve),
            value_mm,
        },
        Residual::EqualRadius3(a, b) => {
            Residual::EqualRadius3(remap_curve3(m, a), remap_curve3(m, b))
        }
        Residual::EqualLength3(a, b) => {
            Residual::EqualLength3(remap_line3(m, a), remap_line3(m, b))
        }
        Residual::TangentLineCurve3 { line: l, curve } => Residual::TangentLineCurve3 {
            line: remap_line3(m, l),
            curve: remap_curve3(m, curve),
        },
        Residual::TangentCurveCurve3 { a, b, internal } => Residual::TangentCurveCurve3 {
            a: remap_curve3(m, a),
            b: remap_curve3(m, b),
            internal,
        },
        Residual::Symmetric3 { a, b, axis } => Residual::Symmetric3 {
            a: remap_point3(m, a),
            b: remap_point3(m, b),
            axis: remap_line3(m, axis),
        },
        Residual::Pin { .. } => {
            unreachable!("remap_residual3 is only reached for 3D residual variants")
        }
    }
}
