use bone_types::{ParameterIndex, ResidualIndex};

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct PointHandle {
    pub x: ParameterIndex,
    pub y: ParameterIndex,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct LineHandle {
    pub a: PointHandle,
    pub b: PointHandle,
}

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum CurveRadius {
    Explicit {
        center: PointHandle,
        radius: ParameterIndex,
    },
    FromSpoke {
        center: PointHandle,
        spoke: PointHandle,
    },
}

impl CurveRadius {
    #[must_use]
    pub fn center(self) -> PointHandle {
        match self {
            Self::Explicit { center, .. } | Self::FromSpoke { center, .. } => center,
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum Residual {
    Pin {
        param: ParameterIndex,
        target: f64,
    },
    Horizontal(LineHandle),
    Vertical(LineHandle),
    Parallel(LineHandle, LineHandle),
    Perpendicular(LineHandle, LineHandle),
    TangentLineCurve {
        line: LineHandle,
        curve: CurveRadius,
    },
    TangentCurveCurve {
        a: CurveRadius,
        b: CurveRadius,
    },
    CoincidentPointPoint(PointHandle, PointHandle),
    CoincidentPointLine {
        point: PointHandle,
        line: LineHandle,
    },
    CoincidentPointCurve {
        point: PointHandle,
        curve: CurveRadius,
    },
    MidpointPointLine {
        point: PointHandle,
        line: LineHandle,
    },
    EqualLength(LineHandle, LineHandle),
    EqualRadius(CurveRadius, CurveRadius),
    LinearDistance {
        a: PointHandle,
        b: PointHandle,
        value_mm: f64,
    },
    AngularBetweenLines {
        a: LineHandle,
        b: LineHandle,
        angle_rad: f64,
    },
    RadiusCurve {
        curve: CurveRadius,
        value_mm: f64,
    },
    Symmetric {
        a: PointHandle,
        b: PointHandle,
        axis: LineHandle,
    },
}

pub type Triplet = (ResidualIndex, ParameterIndex, f64);

impl Residual {
    #[must_use]
    pub fn rows(&self) -> usize {
        match self {
            Self::CoincidentPointPoint(..)
            | Self::MidpointPointLine { .. }
            | Self::Symmetric { .. } => 2,
            _ => 1,
        }
    }

    #[must_use]
    pub fn parameters(&self) -> Vec<ParameterIndex> {
        match *self {
            Self::Pin { param, .. } => vec![param],
            Self::Horizontal(l) | Self::Vertical(l) => line_params(l),
            Self::Parallel(l1, l2) | Self::Perpendicular(l1, l2) | Self::EqualLength(l1, l2) => {
                line_params(l1).into_iter().chain(line_params(l2)).collect()
            }
            Self::TangentLineCurve { line, curve } => line_params(line)
                .into_iter()
                .chain(curve_params(curve))
                .collect(),
            Self::TangentCurveCurve { a, b } | Self::EqualRadius(a, b) => {
                curve_params(a).into_iter().chain(curve_params(b)).collect()
            }
            Self::CoincidentPointPoint(p, q) => {
                point_params(p).into_iter().chain(point_params(q)).collect()
            }
            Self::CoincidentPointLine { point, line } | Self::MidpointPointLine { point, line } => {
                point_params(point)
                    .into_iter()
                    .chain(line_params(line))
                    .collect()
            }
            Self::CoincidentPointCurve { point, curve } => point_params(point)
                .into_iter()
                .chain(curve_params(curve))
                .collect(),
            Self::LinearDistance { a, b, .. } => {
                point_params(a).into_iter().chain(point_params(b)).collect()
            }
            Self::AngularBetweenLines { a, b, .. } => {
                line_params(a).into_iter().chain(line_params(b)).collect()
            }
            Self::RadiusCurve { curve, .. } => curve_params(curve),
            Self::Symmetric { a, b, axis } => point_params(a)
                .into_iter()
                .chain(point_params(b))
                .chain(line_params(axis))
                .collect(),
        }
    }

    pub fn evaluate(&self, params: &[f64], out: &mut [f64]) {
        match *self {
            Self::Pin { param, target } => {
                out[0] = params[param.as_usize()] - target;
            }
            Self::Horizontal(l) => out[0] = get(params, l.a.y) - get(params, l.b.y),
            Self::Vertical(l) => out[0] = get(params, l.a.x) - get(params, l.b.x),
            Self::Parallel(l1, l2) => {
                let (d1x, d1y) = delta(params, l1);
                let (d2x, d2y) = delta(params, l2);
                out[0] = d1x * d2y - d1y * d2x;
            }
            Self::Perpendicular(l1, l2) => {
                let (d1x, d1y) = delta(params, l1);
                let (d2x, d2y) = delta(params, l2);
                out[0] = d1x * d2x + d1y * d2y;
            }
            Self::TangentLineCurve { line, curve } => {
                let (dx, dy) = delta(params, line);
                let cx = get(params, curve.center().x);
                let cy = get(params, curve.center().y);
                let ex = cx - get(params, line.a.x);
                let ey = cy - get(params, line.a.y);
                let cross = ex * dy - ey * dx;
                let len2 = dx * dx + dy * dy;
                let r2 = radius_squared(params, curve);
                out[0] = cross * cross - r2 * len2;
            }
            Self::TangentCurveCurve { a, b } => {
                let ax = get(params, a.center().x);
                let ay = get(params, a.center().y);
                let bx = get(params, b.center().x);
                let by = get(params, b.center().y);
                let ra = radius_squared(params, a).sqrt();
                let rb = radius_squared(params, b).sqrt();
                let u = ax - bx;
                let v = ay - by;
                let rsum = ra + rb;
                out[0] = u * u + v * v - rsum * rsum;
            }
            Self::CoincidentPointPoint(p, q) => {
                out[0] = get(params, p.x) - get(params, q.x);
                out[1] = get(params, p.y) - get(params, q.y);
            }
            Self::CoincidentPointLine { point, line } => {
                let (dx, dy) = delta(params, line);
                let ex = get(params, point.x) - get(params, line.a.x);
                let ey = get(params, point.y) - get(params, line.a.y);
                out[0] = ex * dy - ey * dx;
            }
            Self::MidpointPointLine { point, line } => {
                out[0] =
                    get(params, point.x) - 0.5 * (get(params, line.a.x) + get(params, line.b.x));
                out[1] =
                    get(params, point.y) - 0.5 * (get(params, line.a.y) + get(params, line.b.y));
            }
            Self::CoincidentPointCurve { point, curve } => {
                let cx = get(params, curve.center().x);
                let cy = get(params, curve.center().y);
                let ex = get(params, point.x) - cx;
                let ey = get(params, point.y) - cy;
                let r2 = radius_squared(params, curve);
                out[0] = ex * ex + ey * ey - r2;
            }
            Self::EqualLength(l1, l2) => {
                let (d1x, d1y) = delta(params, l1);
                let (d2x, d2y) = delta(params, l2);
                out[0] = d1x * d1x + d1y * d1y - (d2x * d2x + d2y * d2y);
            }
            Self::EqualRadius(a, b) => {
                out[0] = radius_squared(params, a) - radius_squared(params, b);
            }
            Self::LinearDistance { a, b, value_mm } => {
                let dx = get(params, b.x) - get(params, a.x);
                let dy = get(params, b.y) - get(params, a.y);
                out[0] = dx * dx + dy * dy - value_mm * value_mm;
            }
            Self::AngularBetweenLines { a, b, angle_rad } => {
                let (d1x, d1y) = delta(params, a);
                let (d2x, d2y) = delta(params, b);
                let dot = d1x * d2x + d1y * d2y;
                let cross = d1x * d2y - d1y * d2x;
                let (s, c) = angle_rad.sin_cos();
                out[0] = s * dot - c * cross;
            }
            Self::RadiusCurve { curve, value_mm } => {
                out[0] = radius_squared(params, curve) - value_mm * value_mm;
            }
            Self::Symmetric { a, b, axis } => {
                let (ax, ay) = (get(params, a.x), get(params, a.y));
                let (bx, by) = (get(params, b.x), get(params, b.y));
                let (dx, dy) = delta(params, axis);
                let lax = get(params, axis.a.x);
                let lay = get(params, axis.a.y);
                out[0] = (bx - ax) * dx + (by - ay) * dy;
                let mx = 0.5 * (ax + bx) - lax;
                let my = 0.5 * (ay + by) - lay;
                out[1] = mx * dy - my * dx;
            }
        }
    }

    pub fn jacobian(&self, params: &[f64], row: ResidualIndex, sink: &mut Vec<Triplet>) {
        match *self {
            Self::Pin { param, .. } => sink.push((row, param, 1.0)),
            Self::Horizontal(l) => {
                sink.push((row, l.a.y, 1.0));
                sink.push((row, l.b.y, -1.0));
            }
            Self::Vertical(l) => {
                sink.push((row, l.a.x, 1.0));
                sink.push((row, l.b.x, -1.0));
            }
            Self::Parallel(l1, l2) => jacobian_parallel(params, row, l1, l2, sink),
            Self::Perpendicular(l1, l2) => jacobian_perpendicular(params, row, l1, l2, sink),
            Self::TangentLineCurve { line, curve } => {
                tangent_line_curve_jacobian(params, row, line, curve, sink);
            }
            Self::TangentCurveCurve { a, b } => {
                tangent_curve_curve_jacobian(params, row, a, b, sink);
            }
            Self::CoincidentPointPoint(p, q) => {
                let row2 = row.next();
                sink.push((row, p.x, 1.0));
                sink.push((row, q.x, -1.0));
                sink.push((row2, p.y, 1.0));
                sink.push((row2, q.y, -1.0));
            }
            Self::CoincidentPointLine { point, line } => {
                jacobian_coincident_point_line(params, row, point, line, sink);
            }
            Self::MidpointPointLine { point, line } => {
                let row2 = row.next();
                sink.push((row, point.x, 1.0));
                sink.push((row, line.a.x, -0.5));
                sink.push((row, line.b.x, -0.5));
                sink.push((row2, point.y, 1.0));
                sink.push((row2, line.a.y, -0.5));
                sink.push((row2, line.b.y, -0.5));
            }
            Self::CoincidentPointCurve { point, curve } => {
                jacobian_coincident_point_curve(params, row, point, curve, sink);
            }
            Self::EqualLength(l1, l2) => jacobian_equal_length(params, row, l1, l2, sink),
            Self::EqualRadius(a, b) => {
                add_radius_squared_jacobian(params, row, a, sink);
                subtract_radius_squared_jacobian(params, row, b, sink);
            }
            Self::LinearDistance { a, b, .. } => {
                jacobian_linear_distance(params, row, a, b, sink);
            }
            Self::AngularBetweenLines { a, b, angle_rad } => {
                jacobian_angular_between_lines(params, row, a, b, angle_rad, sink);
            }
            Self::RadiusCurve { curve, .. } => {
                add_radius_squared_jacobian(params, row, curve, sink);
            }
            Self::Symmetric { a, b, axis } => {
                jacobian_symmetric(params, row, a, b, axis, sink);
            }
        }
    }
}

fn jacobian_symmetric(
    params: &[f64],
    row: ResidualIndex,
    a: PointHandle,
    b: PointHandle,
    axis: LineHandle,
    sink: &mut Vec<Triplet>,
) {
    let (dx, dy) = delta(params, axis);
    let (ax, ay) = (get(params, a.x), get(params, a.y));
    let (bx, by) = (get(params, b.x), get(params, b.y));
    let lax = get(params, axis.a.x);
    let lay = get(params, axis.a.y);
    let mx = 0.5 * (ax + bx) - lax;
    let my = 0.5 * (ay + by) - lay;
    let bxa_x = bx - ax;
    let bxa_y = by - ay;
    sink.push((row, a.x, -dx));
    sink.push((row, a.y, -dy));
    sink.push((row, b.x, dx));
    sink.push((row, b.y, dy));
    sink.push((row, axis.a.x, -bxa_x));
    sink.push((row, axis.a.y, -bxa_y));
    sink.push((row, axis.b.x, bxa_x));
    sink.push((row, axis.b.y, bxa_y));
    let row2 = row.next();
    sink.push((row2, a.x, 0.5 * dy));
    sink.push((row2, a.y, -0.5 * dx));
    sink.push((row2, b.x, 0.5 * dy));
    sink.push((row2, b.y, -0.5 * dx));
    sink.push((row2, axis.a.x, -dy + my));
    sink.push((row2, axis.a.y, -mx + dx));
    sink.push((row2, axis.b.x, -my));
    sink.push((row2, axis.b.y, mx));
}

fn jacobian_parallel(
    params: &[f64],
    row: ResidualIndex,
    l1: LineHandle,
    l2: LineHandle,
    sink: &mut Vec<Triplet>,
) {
    let (d1x, d1y) = delta(params, l1);
    let (d2x, d2y) = delta(params, l2);
    sink.push((row, l1.a.x, -d2y));
    sink.push((row, l1.b.x, d2y));
    sink.push((row, l1.a.y, d2x));
    sink.push((row, l1.b.y, -d2x));
    sink.push((row, l2.a.x, d1y));
    sink.push((row, l2.b.x, -d1y));
    sink.push((row, l2.a.y, -d1x));
    sink.push((row, l2.b.y, d1x));
}

fn jacobian_perpendicular(
    params: &[f64],
    row: ResidualIndex,
    l1: LineHandle,
    l2: LineHandle,
    sink: &mut Vec<Triplet>,
) {
    let (d1x, d1y) = delta(params, l1);
    let (d2x, d2y) = delta(params, l2);
    sink.push((row, l1.a.x, -d2x));
    sink.push((row, l1.b.x, d2x));
    sink.push((row, l1.a.y, -d2y));
    sink.push((row, l1.b.y, d2y));
    sink.push((row, l2.a.x, -d1x));
    sink.push((row, l2.b.x, d1x));
    sink.push((row, l2.a.y, -d1y));
    sink.push((row, l2.b.y, d1y));
}

fn jacobian_coincident_point_line(
    params: &[f64],
    row: ResidualIndex,
    point: PointHandle,
    line: LineHandle,
    sink: &mut Vec<Triplet>,
) {
    let (dx, dy) = delta(params, line);
    let ex = get(params, point.x) - get(params, line.a.x);
    let ey = get(params, point.y) - get(params, line.a.y);
    sink.push((row, point.x, dy));
    sink.push((row, point.y, -dx));
    sink.push((row, line.a.x, ey - dy));
    sink.push((row, line.a.y, dx - ex));
    sink.push((row, line.b.x, -ey));
    sink.push((row, line.b.y, ex));
}

fn jacobian_coincident_point_curve(
    params: &[f64],
    row: ResidualIndex,
    point: PointHandle,
    curve: CurveRadius,
    sink: &mut Vec<Triplet>,
) {
    let center = curve.center();
    let ex = get(params, point.x) - get(params, center.x);
    let ey = get(params, point.y) - get(params, center.y);
    sink.push((row, point.x, 2.0 * ex));
    sink.push((row, point.y, 2.0 * ey));
    sink.push((row, center.x, -2.0 * ex));
    sink.push((row, center.y, -2.0 * ey));
    subtract_radius_squared_jacobian(params, row, curve, sink);
}

fn jacobian_equal_length(
    params: &[f64],
    row: ResidualIndex,
    l1: LineHandle,
    l2: LineHandle,
    sink: &mut Vec<Triplet>,
) {
    let (d1x, d1y) = delta(params, l1);
    let (d2x, d2y) = delta(params, l2);
    sink.push((row, l1.a.x, -2.0 * d1x));
    sink.push((row, l1.b.x, 2.0 * d1x));
    sink.push((row, l1.a.y, -2.0 * d1y));
    sink.push((row, l1.b.y, 2.0 * d1y));
    sink.push((row, l2.a.x, 2.0 * d2x));
    sink.push((row, l2.b.x, -2.0 * d2x));
    sink.push((row, l2.a.y, 2.0 * d2y));
    sink.push((row, l2.b.y, -2.0 * d2y));
}

fn jacobian_linear_distance(
    params: &[f64],
    row: ResidualIndex,
    a: PointHandle,
    b: PointHandle,
    sink: &mut Vec<Triplet>,
) {
    let dx = get(params, b.x) - get(params, a.x);
    let dy = get(params, b.y) - get(params, a.y);
    sink.push((row, a.x, -2.0 * dx));
    sink.push((row, a.y, -2.0 * dy));
    sink.push((row, b.x, 2.0 * dx));
    sink.push((row, b.y, 2.0 * dy));
}

fn jacobian_angular_between_lines(
    params: &[f64],
    row: ResidualIndex,
    a: LineHandle,
    b: LineHandle,
    angle_rad: f64,
    sink: &mut Vec<Triplet>,
) {
    let (d1x, d1y) = delta(params, a);
    let (d2x, d2y) = delta(params, b);
    let (sin_t, cos_t) = angle_rad.sin_cos();
    let d1_partial_x = sin_t * d2x - cos_t * d2y;
    let d1_partial_y = sin_t * d2y + cos_t * d2x;
    let d2_partial_x = sin_t * d1x + cos_t * d1y;
    let d2_partial_y = sin_t * d1y - cos_t * d1x;
    sink.push((row, a.a.x, -d1_partial_x));
    sink.push((row, a.b.x, d1_partial_x));
    sink.push((row, a.a.y, -d1_partial_y));
    sink.push((row, a.b.y, d1_partial_y));
    sink.push((row, b.a.x, -d2_partial_x));
    sink.push((row, b.b.x, d2_partial_x));
    sink.push((row, b.a.y, -d2_partial_y));
    sink.push((row, b.b.y, d2_partial_y));
}

fn line_params(l: LineHandle) -> Vec<ParameterIndex> {
    vec![l.a.x, l.a.y, l.b.x, l.b.y]
}

fn point_params(p: PointHandle) -> Vec<ParameterIndex> {
    vec![p.x, p.y]
}

fn curve_params(c: CurveRadius) -> Vec<ParameterIndex> {
    match c {
        CurveRadius::Explicit { center, radius } => vec![center.x, center.y, radius],
        CurveRadius::FromSpoke { center, spoke } => vec![center.x, center.y, spoke.x, spoke.y],
    }
}

fn get(params: &[f64], idx: ParameterIndex) -> f64 {
    params[idx.as_usize()]
}

fn delta(params: &[f64], line: LineHandle) -> (f64, f64) {
    (
        get(params, line.b.x) - get(params, line.a.x),
        get(params, line.b.y) - get(params, line.a.y),
    )
}

fn radius_squared(params: &[f64], curve: CurveRadius) -> f64 {
    match curve {
        CurveRadius::Explicit { radius, .. } => {
            let r = get(params, radius);
            r * r
        }
        CurveRadius::FromSpoke { center, spoke } => {
            let dx = get(params, spoke.x) - get(params, center.x);
            let dy = get(params, spoke.y) - get(params, center.y);
            dx * dx + dy * dy
        }
    }
}

fn add_radius_squared_jacobian(
    params: &[f64],
    row: ResidualIndex,
    curve: CurveRadius,
    sink: &mut Vec<Triplet>,
) {
    match curve {
        CurveRadius::Explicit { radius, .. } => {
            sink.push((row, radius, 2.0 * get(params, radius)));
        }
        CurveRadius::FromSpoke { center, spoke } => {
            let dx = get(params, spoke.x) - get(params, center.x);
            let dy = get(params, spoke.y) - get(params, center.y);
            sink.push((row, spoke.x, 2.0 * dx));
            sink.push((row, spoke.y, 2.0 * dy));
            sink.push((row, center.x, -2.0 * dx));
            sink.push((row, center.y, -2.0 * dy));
        }
    }
}

fn subtract_radius_squared_jacobian(
    params: &[f64],
    row: ResidualIndex,
    curve: CurveRadius,
    sink: &mut Vec<Triplet>,
) {
    match curve {
        CurveRadius::Explicit { radius, .. } => {
            sink.push((row, radius, -2.0 * get(params, radius)));
        }
        CurveRadius::FromSpoke { center, spoke } => {
            let dx = get(params, spoke.x) - get(params, center.x);
            let dy = get(params, spoke.y) - get(params, center.y);
            sink.push((row, spoke.x, -2.0 * dx));
            sink.push((row, spoke.y, -2.0 * dy));
            sink.push((row, center.x, 2.0 * dx));
            sink.push((row, center.y, 2.0 * dy));
        }
    }
}

fn tangent_line_curve_jacobian(
    params: &[f64],
    row: ResidualIndex,
    line: LineHandle,
    curve: CurveRadius,
    sink: &mut Vec<Triplet>,
) {
    let (dx, dy) = delta(params, line);
    let cx = get(params, curve.center().x);
    let cy = get(params, curve.center().y);
    let ex = cx - get(params, line.a.x);
    let ey = cy - get(params, line.a.y);
    let cross = ex * dy - ey * dx;
    let len2 = dx * dx + dy * dy;
    let r2 = radius_squared(params, curve);
    let two_cross = 2.0 * cross;

    sink.push((
        row,
        line.a.x,
        two_cross * (cy - get(params, line.b.y)) + r2 * 2.0 * dx,
    ));
    sink.push((
        row,
        line.a.y,
        two_cross * (get(params, line.b.x) - cx) + r2 * 2.0 * dy,
    ));
    sink.push((
        row,
        line.b.x,
        two_cross * (get(params, line.a.y) - cy) - r2 * 2.0 * dx,
    ));
    sink.push((
        row,
        line.b.y,
        two_cross * (cx - get(params, line.a.x)) - r2 * 2.0 * dy,
    ));
    sink.push((row, curve.center().x, two_cross * dy));
    sink.push((row, curve.center().y, -two_cross * dx));

    let neg_len2 = -len2;
    match curve {
        CurveRadius::Explicit { radius, .. } => {
            sink.push((row, radius, neg_len2 * 2.0 * get(params, radius)));
        }
        CurveRadius::FromSpoke { center, spoke } => {
            let sdx = get(params, spoke.x) - get(params, center.x);
            let sdy = get(params, spoke.y) - get(params, center.y);
            sink.push((row, spoke.x, neg_len2 * 2.0 * sdx));
            sink.push((row, spoke.y, neg_len2 * 2.0 * sdy));
            sink.push((row, center.x, neg_len2 * -2.0 * sdx));
            sink.push((row, center.y, neg_len2 * -2.0 * sdy));
        }
    }
}

fn tangent_curve_curve_jacobian(
    params: &[f64],
    row: ResidualIndex,
    a: CurveRadius,
    b: CurveRadius,
    sink: &mut Vec<Triplet>,
) {
    let ra = radius_squared(params, a).sqrt();
    let rb = radius_squared(params, b).sqrt();
    let rsum = ra + rb;
    let u = get(params, a.center().x) - get(params, b.center().x);
    let v = get(params, a.center().y) - get(params, b.center().y);

    sink.push((row, a.center().x, 2.0 * u));
    sink.push((row, a.center().y, 2.0 * v));
    sink.push((row, b.center().x, -2.0 * u));
    sink.push((row, b.center().y, -2.0 * v));

    let coef = -2.0 * rsum;
    push_radius_linear_jacobian(params, row, a, coef, sink);
    push_radius_linear_jacobian(params, row, b, coef, sink);
}

fn push_radius_linear_jacobian(
    params: &[f64],
    row: ResidualIndex,
    curve: CurveRadius,
    coef: f64,
    sink: &mut Vec<Triplet>,
) {
    if coef == 0.0 {
        return;
    }
    match curve {
        CurveRadius::Explicit { radius, .. } => {
            sink.push((row, radius, coef));
        }
        CurveRadius::FromSpoke { center, spoke } => {
            let dx = get(params, spoke.x) - get(params, center.x);
            let dy = get(params, spoke.y) - get(params, center.y);
            let len = (dx * dx + dy * dy).sqrt();
            if len == 0.0 {
                return;
            }
            let k = coef / len;
            sink.push((row, spoke.x, k * dx));
            sink.push((row, spoke.y, k * dy));
            sink.push((row, center.x, -k * dx));
            sink.push((row, center.y, -k * dy));
        }
    }
}
