program ion_emitter

use bmad
use random_mod
use ion_emitter_mod
use particle_species_mod

implicit none

type (lat_struct) lat
type (coord_struct), allocatable :: orb(:)
type (random_state_struct) randstate
type (ion_emitter_param_struct) param
character(100) :: in_file, latname, outname, formato, gas_name
character(5) :: setc, getc
real(rp) :: setsigmacut, getsigmacut
real(rp) :: max_time, sspacing
real(rp), parameter :: boltz_const = 1.38064852D-23
real(rp), parameter :: boltz_const_ev = 8.6173324D-5
real(rp), parameter :: room_temp = 2.98D2
real(rp) :: ion_mass, boltzparam, gas_pressure, gas_density, kt
real(rp) :: col_cross, tau, n_electrons, t_step1, ion_charge1, kt_ev
real(rp) :: x, y, s, theta, phi, cpx, cpy, cpz, centile, sq
real(rp) :: center_x, center_y, hspacing, local_size
real(rp) :: v_magnitude, p_magnitude, n_ions, n_ion_slice
real(rp) :: macro_mass, macro_charge, s_start, s_end
real(rp) :: t_created = 0.D0
integer :: species, gas_species, open_status, ix_branch1, status
integer :: ix, seedn, snmax, option, n_slices
integer :: jjj, kkk, lll, n_included
integer, parameter :: outfile = 12, namelist_file = 13
integer, parameter :: linenum = 14, append = 15
logical :: truth2
namelist / ion_emitter_params / option, latname, outname, max_time, sspacing, &
gas_name, n_ions, n_slices, s_start, s_end, hspacing, n_included

kt = boltz_const * room_temp !Units of J
kt_ev = boltz_const_ev * room_temp !Units of eV
formato = "(es19.10,7(3x,es19.10))"

!Defaults for namelist
option = 1
latname = 'lat.bmad'
outname = 'lat.out'
max_time = 0.01
sspacing = 0.1
gas_name = 'H2'
n_ions = 1D8
n_slices = 1
s_start = 0.1
s_end = 0.1
hspacing = 0.1
n_included = 8

!Read namelist
in_file = 'ion_emitter.in'
if (command_argument_count() > 0) call get_command_argument(1, in_file)

open(namelist_file, file = in_file, status = "old", iostat=open_status)
if (open_status /= 0) then
  print *, 'File missing: ', in_file
  print *, 'Stopping'
  stop
end if
read (namelist_file, nml = ion_emitter_params)
close (namelist_file)

!Parse Lattice
call bmad_parser (latname, lat)


!Calculate tau
!Note: Until collision cross-sections are readily available, this is only true
!for H2 at 5GeV. I.e. though the gas_name may be arbitrary, this calculation is
!valid only if gas_name = 'H2' and you are considering H2 at 5GeV.
gas_species = species_id(gas_name)
ion_mass = 1.862988D9  !Units eV/c^2 for H2
gas_pressure = 6 *1D-10 * 1.33322D2 !Units Pa
gas_density = gas_pressure / kt !Units m^-3
col_cross = 3.1D-23 !Units m^2 (H2 at 5GeV)
tau = 1 / (col_cross * gas_density * c_light) !Units s
boltzparam = sqrt(kt_ev / ion_mass) * c_light !Units m/s

!Number of electrons in the beam
n_electrons = 1D8

!Time step
t_step1 = tau / n_electrons

ion_charge1 = e_charge

!-----------------------------------------------------------------------
print '(a)', "o                .oPYo.          o   o    o              "
print '(a)', "8                8                   8    8              "
print '(a)', "8 .oPYo. odYo.   `boo   ooYoYo. o8  o8P  o8P .oPYo. oPyo."
print '(a)', "8 8    8 8   8   .P     8' 8 8   8   8    8  8oooo8 8  `'"
print '(a)', "8 8    8 8   8   8      8  8 8   8   8    8  8.     8    "
print '(a)', "8 `YooP' 8   8   `YooP' 8  8 8   8   8    8  `Yooo' 8    "
!-----------------------------------------------------------------------

truth2 = .true.
allocate(orb(0:10))
ix_branch1 = 0
call twiss_and_track(lat, orb, status, ix_branch1, truth2)

!Set up randstate
seedn = 0
setc = 'exact'
setsigmacut = 6.d0
call ran_seed_put(seedn,randstate)
call ran_gauss_converter(setc,setsigmacut,getc,getsigmacut,randstate)
cpz = 0.D0


!Option 1
if (option == 1) then

!Finding maximum s value of lattice
snmax = calc_snmax(lat, sspacing)
print *, "boltzparam = ", boltzparam
call init_ion_param_arrays(lat, orb, sspacing, snmax, param)

open(outfile, file = outname, status = 'replace')
!Loop until over max_time
write(outfile,*)
do while (t_created < max_time)
  do
    call ran_uniform(sq, randstate)
    sq = sq * snmax
    ix = ceiling(sq)
    if ((param%ea(ix) .eqv. .false.) .and. (ix >= param%bstart)) exit
  end do
 
  call ran_uniform(centile, randstate)
  call calc_ion_out(param, randstate, ix, x, y, s, theta)
  v_magnitude = inverse(boltzcdf, centile, 0.D0, 2.D4, 1.D-8)
  p_magnitude = ion_mass * v_magnitude / c_light !magnitude of cp in eV
  cpx = p_magnitude * cos(theta)
  cpy = p_magnitude * sin(theta)
  write(outfile, formato) x, cpx, y, cpy, s, cpz, t_created, ion_charge1
  t_created = t_created + t_step1

end do
close(outfile)

!Option 2
else if (option == 2) then
n_ion_slice = n_ions / real(n_slices) !Number of ions per slice
if (s_end < s_start) then
  print '(a)', "Error! s_end < s_start"
else
  if (n_slices < 1) then
    print '(a)', "n_slices must be  greater than or equal to 1!"
  else
    if (n_slices == 1) then
      sspacing = 0
    else 
      sspacing = (s_end - s_start ) / (n_slices - 1)
    end if
  call init_ion_param_arrays2(lat,orb,s_start,sspacing,n_slices,param)
  end if
end if 

open(outfile, file = outname, status = 'replace')
t_created = 0.D0 !Neglect time created
write(outfile,*)
!Loop over slices
do jjj = 1, n_slices
  s = s_start + (jjj - 1)*param%sspacing
  if (jjj >= param%bstart) then
    center_x = param%xa(jjj)
    center_y = param%ya(jjj)
    do kkk = -n_included, n_included
      do lll = -n_included, n_included
        local_size = calc_number_density(param, jjj, hspacing, kkk, lll, &
        n_ion_slice)
        if (local_size <= -4.D0) then
        cycle
        end if
        !print *, "local_size = ", local_size
        macro_mass = local_size * ion_mass
        !print *, "macro_mass = ", macro_mass
        macro_charge =local_size * ion_charge1
        boltzparam = sqrt(kt_ev / macro_mass) * c_light
        call ran_uniform(centile, randstate)
        call ran_uniform(theta, randstate)
        call ran_uniform(phi, randstate)
        phi = phi * 2 * PI
        theta = theta * PI
        !print *, "boltzparam = ", boltzparam
        v_magnitude = inverse(boltzcdf, centile, 0.D0, 3.D8, 1.D-4)
        p_magnitude = macro_mass * v_magnitude / c_light !Magnitude of cp in eV
        cpx = p_magnitude * sin(theta) * cos(phi)
        cpy = p_magnitude * sin(theta) * sin(phi)
        cpz = p_magnitude * cos(theta)
        x = center_x + (kkk * hspacing)
        y = center_y + (lll * hspacing)
        write(outfile,formato) x, cpx, y, cpy, s, cpz, t_created, macro_charge
      end do
    end do
  end if
end do

close(outfile)

end if !Option if


contains

!Cumulative distribution function of Maxwell-Boltzmann distribution
function boltzcdf(x) result(y)
  real(rp) :: x
  real(rp) :: y
  real(rp) :: tmp, tmp2
  tmp = erf(x/(sqrt(2.d0)*boltzparam))
  tmp2 = sqrt(2.d0/PI)*x*exp(-x*x/(2*boltzparam*boltzparam))/boltzparam
  y = tmp - tmp2
end function boltzcdf

end program