module ion_tracker_mod

use sim_utils
use bmad
use beam_mod
use ion_tracker_struct

implicit none

contains

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+
! Subroutine turn_track1_custom_on_for_all_elements(lat,branch)
!
! Routine sets ele(i)%tracking_method = custom$ for all tracking elements
!
! Input:
!   lat        -- lat_struct: Lattice.
!   branch     -- branch_struct: branch to turn on
!
! Output:
!   lat        -- lat_struct: Lattice with custom$ turned on
!
!-

subroutine turn_track1_custom_on_for_all_elements(lat, branch)

type(lat_struct) :: lat
type(branch_struct) :: branch

integer :: ele_id

print *, 'Diverting default tracking method to track1_custom for all elements.'

do ele_id = 1, branch%n_ele_track
    ! Save previous tracking routine
    branch%ele(ele_id)%ix_pointer = branch%ele(ele_id)%tracking_method
    if ((branch%ele(ele_id)%key /= rfcavity$) .and. &
        (branch%ele(ele_id)%key /= lcavity$)) then
    !Divert EM field calculations to a custom routine
    branch%ele(ele_id)%tracking_method = custom$
    end if
end do

end subroutine turn_track1_custom_on_for_all_elements
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+
! Subroutine special_track_beam_through_ions(lat, param, beam_init, beam, &
!                                            error_flag, ds, version)
!
!  Routine tracks beam through ions
!  Should be called after turn_track1_custom_on_for_all_elements
!  You must delete the old files before calling this subroutine again.
!  Needs an element with 0 length at the start (e.g. marker)
!
!  Input:
!    lat         -- lat_struct: Lattice.
!    param       -- lat_param_struct: General parameters.
!    beam_init   -- beam_init_struct: Starting bunch position
!    ds          -- real(rp): Separation between ion kicks
!    version     -- Integer: Version of ion kick
!    beam        -- beam_struct: Must initialize before
!
!  Output:
!    beam        -- beam_struct
!    error_flag  -- logical: true if there is an error
!
!
!-

subroutine special_track_beam_through_ions(lat, param, beam_init, beam, error_flag, ds, version)


type(lat_struct), target :: lat
type(ele_struct), pointer :: ele1, ele2, ele 
type(ele_struct), target :: mini_ele
type(lat_param_struct) :: param
type(bunch_params_struct) :: bunch_params
type(branch_struct), pointer :: branch
type(beam_struct) :: beam
type(beam_init_struct) :: beam_init
integer :: n_ele_track, i, j, k, filetype, version, nm, ix_branch
logical :: error_flag, new_file, use_fixed_sigs
character(100) :: beam_name, size_name

real(rp) :: ds, length
integer, parameter :: bm = 31


! Deciding beam_name
select case(version)
  case(1)
    beam_name = "ion_kick_v1_beam.out"
    size_name = "size_v1_beam.out"
  case(2)
    beam_name = "ion_kick_v2_beam.out"
    size_name = "size_v2_beam.out"
  case(3)
    beam_name = "ion_kick_v3_beam.out"
    size_name = "size_v3_beam.out"
  case(4)
    beam_name = "ion_kick_v4_beam.out"
    size_name = "size_v4_beam.out"
  case(5)
    beam_name = "ion_kick_v5_beam.out"
    size_name = "size_v5_beam.out"
  case default
    beam_name = "no_ion_beam.out"
    size_name = "size_no_beam.out"
end select

!Initializations
ix_branch = 0
new_file = .false.
filetype = ascii$

print *, 'Tracking through lattice'
ele1 => lat%ele(0)
n_ele_track = lat%n_ele_track
ele2 => lat%ele(n_ele_track)
branch => lat%branch(ele1%ix_branch)

!No ion kicks
if (version == 0) then

!i range determines which elements to track
do i = ele1%ix_ele+1, ele2%ix_ele
  print *, 'i = ', i
  ele => branch%ele(i)
  call track1_beam (beam, lat, ele, beam, error_flag)
  if (error_flag) return
    call write_beam_essentials(beam_name, beam, new_file)
    !call write_beam_size(size_name, beam, new_file, lat, ele)
    !call write_beam_size_condensed(size_name, beam, new_file, lat, ele)
  open (bm, file = beam_name, access = 'append')
  close (bm)
end do

!Ion kicks
else
call twiss_propagate_all(lat, ix_branch, error_flag)
if (error_flag) return

!i range determines which elements to track
do i = ele1%ix_ele+1, ele2%ix_ele
  print *, 'i = ', i
  ele => branch%ele(i)
  length = branch%ele(i)%value(l$)

  !Trouble with energies with cavities
  !Sliced elemtns do not have the appropriate energies when sliced.
  if ((ele%key /= rfcavity$) .and. (ele%key /= lcavity$)) then
    

    !Ekement must be longer than 2*ds to be sliced
    !This ensures that every slice (disregarding unsliced elements)
    !has length at least ds
    if (length < 2.d0*ds) then
      call cbeta_ions_track1_beam_helper(beam, lat, ele, beam)
      call write_beam_essentials(beam_name, beam, new_file)
      !call write_beam_size(size_name, beam, new_file, lat, ele)
      !call write_beam_size_condensed(size_name, beam, new_file, lat, ele)
    open(bm, file = beam_name, access = 'append')
    write(bm, '(es19.10)') branch%ele(i)%s
    close(bm)
    !Elements large enough to slice
    !Divided into beginning slice, middle slices, and end slices in code
    !(j == 1, 2 <= j <= nm-1, j == nm)
    !Artefact of syntax of create_element_slice since it required flagging
    !whether a slice included the upstream or downstream end.
    !Only difference between create_element_slice and
    !create_uniform_element_slice is that create_element_slice can make
    !slices with uneven lengths.
      
    !Comment and uncomment depending on the output needed.
    !All of the output modifying routines such as
    !write_beam_essentials, write_beam_size write_beam_size_condensed,
    !write(bm,'(es19.10)') statements can be commented without creating
    !an error.

    !Unless you need the emittance and twiss parameters,
    !you can comment out the calc_bunch_params statements.
    !If you set sigma to a certain number, then you don't need
    !calc_bunch_params. The error 'No eigensystem found' is caused by
    !the program taking too many steps to find the eigenvalues.
    !Raising the ds will cause the error to go away.

    else
      !nm is the number of slices of element i (integer)
      nm = floor(length / ds)
      use_fixed_sigs = return_usefixedsig()
      if(.not. use_fixed_sigs) then
        call calc_bunch_params(beam%bunch(1), bunch_params, error_flag)
        !Emittance passed to calculate sigmas of the ion distribution
        ele%a%emit = bunch_params%a%emit
        ele%b%emit = bunch_params%b%emit
        ele%z%emit = bunch_params%c%emit
      end if
      call create_uniform_element_slice(ele, param, 1, nm, mini_ele)
      call twiss_propagate1(ele, mini_ele, error_flag)
      if (error_flag) return
      call cbeta_ions_track1_beam_helper(beam, lat, mini_ele, beam, ele)
      call write_beam_essentials(beam_name, beam, new_file)
      open(bm, file = beam_name, access = 'append')
      write(bm, '(es19.10)') mini_ele%s
      close(bm)
      !call write_beam_size(size_name, beam, new_file, lat, ele)
      !call write_beam_size_condensed(size_name, beam, new_file, lat, mini_ele)
      if (nm >= 3) then
        do j = 2, nm-1
          call create_uniform_element_slice(ele, param, j, nm, mini_ele)
          call twiss_propagate1(ele, mini_ele, error_flag)
          if (error_flag) return
          call cbeta_ions_track1_beam_helper(beam, lat, mini_ele, beam, ele)
          call write_beam_essentials(beam_name, beam, new_file)
          !call write_beam_size_condensed(size_name, beam, new_file, lat, mini_ele)
          open(bm, file = beam_name, access = 'append')
          write(bm, '(es19.10)') mini_ele%s
          close(bm)
        
        end do
      end if
      call create_uniform_element_slice(ele, param, nm, nm, mini_ele)
      call twiss_propagate1(ele, mini_ele, error_flag)
      if (error_flag) return
      call cbeta_ions_track1_beam_helper(beam, lat, mini_ele, beam, ele)
      call write_beam_essentials(beam_name, beam, new_file)
      !call write_beam_size_condensed(size_name, beam, new_file, lat, mini_ele)
      open(bm, file = beam_name, access = 'append')
      write(bm, '(es19.10)') mini_ele%s
      close(bm)
    end if
  else
    call track1_beam(beam, lat, ele, beam, error_flag)
    if (error_flag) return
    call write_beam_essentials(beam_name, beam, new_file)
    open(bm, file = beam_name, access = 'append')
    write(bm, '(es19.10)') branch%ele(i)%s
    close(bm)

    !call write_beam_size_condensed(size_name, beam, new_file, lat, ele)
  end if
end do
end if
end subroutine special_track_beam_through_ions

subroutine cbeta_ions_header()

print *, ' '
print '(a)', " _____              _____              _             "
print '(a)', "|_   _|            |_   _|            | |            "
print '(a)', "  | |  ___  _ __     | |_ __ __ _  ___| | _____ _ __ "
print '(a)', "  | | / _ \| '_ \    | | '__/ _` |/ __| |/ / _ \ '__|"
print '(a)', " _| || (_) | | | |   | | | | (_| | (__|   <  __/ |   "
print '(a)', " \___/\___/|_| |_|   \_/_|  \__,_|\___|_|\_\___|_|   "
print *, '---------------------------------------------------------------------'
print *, ' '
print *, 'cbeta_ions : Tracking through a given lattice while including ion kicks'
print *, ' '
print *, '---------------------------------------------------------------------'
print *, ' '
end subroutine cbeta_ions_header

!---------------------------------------------------
!
! Subroutine apply_ion_kicks(start_orb, ele, param, end_orb, kick)
!
! Returns kick felt due to ions
!
! Input:
!   start_orb    --   coord_struct: Starting position.
!   ele          --   ele_struct: Element.
!   param        --   lat_param_struct: Lattice parameters.
!   version      --   logical: Version of ion kick. 0 for no ion kick
! Output:
!   end_orb      --   coord_struct: End position.
!   kick(3)      --   real(rp): (x, y, z) kick in kg*m/sec.
!
!---------------------------------------------------

subroutine apply_ion_kicks(start_orb, ele, param, end_orb, kick, version)

type (coord_struct) :: start_orb
type (coord_struct) :: end_orb
type (ele_struct) :: ele
type (lat_param_struct) :: param
real(rp) :: kick(3)

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, kt_ev
real(rp) :: x0, y0, A, r0, phi0, vx0, vy0, v0
real(rp) :: theta0, vz, vz0, tau, x_perp, y_perp, v_perp
real(rp) :: L, Lr, length, omega, charge_density, p0c
real(rp) :: px, py, pz, m_electron, tmp1, xnorm, ynorm
real(rp) :: rho, vx, vy, rtau, vrtau, vxtau, vytau, vztau
real(rp) :: thetatau, vthetatau, pxtau, pytau, pztau, sigma
real(rp) :: halfpi, linear_charge_density, emit, sigma_a, sigma_b
real(rp) :: El_mag, p_mag, dpx, dpy, rel_gamma, tmp, emit_a, emit_b
real(rp) :: El_x, El_y, s, ratio
!Namelist parameters that need to be passed manually for now
real(rp) :: sig_ee = 7.D-5
real(rp) :: n_beam_part = 5.D4
real(rp) :: r_beam(2) = (/ 0, 0 /)

real(rp), parameter :: tol = 1.D-10, tol2 = 2D-3
logical :: zeroness, use_fixed_sigs
integer :: version
halfpi = asin(1.D0)
m_electron = 9.10938536D-31 !Unit kg
kt = boltz_const * room_temp !Units of J
kt_ev = boltz_const_ev * room_temp !Units of eV
call coord_equal_coord(end_orb, start_orb)

s = start_orb%s
x0 = start_orb%vec(1)
y0 = start_orb%vec(3)
length = ele%value(l$)
p0c = start_orb%p0c
px = start_orb%vec(2) * p0c * 5.344286D-28
py = start_orb%vec(4) * p0c * 5.344286D-28
pz = (start_orb%vec(6) * p0c) + p0c
pz = pz * 5.344286D-28
! px, py, pz converted from eV/c to kg*m/s
! Solving for vx0, vy0, vz from px, py, pz
tmp1 = px/m_electron
vx0 = sqrt((tmp1**2)/(1+(tmp1/c_light)**2))
tmp1 = py/m_electron
vy0 = sqrt((tmp1**2)/(1+(tmp1/c_light)**2))
tmp1 = pz/m_electron
vz0 = sqrt((tmp1**2)/(1+(tmp1/c_light)**2))
v0 = sqrt(vx0**2 + vy0**2)
tau = length / vz0
r0 = sqrt(x0**2+ y0**2)
if (r0 < tol) then
  xnorm = 0
  ynorm = 0
  zeroness = .true.
else
  zeroness = .false.
  xnorm = x0/r0
  ynorm = y0/r0
end if
x_perp = -ynorm
y_perp = xnorm
if (zeroness .eq. .false.) then
  if (x0 > tol) then
    theta0 = atan(y0/x0)
  elseif (x0 < -tol) then
    theta0 = atan(y0/x0) + PI
  elseif (y0 > tol) then
    theta0 = halfpi
  else
    theta0 = -halfpi
  end if
else 
  if (vx0 > tol) then
    theta0 = atan(vy0/vx0)
  elseif (vx0 < -tol) then
    theta0 = atan(vy0/vx0) + PI
  elseif (vy0 > tol) then
    theta0 = halfpi
  else
    theta0 = -halfpi
  end if
end if
v_perp = vx0 * x_perp + vy0 * y_perp

!--------------------------------------------------------
!Version 1   Integrated Uniform Cylindrical
!--------------------------------------------------------

if (version == 1) then
vz = vz0

L = m_electron * r0 * v_perp
Lr = r0 * v_perp
!Arbitrarily chosen ion density rho
rho = 1D-10 ! Coulomb / m^3
omega = sqrt(e_charge*rho/(2.D0*m_electron*8.85418782D-12))
if (r0 < tol) then
  phi0 = sign(halfpi,-v0/omega)
else
  phi0 = atan(-v0/omega*r0)
end if
A = r0/cos(phi0)
rtau = A*cos(omega*tau + phi0)
vrtau = -omega*A*sin(omega*tau + phi0)
if (zeroness .eq. .true.) then
  vthetatau = 0.D0
else
  vthetatau = Lr/rtau
end if
if (zeroness .eq. .true.) then
  thetatau = theta0
else
  thetatau = theta0 + Lr/(omega*(A**2))*(tan(omega*tau+phi0)-tan(phi0))
end if
vxtau = cos(thetatau)*vrtau - sin(thetatau)*vthetatau
vytau = sin(thetatau)*vrtau + cos(thetatau)*vthetatau
kick(1) = vxtau - vx0
kick(2) = vytau - vy0
kick(3) = 0
end_orb = start_orb

tmp = vxtau**2 + vytau**2 + vz**2
rel_gamma = 1.D0/sqrt(1.D0 - tmp/(c_light**2))
pxtau = rel_gamma*m_electron*vxtau
pytau = rel_gamma*m_electron*vytau

!Convert to eV/c
pxtau = pxtau / 5.344286D-28
pytau = pytau / 5.344286D-28
!Convert to normalized px, py
pxtau = pxtau / p0c
pytau = pytau / p0c
end_orb%vec(2) = pxtau
end_orb%vec(4) = pytau

!------------------------------------------------------
!Version 2     Uniform Cylindrical
!------------------------------------------------------

elseif (version == 2) then
!linear_charge_density = 5.79D-10 !4pass
!linear_charge_density = 3.62D-10 !1pass

!emit_a = ele%a%emit
!emit_b = ele%b%emit
!emit = max(emit_a, emit_b)

!sigma_a = sqrt(emit * ele%a%beta)
!sigma_b = sqrt(emit * ele%b%beta)
!sigma = max(sigma_a, sigma_b)

use_fixed_sigs = return_usefixedsig()
if (use_fixed_sigs) then
  sigma = return_sig()
else
  emit_a = ele%a%emit
  emit_b = ele%b%emit
  emit = max(emit_a, emit_b)
  sigma_a = sqrt(emit * ele%a%beta)
  sigma_b = sqrt(emit * ele%b%beta)
  sigma = max(sigma_a, sigma_b)
end if

linear_charge_density = linear_density_calculator(s, 1)
rho = linear_charge_density/(PI * sigma * sigma)
kick(3) = 0.D0
!Calculating kicks
if (zeroness .eqv. .true.) then
  kick(1) = 0.D0
  kick(2) = 0.D0
else
  El_mag = rho*r0/(2.D0 * 8.85418782D-12)
  p_mag = El_mag*e_charge*tau
  kick(1) = -cos(theta0) * p_mag
  kick(2) = -sin(theta0) * p_mag
end if

px = px + kick(1)
py = py + kick(2)

tmp = sqrt(px**2 + py**2 + pz**2)

!Convert to eV/c
pxtau = px / 5.344286D-28
pytau = py / 5.344286D-28
!Convert to normalized px, py
pxtau = pxtau / p0c
pytau = pytau / p0c
end_orb%vec(2) = pxtau
end_orb%vec(4) = pytau

!-----------------------------------------------------------------
!Version 3        Gaussian distribution
!-----------------------------------------------------------------
elseif (version == 3) then
!linear_charge_density = 5.79D-10 !4pass
!linear_charge_density = 3.62D-10 !1pass
!linear_charge_density = 0
linear_charge_density = linear_density_calculator(s, 2)

use_fixed_sigs = return_usefixedsig()
if (use_fixed_sigs) then
  sigma_a = return_sigx()
  sigma_b = return_sigy()
else
  emit_a = ele%a%emit
  emit_b = ele%b%emit
  sigma_a = sqrt(emit_a * ele%a%beta)
  sigma_b = sqrt(emit_b * ele%b%beta)
  ratio = return_sigmaratio()
  sigma_a = sigma_a * ratio
  sigma_b = sigma_b * ratio
end if

call gaussian_charge_electric_field(El_x, El_y, x0, y0, linear_charge_density, &
     sigma_a, sigma_b)
kick(1) = -El_x * e_charge * tau
kick(2) = -El_y * e_charge * tau
kick(3) = 0.D0

pxtau = px + kick(1)
pytau = py + kick(2)

!Convert to eV/c

pxtau = px / 5.344286D-28
pytau = py / 5.344286D-28

!Convert to normalized px, py
pxtau = pxtau / p0c
pytau = pytau / p0c
end_orb%vec(2) = pxtau
end_orb%vec(4) = pytau



!----------------------------------------

else
  print *, 'Incorrect version'
  call err_exit
end if
end subroutine apply_ion_kicks
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!+
!
! Subroutine write_beam_essentials (file_name, beam, new_file)
!
! Writes essential information of beam
!
! Input:
!   file_name    --   character(*): Name of file.
!   beam         --   beam_struct: Beam to write
!   new_file     --   logical, optional: New file or append? Default = True.
!-

subroutine write_beam_essentials (file_name, beam, new_file)
character(100) :: file_name, formato
type(beam_struct), target :: beam
type (bunch_struct), pointer :: bunch
type (coord_struct), pointer :: p
integer j, k, iu
logical :: new_file

iu = lunget()
formato = '(6es19.10, es14.5, i6)'


if (new_file) then
  open (iu, file = file_name)
else
  open (iu, file = file_name, access = 'append')
endif

write (iu, '(a)') 'BEGIN_BUNCH'
do j = 1, size(beam%bunch)
  bunch => beam%bunch(j)
  do k = 1, size(bunch%particle)
    p => bunch%particle(k)
    write (iu, formato) p%vec, p%charge, p%state
  end do
end do
write (iu, '(a)') 'END_BUNCH'

close (iu)

end subroutine write_beam_essentials
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!+
!
! Subroutine write_beam_size (file_name, beam, new_file, lat, ele)
!
! Writes beam sigmas, emittances, beta, alpha
!
! Input:
!   file_name    --   character(*): Name of file.
!   beam         --   beam_struct: Beam to write
!   new_file     --   logical, optional: New file or append? Default = True.
!   lat          --   lat_struct: Lattice.
!   ele          --   ele_struct: Element.
!-

subroutine write_beam_size (file_name, beam, new_file, lat, ele)
character(100) :: file_name, formati
type(beam_struct), target :: beam
type(lat_struct) :: lat
type(ele_struct) :: ele
type (bunch_struct), pointer :: bunch
type (bunch_params_struct) :: prm
integer j, iu
logical :: new_file, error_flag

iu = lunget()
formati = '(i6,15es19.10)'
! Format is sigmas, sigma_ps, emittances, betas, alphas

if (new_file) then
  open (iu, file = file_name)
else
  open (iu, file = file_name, access = 'append')
endif

do j = 1, size(beam%bunch)
  bunch => beam%bunch(j)
  call calc_bunch_params(beam%bunch(j), prm, error_flag)
  if (error_flag) then
    print *, 'Error in calc_bunch_params'
    call exit(0)
  end if
  write (iu, formati) prm%n_particle_live, prm%a%sigma, prm%b%sigma, &
  prm%c%sigma, prm%a%sigma_p, prm%b%sigma_p, prm%c%sigma_p, &
  prm%a%emit, prm%b%emit, prm%c%emit, prm%a%beta, prm%b%beta, &
  prm%c%beta, prm%a%alpha, prm%b%alpha, prm%c%alpha
end do

close (iu)

end subroutine write_beam_size

subroutine write_beam_size_condensed (file_name, beam, new_file, lat, ele)
character(100) :: file_name, formati
type(beam_struct), target :: beam
type(lat_struct) :: lat
type(ele_struct) :: ele
type (bunch_struct), pointer :: bunch
type (bunch_params_struct) :: prm
integer j, iu
logical :: new_file, error_flag

iu = lunget()
formati = '(3es19.10)'
! Format is sigmas, sigma_ps, emittances, betas, alphas

if (new_file) then
  open (iu, file = file_name)
else
  open (iu, file = file_name, access = 'append')
endif

do j = 1, size(beam%bunch)
  bunch => beam%bunch(j)
  call calc_bunch_params(beam%bunch(j), prm, error_flag)
  if (error_flag) then
    print *, 'Error in calc_bunch_params'
    call exit(0)
  end if
  write (iu, formati) ele%s, prm%a%emit, prm%a%beta
end do

close (iu)

end subroutine write_beam_size_condensed



!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!+
!
! Subroutine cbeta_ions_track1_beam_helper
!
! Helper routine
!
! Input:
!   beam_start   --   Beam_struct: Initial beam
!   lat          --   Lat_struct
!   ele          --   Ele_struct
!   or_ele       --   Ele_struct, optional: Holds the original element in case
!                     the element is a slice.
!
! Output:
!   beam_end     --   beam_struct: Tracked beam
!-
subroutine cbeta_ions_track1_beam_helper(beam_start, lat, ele, beam_end, or_ele)

type(beam_struct) :: beam_start, beam_end
type(lat_struct) :: lat
type(ele_struct) :: ele
type(bunch_params_struct) :: bunch_params
integer :: i, j, debug1, debug2
logical :: error_flag, use_fixed_sigs

type(ele_struct), optional :: or_ele


use_fixed_sigs = return_usefixedsig()

do i = 1, size(beam_start%bunch)
  if (.not. use_fixed_sigs) then
    call calc_bunch_params(beam_start%bunch(i), bunch_params, error_flag)
    if (error_flag) then
      print *, 'Error in calc_bunch_params'
      !j = 1/0  --  Unknown whether it truly acts like an error or is a warning
    end if
    ele%a%emit = bunch_params%a%emit
    ele%b%emit = bunch_params%b%emit 
    ele%z%emit = bunch_params%c%emit
  end if

  do j = 1, size(beam_start%bunch(i)%particle)
    if (beam_start%bunch(i)%particle(j)%state /= alive$) cycle
    call track1(beam_start%bunch(i)%particle(j), ele, lat%param, &
                beam_start%bunch(i)%particle(j))
    if(present(or_ele)) beam_start%bunch(i)%particle(j)%ix_ele = or_ele%ix_ele
    beam_start%bunch(i)%charge_live = &
    sum(beam_start%bunch(i)%particle(:)%charge, mask = &
    (beam_start%bunch(i)%particle(:)%state == alive$))
  end do
end do

end subroutine cbeta_ions_track1_beam_helper

!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!+
!
!  Subroutine full_complex_error_function(wr,wi,zr,zi)
!
! Expands complex_error_function to the whole complex plane
!
! Input:
!   zr -- Real(rp): Real part of z.
!   zi -- Real(rp): Imaginary part of z.
!
! Output:
!   wr -- Real(rp): Real part of w.
!   wi -- Real(rp): Imaginary part of w.
!-

subroutine full_complex_error_function(wr, wi, zr, zi)
real(rp) :: zr, zi, wr, wi, tmpx, tmpy
complex(rp) :: z, zz, tmpz


if (zr >= 0 .and. zi >= 0) then
  call complex_error_function(wr, wi, zr, zi)
elseif (zr < 0 .and. zi >= 0) then
  call complex_error_function(tmpx, tmpy, -zr, zi)
  wr = tmpx
  wi = -tmpy
elseif (zr < 0 .and. zi < 0) then
  call complex_error_function(tmpx, tmpy, -zr, -zi)
  zz = cmplx(zr, zi, kind(rp))
  z = 2*exp(-1*(zz**2))
  tmpz = cmplx(tmpx, tmpy, kind(rp))
  z = z - tmpz
  wr = real(z, kind(rp))
  wi = aimag(z)
elseif (zr > 0 .and. zi < 0) then
  call complex_error_function(tmpx, tmpy, zr, -zi)
  zz = cmplx(zr, zi, kind(rp))
  z = 2*exp(-1*(zz**2))
  tmpz = cmplx(tmpx, -tmpy, kind(rp))
  z = z - tmpz
  wr = real(z, kind(rp))
  wi = aimag(z)
end if

end subroutine full_complex_error_function

!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!-----------------------------------------------------------------------------
!+
!
!  Subroutine gaussian_charge_electric_field(ex,ey,xx,yy,lambda,sigx,sigy)
!
! Gives transverse electric field at a position (xx,yy) from
! elliptical gaussian charge distribution with linear charge density lambda
! Assumes the charge distribution is centered at (0,0)
!
! Input:
!   xx --     Real(rp): x position to evaluate electric field
!   yy --     Real(rp): y position to evaluate electric field
!   lambda -- Real(rp): Linear charge density
!   sigx --   Real(rp): x sigma of gaussian charge distribution
!   sigy --   Real(rp): y sigma of gaussian charge distribution 
!
! Output:
!   ex -- Real(rp): x component of electric field at (xx,yy)
!   ey -- Real(rp): y component of electric field at (xx,yy)
!-

subroutine gaussian_charge_electric_field(ex,ey,xx,yy,lambda,sigx,sigy)

real(rp) :: xx, yy, rr, lambda, ex, ey, sigx, sigy, tmpx, tmpy, const, er
real(rp) :: halfpi, theta, cx, cy, tmp, tmp2
real(rp), parameter :: tol = 1.D-7, tol2 = 2.D-3
complex(rp) :: zz, tmpz1, tmpz2
logical :: zeroness

halfpi = pi/2
ex = 0
ey = 0
rr = sqrt(xx**2 + yy**2)
if (rr < tol) then
  zeroness = .true.
else
  zeroness = .false.
endif


if (zeroness .eq. .false.) then
  if (xx > tol) then
    theta = atan(yy/xx)
  elseif (xx < -tol) then
    theta = atan(yy/xx) + PI
  elseif (yy > tol) then
    theta = halfpi
  else
    theta = -halfpi
  end if


  if (abs(sigx - sigy) < tol2) then
    const = lambda*e_charge/(twopi*rr*8.85418782D-12)
    er = const * (1.D0 - exp(-5.D-1*((rr/sigx)**2)))
    ex = er * cos(theta)
    ey = er * sin(theta)
  else
    if (sigy > sigx) then
      tmpx = sigx
      sigx = sigy
      sigy = tmpx
      tmpx = xx
      xx = yy
      yy = xx
    endif
    const = lambda*e_charge/(2.D0*8.85418782D-12*sqrt(twopi*(sigx**2-sigy**2)))
    tmp2 = 1.D0/sqrt(2.D0*(sigx**2-sigy**2))
    cx = xx*tmp2
    cy = yy*tmp2
    call full_complex_error_function(tmpx, tmpy, cx, cy)
    tmpz1 = cmplx(tmpx, tmpy, kind(rp))
    tmp = -1.D0*exp((-5.D-1*(xx/sigx)**2)+(-5.D-1*(yy/sigy)**2))
    cx = xx*tmp2*sigy/sigx
    cy = yy*tmp2*sigx/sigy
    call full_complex_error_function(tmpx, tmpy, cx, cy)
    tmpz2 = cmplx(tmpx, tmpy, kind(rp))
    zz = tmpz1 + tmp*tmpz2
    ex = const * aimag(zz)
    ey = const * real(zz, kind(rp))
    if (sigy > sigx) then
      tmpx = ex
      ex = ey
      ey = tmpx
    endif
  endif
endif



end subroutine



end module