module ion_emitter_mod

use bmad
use random_mod
use twiss_and_track_mod
use ion_emitter_struct

implicit none

contains
!-------------------------------------------------------------------------
!-------------------------------------------------------------------------
!-------------------------------------------------------------------------
!+
! Subroutine sigma_at_s(lat, orb, s, sig_x, sig_y, x, y, error)
!
! Routine to calculate sigma_x and sigma_y as well as the orbit at s.
! Assumes only one branch.
!
! Input:
!   lat        -- lat_struct: Lattice.
!   orb(0:)    -- Coord_struct: Orbit through the lattice
!   s          -- s value to calculate sigma_x, simga_y
!
! Output:
!   sig_x      -- Real(rp): x sigma
!   sig_y      -- Real(rp): y sigma
!   x          -- Real(rp): x offset in orbit
!   y          -- Real(rp): y offset in orbit
!   error      -- Logical : .true. if there was an error in calculation
!-

subroutine sigma_at_s(lat, orb, s, sig_x, sig_y, x, y, error)

type (lat_struct) :: lat
type (coord_struct) :: orb(:)
type (coord_struct) :: orbit_end
type (ele_struct) :: end_ele
real (rp) s, x_beta, y_beta, x, y, sig_x, sig_y, norm_emit, emit
integer species, ix_branch
logical error

species = lat%param%particle
error = .false.
ix_branch = 0
call twiss_and_track_at_s(lat, s, end_ele, orb, orbit_end, ix_branch, error)

x_beta = end_ele%a%beta
y_beta = end_ele%b%beta
x = orbit_end%vec(1)
y = orbit_end%vec(3)
norm_emit = 1D-12
emit = norm_emit * mass_of(species) / end_ele%value(e_tot$)
sig_x = sqrt(x_beta * emit)
sig_y = sqrt(y_beta * emit)

end subroutine

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+
!  Subroutine init_ion_param_arrays(lat, orb, sspacing, nmax, param)
!
!  Routine to initialize the arrays needed for ion tracker
!
!  Input:
!    lat         -- lat_struct: Lattice
!    orb(0:)     -- coord_struct: Orbit through the lattice.
!    sspacing    -- Real(rp): s spacing given by user.
!    nmax        -- Integer : Maximum value of index; s_max = sspacing * nmax.
!  
!  Output:
!    param       -- ion_emitter_param_struct
!-

subroutine init_ion_param_arrays(lat, orb, sspacing, nmax, param)

type (lat_struct) :: lat
type (coord_struct) :: orb(:)
real(rp) :: sspacing, s
integer :: nmax, bstart, iii
type (ion_emitter_param_struct) :: param
real(rp) :: x, y, sig_x, sig_y 
logical :: error

bstart = 0

allocate(param%xa(1:nmax))
allocate(param%ya(1:nmax))
allocate(param%sxa(1:nmax))
allocate(param%sya(1:nmax))
allocate(param%ea(1:nmax))
param%sspacing = sspacing
param%nmax = nmax

do iii = 1, nmax
s = sspacing * iii
call sigma_at_s(lat, orb, s, sig_x, sig_y, x, y, error)
param%xa(iii) = x
param%ya(iii) = y
param%sxa(iii) = sig_x
param%sya(iii) = sig_y
param%ea(iii) = error
if ((.not. error) .and. (abs(sig_x) > 1D-30) .and. (abs(sig_y) > 1D-30) &
     .and. (bstart == 0)) bstart = iii
end do

param%bstart = bstart
end subroutine

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+
!  Subroutine init_ion_param_arrays2(lat, orb, s_start, sspacing, n_slices,
!  param)
!
!  Input:
!    lat          -- lat_struct: Lattice
!    orb(0:)      -- coord_struct: Orbit through lattice
!    s_start      -- Real(rp): s position of first slice
!    sspacing     -- Real(rp): interval in meters between slices
!    n_slices     -- Intger: Number of slices
!
!  Output:   
!    param        -- ion_emitter_param_struct
!-

subroutine init_ion_param_arrays2(lat,orb,s_start,sspacing,n_slices,param)
type (lat_struct) :: lat
type (coord_struct) :: orb(:)
real(rp) :: s_start, sspacing, s, sigx, sigy, x, y
logical :: error
integer :: n_slices, bstart, iii
type (ion_emitter_param_struct) :: param
param%nmax = n_slices
param%sspacing = sspacing
allocate(param%xa(1:n_slices))
allocate(param%ya(1:n_slices))
allocate(param%sxa(1:n_slices))
allocate(param%sya(1:n_slices))
allocate(param%ea(1:n_slices))

if (sspacing < 0) then
print '(a)', "sspacing < 0 !!!"
end if

bstart = 0
do iii = 1, n_slices
s = s_start + (iii-1) * sspacing
call sigma_at_s(lat, orb, s, sigx, sigy, x, y, error)
param%xa(iii) = x
param%ya(iii) = y
param%sxa(iii) = sigx
param%sya(iii) = sigy
param%ea(iii) = error
if ((.not. error) .and. (abs(sigx) > 1D-30) .and. (abs(sigy) > 1D-30) &
   .and. (bstart == 0)) bstart = iii
end do
param%bstart = bstart

end subroutine 
 

!------------------------------------------------------------------------
!------------------------------------------------------------------------
!------------------------------------------------------------------------
!+
!  function calc_snmax(lat, sspacing) result (snmax)
!    
!    Input:
!      lat      -- lat_struct: Lattice.
!      sspacing -- Real(rp)  : s spacing
!
!    Output:
!      snmax    -- Integer   : maximum index for ion_emitter_param_struct
!-
function calc_snmax(lat, sspacing) result (snmax)
  type (lat_struct) :: lat
  real(rp) :: sspacing, s_end, snmaxr
  integer :: snmax, nelemax

nelemax = lat%branch(0)%n_ele_max
s_end = lat%branch(0)%ele(nelemax)%s
snmaxr = s_end / sspacing
snmax = floor(snmaxr)

end function calc_snmax

!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!+
!  function boltzcdf2(x, boltzparam) result (y)
!    
!    Cumulative distribution function of Maxwell-Boltzmann distribution
!    Not used in ion emitter program because inverse needs
!    a function with only one parameter. Can be used to define a one
!    parameter boltzcdf function within a program.
!
!    Input:
!      x          -- Real(rp): velocity (m/s)
!      boltzparam -- Real(rp): kt/ ion_mass
!
!    Output:
!      y          -- Real(rp): probability of (velociy <= y)
!-
function boltzcdf2(x, boltzparam) result(y)
  real(rp) :: x, y, boltzparam
  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 boltzcdf2

!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!+
!   Subroutine calc_ion_pos(param, randstate, ix, x, y, angle)
!   
!     Routine that calculates x, y of one ion at s = ix * sspacing.
!     Also generates angle (tan^-1(py/px))
!     Doesn't check if ix >= bstart.
!
!     Input:
!       param  -- ion_emitter_param_struct
!       randstate -- random_state_struct
!       ix     -- Integer: index where values are calculated
!
!     Output:
!       x      -- Real(rp) Units: m
!       y      -- Real(rp) Units: m
!       s      -- Real(rp) Units: m      s = ix * param%sspacing
!       angle  -- Real(rp) Units: rad (unitless)
!-
subroutine calc_ion_out(param, randstate, ix, x, y, s, angle)
type (ion_emitter_param_struct) :: param
type (random_state_struct) randstate
integer :: ix
real(rp) :: ion_mass, x, y, s, angle, tmp1, tmp2
real(rp) :: boltzparam

call ran_gauss(tmp1, randstate)
call ran_gauss(tmp2, randstate)
call ran_uniform(angle, randstate)

angle = 2 * pi * angle
x = tmp1 * param%sxa(ix) + param%xa(ix)
y = tmp2 * param%sya(ix) + param%ya(ix)
s = ix * param%sspacing

end subroutine

!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!-----------------------------------------------------------------------
!+
!  function calc_number_density(param, ix, h, n, m, tot_ion) result n_ion
!    Calculates the number of ions near
!    (x_center,y_center)+(n*hspacing,m*hspacing)
!
!    Input:
!      param     -- ion_emitter_param_struct
!      ix        -- integer : index (s = sspacing * ix)
!      h  -- real(rp): horizontal spacing
!      n         -- integer : x = n*hspacing
!      m         -- integer : y = m*hspacing
!      tot_ion   -- real(rp): total number of ions in the whole slice
!
!    Output:
!      n_ion      -- real(rp): number of ions in the region
!-
function calc_number_density(param, ix, h, n, m, tot_ion) result(n_ion)
  type (ion_emitter_param_struct) :: param
  integer :: ix, n, m
  real(rp) :: h, sigx, sigy, tot_ion
  real(rp) :: x, y, tmp, tmp2, tmp3, tmp4, tmp5, n_ion
  sigx = param%sxa(ix)
  sigy = param%sya(ix)
  x = param%xa(ix)
  y = param%ya(ix)
  tmp = 1.D0/(2*PI*sigx*sigy)
  tmp2 = exp(-5.D-1*(n*h/sigx)*(n*h/sigx))
  tmp3 = exp(-5.D-1*(m*h/sigx)*(m*h/sigy))
  n_ion = tot_ion* tmp * tmp2 * tmp3 * h * h
  if (n_ion < 2.D0) then
  n_ion = -5.D0
  !print *, "n_ion < 2.D0"
  else if (n_ion > tot_ion) then
  n_ion = -6.D0
  !print *, "n_ion > tot_ion: Make hspacing bigger"
  end if

end function calc_number_density



end module