use nalgebra::{DVector, RowDVector};

use crate::algebra::{Point, Scalar, N};
use crate::particle_system::ParticleSystem;

use super::Constraint;

#[derive(Clone, Copy)]
enum BeamPoint {
    Static(Point),
    FromParticle(usize),
}

pub struct BeamConstraint {
    points: [BeamPoint; 2],
    pub length: Scalar,

    jacobian: RowDVector<Scalar>,
}

impl ParticleSystem {
    pub fn add_beam_constraint(&mut self, particle_ids: [usize; 2]) {
        let a = &self.particles[particle_ids[0]];
        let b = &self.particles[particle_ids[1]];

        self.constraints.push(Box::new(BeamConstraint {
            points: [
                BeamPoint::FromParticle(particle_ids[0]),
                BeamPoint::FromParticle(particle_ids[1]),
            ],
            length: (a.position - b.position).norm(),
            jacobian: RowDVector::zeros(self.particles.len() * N),
        }));
    }

    pub fn add_anchor_constraint(&mut self, particle_id: usize, point: Point) {
        let position = &self.particles[particle_id].position;

        self.constraints.push(Box::new(BeamConstraint {
            points: [
                BeamPoint::FromParticle(particle_id),
                BeamPoint::Static(point),
            ],
            length: (position - point).norm(),
            jacobian: RowDVector::zeros(self.particles.len() * N),
        }));
    }
}

impl Constraint for BeamConstraint {
    fn get_particles(&self) -> Vec<usize> {
        self.points.iter().filter_map(|p| match p {
            BeamPoint::FromParticle(id) => Some(*id),
            BeamPoint::Static(_) => None,
        }).collect()
    }

    fn c(&self, q: &DVector<Scalar>) -> Scalar {
        let [a, b] = self.points.map(|p| {
            match p {
                BeamPoint::Static(p) => p.coords,
                BeamPoint::FromParticle(id) => q.fixed_rows::<N>(id * N).into(),
            }
        });

        (a - b).norm() - self.length
    }

    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>) {
        self.jacobian = jacobian
    }

    fn jacobian_prev(&self) -> RowDVector<Scalar> {
        self.jacobian.clone()
    }
}