///|
pub(all) enum CollisionKind {
  Elastic
  Drag
  Langevin
  Thermalize
} derive(Eq, Debug, ToJson)

///|
pub(all) struct CollisionModel {
  kind : CollisionKind
  rate : Double
  background_velocity : Double
  thermal_speed : Double
} derive(Debug, ToJson)

///|
pub fn collision_model(
  kind : CollisionKind,
  rate : Double,
  background_velocity : Double,
  thermal_speed : Double,
) -> CollisionModel {
  {
    kind,
    rate: rate.max(0.0),
    background_velocity,
    thermal_speed: thermal_speed.max(0.0),
  }
}

///|
pub fn coulomb_logarithm(
  debye_length_m : Double,
  impact_parameter_m : Double,
) -> Double {
  if debye_length_m <= impact_parameter_m || impact_parameter_m <= 0.0 {
    0.0
  } else {
    approximate_log(debye_length_m / impact_parameter_m)
  }
}

///|
pub fn collision_frequency(
  density : Double,
  charge : Double,
  mass : Double,
  temperature : Double,
  log_lambda : Double,
) -> Double {
  if mass <= 0.0 || temperature <= 0.0 {
    0.0
  } else {
    density *
    charge *
    charge *
    log_lambda.abs() /
    (mass.sqrt() * @math.pow(temperature, 1.5).max(1.0e-30))
  }
}

///|
pub fn drag_factor(rate : Double, dt : Double) -> Double {
  @math.exp(-rate.max(0.0) * dt.max(0.0))
}

///|
pub fn apply_drag(
  particle : Particle,
  rate : Double,
  dt : Double,
  background_velocity : Double,
) -> Particle {
  let factor = drag_factor(rate, dt)
  let velocity = background_velocity +
    (particle.v - background_velocity) * factor
  Particle::new(
    x=particle.x + velocity * dt,
    v=velocity,
    weight=particle.weight,
    charge=particle.charge,
    mass=particle.mass,
  )
}

///|
pub fn apply_collision(
  particle : Particle,
  model : CollisionModel,
  dt : Double,
  noise : Double,
) -> Particle {
  match model.kind {
    Elastic => particle
    Drag => apply_drag(particle, model.rate, dt, model.background_velocity)
    Langevin => {
      let factor = drag_factor(model.rate, dt)
      let thermal = noise *
        model.thermal_speed *
        (1.0 - factor * factor).max(0.0).sqrt()
      let velocity = model.background_velocity +
        (particle.v - model.background_velocity) * factor +
        thermal
      Particle::new(
        x=particle.x + velocity * dt,
        v=velocity,
        weight=particle.weight,
        charge=particle.charge,
        mass=particle.mass,
      )
    }
    Thermalize => {
      let velocity = model.background_velocity + noise * model.thermal_speed
      Particle::new(
        x=particle.x + velocity * dt,
        v=velocity,
        weight=particle.weight,
        charge=particle.charge,
        mass=particle.mass,
      )
    }
  }
}

///|
pub fn apply_collisions(
  particles : ArrayView[Particle],
  model : CollisionModel,
  dt : Double,
  noise : ArrayView[Double],
) -> Array[Particle] {
  particles.map(fn(particle) {
    let index = if noise.length() == 0 {
      0
    } else {
      modulo_index(particles.length(), noise.length())
    }
    apply_collision(
      particle,
      model,
      dt,
      if noise.length() == 0 {
        0.0
      } else {
        noise[index]
      },
    )
  })
}

///|
pub fn relaxation_time(rate : Double) -> Double {
  if rate <= 0.0 {
    0.0
  } else {
    1.0 / rate
  }
}

///|
pub fn collision_probability(rate : Double, dt : Double) -> Double {
  1.0 - drag_factor(rate, dt)
}

///|
pub fn collision_energy_loss(particle : Particle, after : Particle) -> Double {
  (kinetic_energy([particle]) - kinetic_energy([after])).max(0.0)
}

///|
pub fn thermal_speed_from_temperature(
  temperature : Double,
  mass : Double,
) -> Double {
  thermal_velocity(temperature, mass_kg=mass)
}

///|
pub fn collision_temperature(
  particles : ArrayView[Particle],
  mass : Double,
) -> Double {
  velocity_variance(particles) * mass / boltzmann_constant
}

///|
pub fn collision_rate_for_species(
  species : Species,
  density : Double,
  temperature : Double,
  log_lambda : Double,
) -> Double {
  collision_frequency(
    density,
    species.charge.abs(),
    species.mass,
    temperature,
    log_lambda,
  )
}

///|
pub fn collision_summary(
  particles : ArrayView[Particle],
  model : CollisionModel,
  dt : Double,
) -> (Double, Double) {
  let after = apply_collisions(particles, model, dt, [])
  let before_energy = kinetic_energy(particles)
  let after_energy = kinetic_energy(after)
  (before_energy, after_energy)
}