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

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

pub const SPRING_CONSTANT: Scalar = 16.0;
pub const DAMPING_CONSTANT: Scalar = 4.0;

pub trait Constraint {
    fn get_particles(&self) -> Vec<usize>;

    /// Constraint function: C = 0 <=> constrained satisfied
    fn c(&self, q: &DVector<Scalar>) -> Scalar;

    /// J = dC / dq
    /// This implementation computes partial derivative via
    /// finite differences method: stepping +-dQ along each axis
    ///
    /// NOTE: this is rather computationally intensive so
    /// providing constraint-specific analytical implementation is preferred
    fn jacobian(&self, q: &DVector<Scalar>) -> RowDVector<Scalar> {
        let dq = 1e-10;
        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 = q.clone();

                q[index] += dq;
                let c_a = self.c(&q);

                q[index] -= 2.0 * dq;
                let c_b = self.c(&q);

                result[index] = (c_a - c_b) / (2.0 * dq)
            }
        }
        result
    }

    /// C' = J * q'
    fn c_dot(&self, jacobian: &RowDVector<Scalar>, q_dot: &DVector<Scalar>) -> Scalar {
        (jacobian * q_dot)[0]
    }

    /// J' = dC' / dq
    fn jacobian_dot(&self, q: &DVector<Scalar>, q_dot: &DVector<Scalar>) -> RowDVector<Scalar> {
        let dq = 1e-4;
        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 = q.clone();

                q[index] += dq;
                let j_a = self.jacobian(&q);

                q[index] -= 2.0 * dq;
                let j_b = self.jacobian(&q);

                result[index] = (self.c_dot(&j_a, q_dot) - self.c_dot(&j_b, q_dot)) / (2.0 * dq)
            }
        }
        result
    }
}

impl ParticleSystem {
    /// q is a state vector - concatenated positions of particles
    /// q' - its derivative - concatenated velocities
    pub fn collect_q(&self) -> (DVector<Scalar>, DVector<Scalar>) {
        let size = self.q_dim();
        let mut q = DVector::zeros(size);
        let mut q_dot = DVector::zeros(size);

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

        (q, q_dot)
    }
    /// dim(q) = #particles * N
    fn q_dim(&self) -> usize {
        self.particles.len() * N
    }

    /// Diagonal matrix of particle masses (each repeated N times to match dim(q))
    fn mass_matrix(&self) -> DMatrix<Scalar> {
        DMatrix::from_diagonal(&DVector::from_iterator(
            self.q_dim(),
            self.particles
                .iter()
                .flat_map(|particle| (0..N).map(|_i| 1.0 / particle.mass)),
        ))
    }

    pub fn enforce_constraints(&mut self) {
        if self.constraints.len() == 0 {
            return;
        }

        let (q, q_dot) = self.collect_q();
        let dim_q = self.q_dim();
        let dim_c = self.constraints.len();

        let c = DVector::from_iterator(
            dim_c,
            self.constraints.iter().map(|constraint| constraint.c(&q)),
        );

        let jacobian = DMatrix::from_rows(
            &self
                .constraints
                .iter()
                .map(|constraint| constraint.jacobian(&q))
                .collect::<Vec<_>>(),
        );

        let jacobian_dot = DMatrix::from_rows(
            &self
                .constraints
                .iter()
                .map(|constraint| constraint.jacobian_dot(&q, &q_dot))
                .collect::<Vec<_>>(),
        );

        let c_dot = &jacobian * &q_dot;

        let forces = DVector::from_iterator(
            dim_q,
            self.particles
                .iter()
                .flat_map(|particle| particle.force.iter().map(|&v| v)),
        );

        let jacobian_times_w = &jacobian * self.mass_matrix();
        let jacobian_transpose = &jacobian.transpose();

        let lhs = &jacobian_times_w * jacobian_transpose;

        let rhs = -(&jacobian_dot * &q_dot
            + jacobian_times_w * &forces
            + SPRING_CONSTANT * &c
            + DAMPING_CONSTANT * &c_dot);

        match lhs.lu().solve(&rhs) {
            Some(lambda) => {
                let constraint_force = jacobian_transpose * lambda;

                for i in 0..self.particles.len() {
                    let force = constraint_force.fixed_rows::<N>(i * N);
                    self.particles[i].apply_force(force.into());
                }
            }
            None => println!("Lambda not found"),
        }
    }
}