use nalgebra::{DVector, RowDVector};

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

use super::Constraint;

pub struct BeamConstraint {
    pub particle_ids: [usize; 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 {
            particle_ids,
            length: (a.position - b.position).norm(),
            jacobian: RowDVector::zeros(self.particles.len() * N),
        }));
    }
}

impl Constraint for BeamConstraint {
    fn get_particles(&self) -> Vec<usize> {
        Vec::from(self.particle_ids)
    }

    fn c(&self, q: &DVector<Scalar>) -> Scalar {
        let a = q.fixed_rows::<N>(self.particle_ids[0] * N);
        let b = q.fixed_rows::<N>(self.particle_ids[1] * N);

        (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()
    }
}