use nalgebra::{DMatrix, DVector, RowDVector};

use crate::particle_system::ParticleSystem;
use crate::algebra::{Scalar, Vector, N};
pub mod anchor;
pub mod slider;
pub mod beam;

const SPRING_CONSTANT: Scalar = 0.65;
const DAMPING_CONSTANT: Scalar = 0.85;

/// SIZE is always 3xN, but this operation can't be done at compile time yet:
/// "generic parameters may not be used in const operations"
pub trait Constraint {
    fn get_particles(&self) -> Vec<usize>;

    fn c(&self, q: &DVector<Scalar>) -> Scalar;

    fn jacobian_prev(&self) -> RowDVector<Scalar>;
    fn set_jacobian(&mut self, jacobian: RowDVector<Scalar>);

    fn partial_derivative(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
        let dq = 0.00001;
        let c_original = self.c(&q);
        let mut result = RowDVector::zeros(q.len());
        // The only non-zero components of derivative vector are
        // the particles that this constraint affects
        for particle_id in self.get_particles() {
            for i in 0..N {
                let index = particle_id * N + i;
                let mut q_partial = q.clone();
                q_partial[index] += dq;

                let c = self.c(&q_partial);

                result[index] = (c - c_original) / dq
            }
        }

        result
    }

    fn compute(
        &mut self,
        q: &DVector<Scalar>,
        q_prev: &DVector<Scalar>,
        dt: Scalar,
    ) -> (Scalar, Scalar, RowDVector<Scalar>, RowDVector<Scalar>) {
        let c = self.c(q);
        let c_prev = self.c(q_prev);

        let c_dot = (c - c_prev) / dt;

        let jacobian = self.partial_derivative(&q);

        let jacobian_dot = (jacobian.clone() - self.jacobian_prev()) / dt;
        self.set_jacobian(jacobian.clone());

        (c, c_dot, jacobian, jacobian_dot)
    }
}

impl ParticleSystem {
    pub fn enforce_constraints(&mut self, dt: Scalar) {
        if self.constraints.len() == 0 {
            return
        }

        let size = self.particles.len() * N;
        let mut q = DVector::zeros(size);
        let mut q_prev = DVector::zeros(size);

        for (p, particle) in self.particles.iter().enumerate() {
            for i in 0..N {
                q[p * N + i] = particle.position[i];
                q_prev[p * N + i] = particle.position[i] - particle.velocity[i] * dt;
            }
        }

        let q_dot = (q.clone() - q_prev.clone()) / dt;

        let mut c = DVector::zeros(self.constraints.len());
        let mut c_dot = DVector::zeros(self.constraints.len());
        let mut jacobian = DMatrix::<Scalar>::zeros(self.constraints.len(), size);
        let mut jacobian_dot = DMatrix::<Scalar>::zeros(self.constraints.len(), size);

        for (constraint_id, constraint) in self.constraints.iter_mut().enumerate() {
            let (value, derivative, j, jdot) = constraint.compute(&q, &q_prev, dt);
            jacobian.set_row(constraint_id, &j);
            jacobian_dot.set_row(constraint_id, &jdot);
            c[constraint_id] = value;
            c_dot[constraint_id] = derivative;
        }

        // println!("C {}\nC' {}", c, c_dot);
        // println!("J  = {}", jacobian.clone());
        // println!("J' = {}", jacobian_dot.clone());

        let mut w = DMatrix::zeros(size, size);
        for (i, particle) in self.particles.iter().enumerate() {
            for j in 0..N {
                *w.index_mut((i * N + j, i * N + j)) = 1.0 / particle.mass;
            }
        }

        let mut forces = DVector::zeros(size);
        for (p, particle) in self.particles.iter().enumerate() {
            for i in 0..N {
                forces[p * N + i] = particle.force[i];
            }
        }

        let lhs = jacobian.clone() * w.clone() * jacobian.transpose();
        // println!("LHS {}", lhs);

        let rhs = -(jacobian_dot * q_dot
            + jacobian.clone() * w * forces
            + SPRING_CONSTANT * c
            + DAMPING_CONSTANT * c_dot);

        // println!("RHS {}", rhs);

        match lhs.lu().solve(&rhs) {
            Some(lambda) => {
                // println!("Lambda = {}", lambda);
                let constraint_force = jacobian.transpose() * lambda;

                for i in 0..self.particles.len() {
                    let mut force = Vector::zeros();
                    for j in 0..N {
                        force[j] = constraint_force[i * N + j];
                    }

                    self.particles[i].apply_force(force);
                }
            }
            None => println!("Lambda not found"),
        }
    }
}