// Just position and velocity along one dimension. We're assuming the particle has a mass of 1.
#[derive(Clone)]
struct State {
  x: f64,
  v: f64,
}


trait Integrator {
  fn step(&mut self, state: State, sim: &Simulation, dt: f64) -> State;
}

trait Simulation {
  fn acc(&self, state: &State) -> f64;
  fn energy(&self, state: &State) -> f64;
  fn exact(&self, time: f64) -> f64;
}


struct HarmonicOscillator {
  k: f64,
  start_state: State
}

impl Simulation for HarmonicOscillator {
  fn acc(&self, state: &State) -> f64 {
      // https://en.wikipedia.org/wiki/Hooke's_law#For_linear_springs
    return -self.k * state.x;
  }

  fn energy(&self, state: &State) -> f64 {
      // Potential + kinetic energy
    return 0.5 * self.k * state.x * state.x + 0.5 * state.v * state.v;
  }

  fn exact(&self, time: f64) -> f64 {
      // See https://en.wikipedia.org/wiki/Simple_harmonic_motion#Dynamics for the corresponding equations.
    let start_state = &self.start_state;
    let omega = self.k.sqrt();
    let magnitude = (start_state.x * start_state.x + (start_state.v / omega) * (start_state.v / omega)).sqrt();
    let phase_offset = 
        // Prevent divide by 0
      if start_state.v == 0.0 {
        3.14159265358979 / 2.0
      } else {
        (start_state.x * omega / start_state.v).atan()
      };


    magnitude * (omega * time + phase_offset).sin()
  }
}


// https://en.wikipedia.org/wiki/Euler_method
#[allow(dead_code)]
struct ExplicitEuler {}
impl Integrator for ExplicitEuler {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a = sim.acc(&state);
    state.x += state.v * dt;
    state.v += a * dt;

    state
  }
}


// https://en.wikipedia.org/wiki/Semi-implicit_Euler_method
#[allow(dead_code)]
struct SemiImplicitEuler {}
impl Integrator for SemiImplicitEuler {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a = sim.acc(&state);
    state.v += a * dt;
    state.x += state.v * dt;

    state
  }
}


// https://en.wikipedia.org/wiki/Verlet_integration#Velocity_Verlet
#[allow(dead_code)]
struct VelocityVerlet {}
impl Integrator for VelocityVerlet {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a = sim.acc(&state);
    state.x = state.x + state.v * dt + 0.5 * a * dt * dt;
    let new_a = sim.acc(&state);
    state.v = state.v + ((a + new_a) / 2.0) * dt;

    state
  }
}

// Integration method given in "Building a Better Jump" - a simplified version of Velocity Verlet that doesn't evaulate acceleration at the updated position.
// http://www.mathforgameprogrammers.com/gdc2016/GDC2016_Pittman_Kyle_BuildingABetterJump.pdf, https://www.youtube.com/watch?v=hG9SzQxaCm8 for a recording of the talk.
#[allow(dead_code)]
struct BABJ {}
impl Integrator for BABJ {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a = sim.acc(&state);
    state.x = state.x + state.v * dt + 0.5 * a * dt * dt;
    state.v = state.v + a * dt;

    state
  }
}


// https://en.wikipedia.org/wiki/Beeman%27s_algorithm
// I don't recommend using my implementation of this method! In my testing it gains energy over time, so I may have misinterpreted the equations. I wasn't able to find the issue with it though..
#[allow(dead_code)]
struct Beeman {
  prev_state: State
}
impl Integrator for Beeman {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a_prev = sim.acc(&self.prev_state);
    let a_cur = sim.acc(&state);

    state.x += state.v * dt + (2.0 / 3.0) * a_cur * dt * dt - (1.0 / 6.0) * a_prev * dt * dt;

      // predicted velocity, in case acceleration calculation depends on velocity
    let original_v = state.v;
    state.v += (3.0 / 2.0) * a_cur * dt - 0.5 * a_prev * dt;

    let a_next = sim.acc(&state);

    state.v = original_v + (5.0 / 12.0) * a_next * dt + (2.0 / 3.0) * a_cur * dt - (1.0 / 12.0) * a_prev * dt;

    self.prev_state = state.clone();

    state
  }
}

struct Derivative {
  dx: f64,
  dv: f64
}

// https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods
// This implementation is a almost direct translation of the one given in https://gafferongames.com/post/integration_basics/
struct RK4 {}
impl RK4 {
  fn eval(state: &State, sim: &Simulation, dt: f64, derivative: &Derivative) -> Derivative {
    let advanced_state = State {
      x: state.x + derivative.dx * dt,
      v: state.v + derivative.dv * dt,
    };

    Derivative {
      dx: advanced_state.v,
      dv: sim.acc(&advanced_state)
    }
  }
}

impl Integrator for RK4 {
  fn step(&mut self, mut state: State, sim: &Simulation, dt: f64) -> State {
    let a = RK4::eval(&state, sim, dt, &Derivative { dx: 0.0, dv: 0.0});
    let b = RK4::eval(&state, sim, dt * 0.5, &a);
    let c = RK4::eval(&state, sim, dt * 0.5, &b);
    let d = RK4::eval(&state, sim, dt, &c);

    let dxdt = (1.0 / 6.0) * (a.dx + 2.0 * (b.dx + c.dx) + d.dx);
    let dvdt = (1.0 / 6.0) * (a.dv + 2.0 * (b.dv + c.dv) + d.dv);

    state.x += dxdt * dt;
    state.v += dvdt * dt;

    state
  }
}



fn main() {
    // pick your delta time
  let dt = 1.0/64.0;

  let total_time = 128.0;
  let steps = (total_time / dt) as usize;

  let mut state = State {
    x: 50.0,
    v: 0.0
  };

  let sim = HarmonicOscillator {
    k: 10.0, // spring stiffness
    start_state: state.clone(),
  };

    // Uncomment the line for whatever integration method you want to use
  let mut integrator = ExplicitEuler {};
  // let mut integrator = SemiImplicitEuler {};
  // let mut integrator = VelocityVerlet {};
  // let mut integrator = BABJ {};
  // let mut integrator = Beeman { prev_state: state.clone() };
  // let mut integrator = RK4 {};

  print_header();
  for i in 0..steps {
    state = integrator.step(state, &sim, dt);
    log(((i+1) as f64) * dt, &state, &sim);
  }
}

fn print_header() {
  println!("step,x,exact,energy");
}

fn log(time: f64, state: &State, sim: &Simulation) {
  let energy = sim.energy(state);
  let exact_x = sim.exact(time);
  println!("{},{},{},{}", time, state.x, exact_x, energy);
}