use nalgebra::{Const, DVector, Dyn, Matrix, Point as PointBase, SVector, ViewStorage};

const N: usize = 2;
pub type Scalar = f32;

pub type Vector = SVector<Scalar, N>;
pub type Point = PointBase<Scalar, N>;

#[derive(Debug)]
pub struct Particle {
    pub mass: Scalar,
    pub position: Point,
    pub velocity: Vector,

    /// Force accumulator
    pub force: Vector,
}

impl Particle {
    pub fn new(position: Point, mass: Scalar) -> Self {
        Self {
            mass,
            position,
            velocity: Vector::zeros(),
            force: Vector::zeros(),
        }
    }

    pub fn apply_force(&mut self, force: Vector) {
        self.force += force;
    }
    pub fn reset_force(&mut self) {
        self.force = Vector::zeros()
    }
}

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

#[derive(Debug)]
pub struct ParticleSystem {
    pub particles: Vec<Particle>,

    /// Simulation clock
    pub t: Scalar,
}

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 fn euler_step(&mut self, dt: Scalar) {
        let derivative = self.compute_derivative();
        let mut state = self.collect_phase_space();

        state.0 += derivative.0 * dt;
        self.scatter_phase_space(&state);

        self.t += dt;
    }
}

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

    #[test]
    fn test_collect_phase_space() {
        let system = ParticleSystem {
            particles: vec![Particle::new(Point::new(2.0, 3.0), 1.0)],
            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), Vector::x());

        let mut system = ParticleSystem {
            particles: vec![
                Particle::new(Point::new(0.0, 0.0), 1.0),
                Particle::new(Point::new(0.0, 0.0), 1.0),
            ],
            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) -> (Vector, Vector) {
        let gravity = -9.8 * Vector::y();

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

        let iterations = (fall_time / dt) as usize;

        for _ in 0..iterations {
            for particle in &mut system.particles {
                particle.reset_force();
            }

            for particle in &mut system.particles {
                particle.apply_force(gravity);
            }

            system.euler_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,
        )
    }

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