///|
pub(all) struct SpeciesPopulation {
  species : Species
  particles : Array[Particle]
} derive(Debug, ToJson)

///|
pub fn SpeciesPopulation::new(
  species : Species,
  particles : Array[Particle],
) -> SpeciesPopulation {
  { species, particles }
}

///|
pub fn SpeciesPopulation::count(population : SpeciesPopulation) -> Int {
  population.particles.length()
}

///|
pub fn SpeciesPopulation::charge(population : SpeciesPopulation) -> Double {
  total_particle_charge(population.particles)
}

///|
pub fn SpeciesPopulation::mass(population : SpeciesPopulation) -> Double {
  population.particles.fold(init=0.0, fn(acc, p) { acc + p.mass * p.weight })
}

///|
pub(all) struct MultiSpeciesState {
  grid : Grid1D
  populations : Array[SpeciesPopulation]
  charge_density : Array[Double]
  field : Field1D
  time : Double
} derive(Debug, ToJson)

///|
pub fn species_particles(
  species : Species,
  count : Int,
  spacing : Double,
  drift : Double,
) -> Array[Particle] {
  let n = if count < 0 { 0 } else { count }
  Array::makei(n, fn(i) {
    let sign = if i % 2 == 0 { 1.0 } else { -1.0 }
    particle_for_species(species, x=i.to_double() * spacing, v=drift * sign)
  })
}

///|
pub fn population_charge_density(
  grid : Grid1D,
  populations : ArrayView[SpeciesPopulation],
  kind : ShapeKind,
) -> Array[Double] {
  let density = zeros(grid.cells)
  for population in populations {
    let contribution = deposit_charge_with_shape(
      grid,
      population.particles,
      kind,
    )
    for i in 0.. MultiSpeciesState {
  let charge_density = population_charge_density(grid, populations, CIC)
  {
    grid,
    populations,
    charge_density,
    field: solve_periodic_field(grid, charge_density),
    time: 0.0,
  }
}

///|
pub fn MultiSpeciesState::particle_count(state : MultiSpeciesState) -> Int {
  state.populations.fold(init=0, fn(acc, population) {
    acc + population.count()
  })
}

///|
pub fn MultiSpeciesState::total_charge(state : MultiSpeciesState) -> Double {
  state.populations.fold(init=0.0, fn(acc, population) {
    acc + population.charge()
  })
}

///|
pub fn MultiSpeciesState::total_mass(state : MultiSpeciesState) -> Double {
  state.populations.fold(init=0.0, fn(acc, population) {
    acc + population.mass()
  })
}

///|
pub fn MultiSpeciesState::species_count(state : MultiSpeciesState) -> Int {
  state.populations.length()
}

///|
pub fn MultiSpeciesState::population(
  state : MultiSpeciesState,
  index : Int,
) -> SpeciesPopulation? {
  if index < 0 || index >= state.populations.length() {
    None
  } else {
    Some(state.populations[index])
  }
}

///|
pub fn MultiSpeciesState::number_density(
  state : MultiSpeciesState,
) -> Array[Double] {
  let density = zeros(state.grid.cells)
  for population in state.populations {
    let contribution = deposit_number_with_shape(
      state.grid,
      population.particles,
      CIC,
    )
    for i in 0.. Array[Double] {
  let current = zeros(state.grid.cells)
  for population in state.populations {
    let contribution = particle_current_density(
      state.grid,
      population.particles,
    )
    for i in 0.. MultiSpeciesState {
  let populations = state.populations.map(fn(population) {
    let particles = population.particles.map(fn(particle) {
      let electric = sample_linear(state.grid, state.field.electric, particle.x)
      let acceleration = particle.charge * electric / particle.mass
      let velocity = particle.v + acceleration * dt
      Particle::new(
        x=state.grid.wrap(particle.x + velocity * dt),
        v=velocity,
        weight=particle.weight,
        charge=particle.charge,
        mass=particle.mass,
      )
    })
    SpeciesPopulation::new(population.species, particles)
  })
  let density = population_charge_density(state.grid, populations, CIC)
  {
    grid: state.grid,
    populations,
    charge_density: density,
    field: solve_periodic_field(state.grid, density),
    time: state.time + dt,
  }
}

///|
pub fn MultiSpeciesState::run(
  state : MultiSpeciesState,
  dt : Double,
  steps : Int,
) -> MultiSpeciesState {
  let mut current = state
  let mut index = 0
  while index < steps {
    current = current.step(dt)
    index = index + 1
  }
  current
}

///|
pub fn MultiSpeciesState::kinetic_energy(state : MultiSpeciesState) -> Double {
  state.populations.fold(init=0.0, fn(acc, population) {
    acc + kinetic_energy(population.particles)
  })
}

///|
pub fn MultiSpeciesState::total_energy(state : MultiSpeciesState) -> Double {
  state.kinetic_energy() + field_energy(state.field)
}

///|
pub fn MultiSpeciesState::flatten(state : MultiSpeciesState) -> Array[Particle] {
  let output : Array[Particle] = []
  for population in state.populations {
    for particle in population.particles {
      output.push(particle)
    }
  }
  output
}

///|
pub fn population_names(state : MultiSpeciesState) -> Array[String] {
  state.populations.map(fn(population) { population.species.name })
}

///|
pub fn population_charge_summary(state : MultiSpeciesState) -> Array[Double] {
  state.populations.map(fn(population) { population.charge() })
}

///|
pub fn population_mass_summary(state : MultiSpeciesState) -> Array[Double] {
  state.populations.map(fn(population) { population.mass() })
}

///|
pub fn merge_populations(
  first : SpeciesPopulation,
  second : SpeciesPopulation,
) -> SpeciesPopulation {
  let particles = first.particles.copy()
  for particle in second.particles {
    particles.push(particle)
  }
  SpeciesPopulation::new(first.species, particles)
}