program ion_tracker

use bmad
use beam_mod
use ion_tracker_mod
use ion_tracker_struct

implicit none

type(lat_struct), target :: lat
type(branch_struct), pointer :: branch
type (ele_struct), pointer :: ele
type (coord_struct), allocatable :: orbit(:)
type(beam_struct) :: beam
type(beam_init_struct) :: beam_init
real(rp) :: ds, tmp, sigma_ratio, fixed_sigx, fixed_sigy
real(rp), allocatable :: s_array1(:), linear_density_array1(:)
integer :: n_particle, version, ele_ix, filetype, ns
integer :: i, jpjp
logical :: new_file, error_flag, truth, manual_file
logical :: use_fixed_sig

integer, parameter :: namelist_file = 11

character(20) :: state_str
character(100) :: in_file, lat_name, out_name, init_name, init_out
character(100) :: density_file


! --------------------------------------------------
!
! Specify parameters in namelist file
!
! --------------------------------------------------

! Notes:
! One cannot change the parameters of track1_custom
! This means that any additional parameters that need to be passed
! must be entered through unused members in lat%param or
! through ion_tracker_struct.
!
! Make sure to delete or move the output every time you do a new run
! of the program since it will not delete the previous output but
! simply append the existing file. 
!
! Make the last s value of the density_file a little larger than the
! max s value of the lattice. The program will give an error if the
! s value passed to linear_density_calculator is less than 0 or greater than
! the last s value in the density_file.
!
! fixed_sig options available for testing purposes. They should not be used for
! ordinary use.

in_file = 'cbeta_ions.in'       !Filename of namelist
density_file = 'ion_density.in' !Filename of charge density file
truth = .true.
manual_file = .false.           !Whether you use charge density file or not
sigma_ratio = 1.D0              !Ratio of sigma of electron beam and sigma of
                                !ions
use_fixed_sig = .false.         !Determines whether to use fixed values for the
                                !sigmas of ions (Keeo .false. if not testing)
fixed_sigx = 1.D-4              !The fixed sigx value
fixed_sigy = 1.D-4              !The fixed sigy value


namelist / cbeta_ions_params / lat_name, beam_init, version, ds, manual_file, &
            sigma_ratio, fixed_sigx, fixed_sigy, use_fixed_sig

!Initialize from namelist and charge density file.

call cbeta_ions_header()

print *, 'Reading namelist file: ', trim(in_file)
open (namelist_file, file = in_file, status = "old")
read (namelist_file, nml = cbeta_ions_params)
close (namelist_file)
print *, 'Namelist file successfully read!'

jpjp = lunget()
print *, 'Reading ion linear density from ion_density.in'
open (jpjp, file = density_file, status = 'old')
read (jpjp, '(i19)') ns
call initialize_num_s(ns)
call initialize_sigmaratio(sigma_ratio)
call initialize_usefixedsig(use_fixed_sig)
call initialize_fixedsigx(fixed_sigx)
call initialize_fixedsigy(fixed_sigy)
call allocate_arrays()

allocate(s_array1(1:ns))
allocate(linear_density_array1(1:ns))
do i = 1, ns
  read(jpjp, '(2es19.10)') s_array1(i), linear_density_array1(i)
end do
close(jpjp)

do i = 1, ns
  call initialize_s_array(s_array1(i), i)
  call initialize_lin_array(linear_density_array1(i), i)
end do

! --------------------------------------------------
!
! Main program begins
!
! --------------------------------------------------

!Parsing lattice
print *, 'Parsing lattice:', lat_name
call bmad_parser (lat_name, lat)
branch => lat%branch(0)

!Initialize
init_name = 'init_beam.out'
init_out = 'read_init_beam.out'
new_file = .true.
filetype = ascii$
n_particle = beam_init%n_particle !Particles per bunch
ele => lat%ele(0)

!Need twiss at start to propagate twiss parameters throughout lattice
if (lat%param%geometry == closed$) call twiss_at_start(lat)

!If manual_file is true, initialize beam from init_name
!If manual_file is false, initialize beam from init_beam (namelist)
if (manual_file) then
call read_beam_file(init_name, beam, beam_init, error_flag)
else
call init_beam_distribution(ele, lat%param, beam_init, beam)
end if
!Write the initial beam distribution to init_out to manually check
!if beam was initialized correctly
call write_beam_file(init_out, ele, beam, new_file, filetype)
call lat_compute_ref_energy_and_time(lat, error_flag)
!If version == 0, then it is bmad standard tracking without ions
if (version > 0) then
  call turn_track1_custom_on_for_all_elements(lat, branch)
end if
!lat%param%ixx can be accessed in track1_custom (lat%param is passed)
lat%param%ixx = version
call special_track_beam_through_ions(lat, lat%param, beam_init, beam, error_flag, ds, version)

end program