module ion_tracker_struct

use bmad

implicit none

integer :: num_s = 2000
logical :: use_fixedsig = .false.
real(rp) :: sigmaratio = 1.D0
real(rp) :: fixedsigx = 1.D-4
real(rp) :: fixedsigy = 1.D-4
real(rp), allocatable :: s_array(:), linear_density_array(:)
save
contains

!num_s = number of lines in the density file

!Calculates linear_density at position s with interpolation pattern
!determined by interp_pattern
!Key:
!interp_pattern = 1 ----- Piecewise constant
!interp_pattern = 2 ----- Linear interpolation
function linear_density_calculator(s, interp_pattern) result (lambda)
  real(rp) :: s, lambda, dis, frac, fdis
  integer :: i, ub, lb, interp_pattern
  logical :: searched
  real(rp), parameter :: tol = 1.D-14

  !Checking if s is within bounds
  !First s position in the file ought to be 0 for the program to work
  
  if ((s < 0) .or. (s > s_array(num_s))) then
    print *, 's is out of bounds'
    i = 1/0
  end if

  !Binary search
  searched = .false.
  ub = num_s
  lb = 1
  i = floor(real(1 + num_s)/2.d0)
  do
    if (s_array(i) < s) lb = i
    if (s_array(i) > s) ub = i
    if (abs(s_array(i) - s) < tol) then
      lambda = linear_density_array(i)
      searched = .true.
      exit 
    end if
    if (ub <= lb) i = 1/0 !Error: Expect lb < ub and lb /= ub
    if (ub - lb == 1) exit
    i = floor(real(lb + ub)/2.d0)
  end do

  if (searched .eqv. .false.) then
    select case (interp_pattern)
      case(1)
        lambda = linear_density_array(lb)
      case(2)
        dis = s_array(ub) - s_array(lb)
        fdis = linear_density_array(ub) - linear_density_array(lb)
        frac = (s - s_array(lb))/dis
        lambda = linear_density_array(lb) + frac*fdis
      case default
        lambda = 0.d0
    end select
  end if
end function linear_density_calculator

function return_sigmaratio() result(ratio)
real(rp) :: ratio
ratio = sigmaratio
end function return_sigmaratio

function return_sig() result(sig)
real(rp) :: sig
sig = max(fixedsigx,fixedsigy)
end function return_sig

function return_sigx() result(sigx)
real(rp) :: sigx
sigx = fixedsigx
end function return_sigx

function return_sigy() result(sigy)
real(rp) :: sigy
sigy = fixedsigy
end function return_sigy

function return_usefixedsig() result(io)
logical :: io
io = use_fixedsig
end function return_usefixedsig

subroutine initialize_num_s(pj)
integer :: pj
num_s = pj
end subroutine

subroutine allocate_arrays()
allocate(s_array(1:num_s))
allocate(linear_density_array(1:num_s))
end subroutine

subroutine initialize_s_array(hee, n)
real(rp) :: hee
integer :: n
s_array(n) = hee
end subroutine

subroutine initialize_lin_array(hep, n)
real(rp) :: hep
integer :: n
linear_density_array(n) = hep
end subroutine

subroutine initialize_sigmaratio(ratio)
real(rp) :: ratio
sigmaratio = ratio
end subroutine

subroutine initialize_fixedsigx(asd)
real(rp) :: asd
fixedsigx = asd
end subroutine

subroutine initialize_fixedsigy(asf)
real(rp) :: asf
fixedsigy = asf
end subroutine

subroutine initialize_usefixedsig(pp)
logical :: pp
use_fixedsig = pp
end subroutine

end module