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,
}

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

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

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 jacobian(&self, q: &DVector<Scalar>) -> RowDVector<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(),
            }
        });
        let mut result = RowDVector::zeros(q.len());
        match (a - b).try_normalize(0.0) {
            Some(normal) => {
                for i in 0..N {
                    match self.points[0] {
                        BeamPoint::FromParticle(id) => {
                            result[id * N + i] = normal[i];
                        }
                        _ => {}
                    }
                    match self.points[1] {
                        BeamPoint::FromParticle(id) => {
                            result[id * N + i] = -normal[i];
                        }
                        _ => {}
                    }
                }
            }
            None => ()
        };

        result
    }
}