program two_stream
  implicit none

  ! Domain settings
  integer, parameter          :: nx            = 400
  double precision, parameter :: domain_length = 1.0d0
  double precision, parameter :: dx            = domain_length / nx
  double precision, parameter :: inv_dx        = 1.0d0/dx

  ! Simulation settings
  double precision, parameter :: end_time         = 1.0d0
  double precision, parameter :: dt_output        = end_time * 1.0d-2

  ! Settings for the beam
  ! Mean velocity of the first beam
  double precision, parameter :: beam1_v0       = 4.0d0
  ! Mean velocity of the first beam
  double precision, parameter :: beam2_v0       = -4.0d0
  ! How large the thermal velocities of the beams are
  double precision, parameter :: beam_vthermal = 1.0d0
  ! The plasma density in the beam
  double precision, parameter :: plasma_density   = 1.0d0

  ! Particle properties
  double precision, parameter :: particle_charge  = -1.6022d-19
  double precision, parameter :: particle_mass    = 9.10938189d-31

  ! Physical constants (SI units)
  double precision, parameter :: boltzmann_const  = 1.3806503d-23
  double precision, parameter :: epsilon_0        = 8.8541878176d-12

  ! Number of particles per stream
  integer, parameter :: n_per_stream = 100*1000

  ! Total number of particles
  integer, parameter :: n_particles  = 2 * n_per_stream

  ! Particle positions (m)
  double precision :: x(n_particles)
  ! Particle velocities (m/s)
  double precision :: v(n_particles)
  ! Particle accelerations (m/s^2)
  double precision :: a(n_particles)

  ! The time step in the simulation
  double precision :: dt

  ! The inverse of the 'transverse area' of the 1D model. This is used to make
  ! sure we obtain the desired plasma density.
  double precision :: inverse_area

  ! Electric field at cell faces, so field_at_faces(1) is centered at x = 0.0 dx
  double precision :: field_at_faces(-1:nx+1)

  ! Particle density
  double precision :: density(-1:nx+1)

  ! Other simulation parameters
  double precision :: time     = 0.0d0
  integer          :: n_output = 0
  integer          :: n, my_unit
  double precision :: debye_length
  double precision :: plasma_frequency

  debye_length = sqrt(epsilon_0 * particle_mass * beam_vthermal**2 / &
       (plasma_density * particle_charge**2))

  plasma_frequency = sqrt(plasma_density * particle_charge**2 / &
       (particle_mass * epsilon_0))

  ! Use a fixed time step
  dt = 0.25d0 * dx / (max(abs(beam1_v0), abs(beam2_v0)) + 3 * beam_vthermal)

  write(*, "(A30,E14.6)") "Domain length:", domain_length
  write(*, "(A30,I14)") "Domain Nx:", nx
  write(*, "(A30,E14.6)") "Domain dx:", dx
  write(*, "(A30,E14.6)") "Time step:", dt
  write(*, "(A30,E14.6)") "Beam 1 mean velocity:", beam1_v0
  write(*, "(A30,E14.6)") "Beam 2 mean velocity:", beam2_v0
  write(*, "(A30,E14.6)") "Beam thermal velocity:", beam_vthermal
  write(*, "(A30,E14.6)") "Beam density:", plasma_density
  write(*, "(A30,E14.6)") "Debye length:", debye_length
  write(*, "(A30,E14.6)") "Plasma frequency:", plasma_frequency
  write(*, "(A30,E14.6)") "Inverse plasma frequency:", 1/plasma_frequency
  write(*, "(A30,E14.6)") "End time:", end_time

  ! Make sure we have the right plasma density, by adjusting the tranvserse area
  ! of the 1D domain
  inverse_area = (plasma_density * domain_length) / n_particles

  call set_initial_conditions()
  call particles_to_density()
  call compute_electric_field()
  call set_acceleration()

  do while(time < end_time)
     if (time >= n_output * dt_output) then
        n_output = n_output + 1
        call write_output(n_output)
        write(*, "(A,I4,A,E11.4,F6.1,A)") "Written output ", &
             n_output, " at t =", time, 100*time/end_time, "%"
     end if

     ! Advance in time with leapfrog scheme. The scheme is written such that
     ! position and velocity can be defined at the same time.
     do n = 1, n_particles
        x(n) = x(n) + v(n) * 0.5d0 * dt
     end do

     call particles_to_density()
     call compute_electric_field()
     call set_acceleration()

     do n = 1, n_particles
        v(n) = v(n) + a(n) * dt
     end do

     do n = 1, n_particles
        x(n) = x(n) + v(n) * 0.5d0 * dt
     end do

     time = time + dt
  end do

contains

  ! Set initial conditions
  subroutine set_initial_conditions()
    integer          :: n
    double precision :: vthermal(2)

    do n = 1, n_per_stream
       x(n)              = (n-0.5d0)/n_per_stream * domain_length
       x(n_per_stream+n) = x(n)
       vthermal          = two_normals()
       v(n)              = beam1_v0 + vthermal(1)
       v(n_per_stream+n) = beam2_v0 + vthermal(2)
    end do
  end subroutine set_initial_conditions

  subroutine particles_to_density()
    integer :: n

    density(:) = 0.0d0

    do n = 1, n_particles
       x(n) = modulo(x(n), domain_length)
       call add_to_dens_cic(inverse_area, x(n))
    end do

    ! Account for periodic domain
    density(1)    = density(1) + density(nx+1)
    density(nx)   = density(nx) + density(0)
    density(0)    = 0.0d0
    density(nx+1) = 0.0d0
  end subroutine particles_to_density

  subroutine add_to_dens_cic(amount, x)
    double precision, intent(in) :: amount, x
    integer              :: low_ix
    double precision             :: tmp, val

    ! Do linear interpolation, note that dens(1) is defined at 0.5 * dx
    tmp    = x * inv_dx - 0.5
    low_ix = floor(tmp) + 1
    ! Coefficient for upper point between 0 and 1
    tmp    = low_ix - tmp
    val    = amount * inv_dx

    if (low_ix < 0 .or. low_ix > nx) then
       error stop "a particle is outside the computational domain"
    end if

    density(low_ix)   = density(low_ix) + (1.0d0 - tmp) * val
    density(low_ix+1) = density(low_ix+1) + tmp * val
  end subroutine add_to_dens_cic

  ! Compute the electric field and store it
  subroutine compute_electric_field()
    double precision :: pot_diff
    double precision :: my_density(1:nx)
    integer          :: n

    my_density = (density(1:nx) - plasma_density) * particle_charge / epsilon_0

    ! Start with this guess
    field_at_faces(1) = 0.0d0
    do n = 2, nx+1
       field_at_faces(n) = field_at_faces(n-1) + dx * my_density(n-1)
    end do

    ! Compute total potential difference in the domain, with a weight of 0.5 for
    ! the boundary field values
    pot_diff = -dx * (sum(field_at_faces(2:nx)) + &
         0.5d0 * (field_at_faces(1) + field_at_faces(nx+1)))

    ! Correct field so that the potential difference is zero (since the domain
    ! is periodic)
    field_at_faces(:) = field_at_faces(:) + pot_diff / domain_length

  end subroutine compute_electric_field

  subroutine set_acceleration()
    integer          :: n
    double precision :: field, factor

    factor = particle_charge / particle_mass

    do n = 1, n_particles
       field = get_field_at(x(n))
       a(n) = field * factor
    end do
  end subroutine set_acceleration

  ! Get the electric field at a position in the domain (useful for the particle
  ! model)
  function get_field_at(x) result(field)
    double precision, intent(in) :: x
    double precision             :: field, tmp
    integer                      :: low_ix

    ! Do linear interpolation, note that field_at_faces(1) is defined at 0.0
    tmp    = x * inv_dx
    low_ix = floor(tmp) + 1
    ! Coefficient for upper point between 0 and 1
    tmp    = low_ix - tmp

    if (low_ix < 1 .or. low_ix > nx) then
       print *, "get_field_at invalid position x = ", x
       error stop
    end if

    field = (1.0d0 - tmp) * field_at_faces(low_ix) + tmp * field_at_faces(low_ix+1)
  end function get_field_at

  function two_normals() result(rands)
    double precision :: rands(2), sum_sq

    do
       call random_number(rands)
       rands = 2.0d0 * rands - 1.0d0
       sum_sq = sum(rands**2)
       if (sum_sq < 1.0d0 .and. sum_sq > 0.0d0) exit
    end do
    rands = rands * sqrt(-2 * log(sum_sq) / sum_sq)
  end function two_normals

  ! Write output to text files
  subroutine write_output(n_output)
    integer, intent(in)         :: n_output
    integer, parameter          :: plot_nx = 400
    character(len=*), parameter :: fmt_string = "(400E16.8)"
    double precision            :: vmax
    double precision            :: plot_invdv
    double precision            :: plot_invdx
    double precision            :: xv_density(plot_nx, plot_nx)
    character(len=200)          :: filename
    integer                     :: i, j

    vmax       = 8.0d0 ! Maximal velocity for the output
    plot_invdv = plot_nx/(2*vmax)
    plot_invdx = plot_nx/domain_length
    xv_density(:, :) = 0.0d0

    write(filename, "(A,I0.4,A)") "twostream_xv_", n_output, ".txt"
    open(newunit=my_unit, file=trim(filename), action="write")
    do n = 1, n_particles
       i = floor(x(n) * plot_invdx) + 1
       j = floor((v(n) + vmax) * plot_invdv) + 1
       if (i > 0 .and. i <= plot_nx .and. j > 0 .and. j < plot_nx) then
          xv_density(i, j) = xv_density(i, j) + inverse_area * plot_invdx * plot_invdv
       end if
    end do
    do n = 1, plot_nx
       write(my_unit, fmt_string) xv_density(:, n)
    end do
    close(my_unit)

    ! Enable this to write particle data
    ! write(filename, "(A,I0.4,A)") "twostream_part_", n_output, ".txt"
    ! open(newunit=my_unit, file=trim(filename), action="write")
    ! do n = 1, n_particles
    !    write(my_unit, "(3E16.8)") x(n), v(n), a(n)
    ! end do
    ! close(my_unit)

    write(filename, "(A,I0.4,A)") "twostream_", n_output, ".txt"
    open(newunit=my_unit, file=trim(filename), action="write")
    do n = 1, nx
       write(my_unit, "(3E16.8)") (n-0.5d0) * dx, density(n), &
            0.5d0 * (field_at_faces(n) + field_at_faces(n+1))
    end do
    close(my_unit)
  end subroutine write_output

end program two_stream
