use crate::particle_system::ParticleSystem;
use crate::algebra::{Point, Scalar, Vector, N};
use nalgebra::{Const, DVector, Dyn, Matrix, ViewStorage};

mod rk4;

/// A vector of concatenated position and velocity components of each particle
#[derive(Debug, Clone)]
pub struct PhaseSpace(DVector<Scalar>);
type ParticleView<'a> = Matrix<
    Scalar,
    Const<{ PhaseSpace::PARTICLE_DIM }>,
    Const<1>,
    ViewStorage<'a, Scalar, Const<{ PhaseSpace::PARTICLE_DIM }>, Const<1>, Const<1>, Dyn>,
>;

impl PhaseSpace {
    /// Each particle spans 2N elements in a vector
    /// first N for position, then N more for velocity
    const PARTICLE_DIM: usize = N * 2;

    pub fn new(particle_count: usize) -> Self {
        let dimension = particle_count * PhaseSpace::PARTICLE_DIM;
        Self(DVector::<Scalar>::zeros(dimension))
    }

    pub fn particle_view(&self, i: usize) -> ParticleView {
        self.0
            .fixed_rows::<{ PhaseSpace::PARTICLE_DIM }>(i * PhaseSpace::PARTICLE_DIM)
    }

    pub fn set_particle(&mut self, i: usize, position: Point, velocity: Vector) {
        let mut view = self
            .0
            .fixed_rows_mut::<{ PhaseSpace::PARTICLE_DIM }>(i * PhaseSpace::PARTICLE_DIM);
        for i in 0..N {
            view[i] = position[i];
            view[i + N] = velocity[i];
        }
    }
}

impl ParticleSystem {
    fn collect_phase_space(&self) -> PhaseSpace {
        let mut phase_space = PhaseSpace::new(self.particles.len());
        for (particle_index, particle) in self.particles.iter().enumerate() {
            phase_space.set_particle(particle_index, particle.position, particle.velocity);
        }
        phase_space
    }

    fn compute_derivative(&self) -> PhaseSpace {
        let mut phase_space = PhaseSpace::new(self.particles.len());
        for (particle_index, particle) in self.particles.iter().enumerate() {
            phase_space.set_particle(
                particle_index,
                particle.velocity.into(),
                particle.force / particle.mass,
            );
        }
        phase_space
    }

    fn scatter_phase_space(&mut self, phase_space: &PhaseSpace) {
        for (particle_index, particle) in &mut self.particles.iter_mut().enumerate() {
            let view = phase_space.particle_view(particle_index);
            for i in 0..N {
                particle.position[i] = view[i];
                particle.velocity[i] = view[i + N];
            }
        }
    }
}

pub trait Solver {
    fn step(&mut self, dt: Scalar);
}

#[cfg(test)]
mod tests {
    use super::{ParticleSystem, PhaseSpace, Point, Scalar, Solver, Vector};
    use crate::particle_system::Particle;

    #[test]
    fn test_collect_phase_space() {
        let system = ParticleSystem {
            particles: vec![Particle::new(Point::new(2.0, 3.0, 0.0), 1.0)],
            constraints: vec![],
            forces: vec![],
            t: 0.0,
        };
        let phase_space = system.collect_phase_space();

        assert!(
            !phase_space.0.is_empty(),
            "Phase space has to contain non-zero values"
        );
    }

    #[test]
    fn test_scatter_phase_space() {
        let mut phase_space = PhaseSpace::new(2);
        phase_space.set_particle(1, Point::new(5.0, 7.0, 0.0), Vector::x());

        let mut system = ParticleSystem {
            particles: vec![
                Particle::new(Point::new(0.0, 0.0, 0.0), 1.0),
                Particle::new(Point::new(0.0, 0.0, 0.0), 1.0),
            ],
            constraints: vec![],
            forces: vec![],
            t: 0.0,
        };

        system.scatter_phase_space(&phase_space);

        assert!(
            !system.particles[1].velocity.is_empty(),
            "Velocity has to be set"
        );
        assert!(
            !system.particles[1].position.is_empty(),
            "Position has to be set"
        );
    }

    fn simulate_falling_ball(
        fall_time: Scalar,
        dt: Scalar,
        mass: Scalar,
    ) -> (Vector, Vector, ParticleSystem) {
        let gravity = -9.8 * Vector::y();

        let mut system = ParticleSystem {
            particles: vec![Particle::new(Point::origin(), mass)],
            constraints: vec![],
            forces: vec![],
            t: 0.0,
        };

        let iterations = (fall_time / dt) as usize;

        for _ in 0..iterations {
            for particle in &mut system.particles {
                particle.reset_force();
                particle.apply_force(gravity * particle.mass);
            }
            system.step(dt);
        }

        let expected_velocity = gravity * fall_time; // vt
        let expected_position = gravity * fall_time * fall_time / 2.0; // at^2 / 2

        (
            system.particles[0].position.coords - expected_position,
            system.particles[0].velocity - expected_velocity,
            system,
        )
    }

    #[test]
    fn ball_should_fall() {
        let (position_error, velocity_error, _) = simulate_falling_ball(10.0, 0.01, 3.0);
        assert!(
            position_error.norm() < 0.01,
            "Position error is too high: {}",
            position_error,
        );
        assert!(
            velocity_error.norm() < 0.01,
            "Velocity error is too high: {}",
            velocity_error,
        );
    }

    #[test]
    fn freefall_different_masses() {
        let (_, _, system1) = simulate_falling_ball(10.0, 0.01, 2.0);
        let (_, _, system2) = simulate_falling_ball(10.0, 0.01, 10.0);

        let diff = system1.particles[0].position - system2.particles[0].position;
        assert!(
            diff.norm() < 0.01,
            "Points with different masses should fall with the same speed, diff: {}",
            diff
        );
    }
}