program toy
  use bmad
  use beam_mod
  use mode3_mod
  implicit none

  type (lat_struct) lat
  type (ele_struct) ele 
  type (coord_struct), allocatable :: orbit(:)
  type (beam_init_struct) beam_init
  type (beam_struct) beam, beam_0
  type (normal_modes_struct) mode
  type (rad_int_all_ele_struct) rad_int

  integer unit, number, m, i, ix, nargs,j,k, n
  integer ix_cache/0/
  integer lun
  integer nturns, n_bunch, n_particle
  integer ios
  integer multiuse_counter
  integer counter

  real(rp) bunch_spacing, bunch_charge
  real(rp) tune_ex
  real(rp) bound_x

  character*19 file_name
  character*50 filename
  character*50 filename1
  character*50 filename2
  character*3 num

  character*140 lat_file
  character*120 line, last_line
  logical error/.false./

  namelist /beam_def/nturns, n_bunch, bunch_spacing, n_particle, bunch_charge

  nargs = cesr_iargc()
  if(nargs == 1)then
    call cesr_getarg(1,lat_file)
    print *, 'Using ', trim(lat_file)
  else
    lat_file = 'bmad.'
    print '(a,$)',' Lattice file name ? (default= bmad.) '
    read(5,'(a)') line
    call string_trim(line, line, ix)
    lat_file = line
    if(ix == 0) lat_file = 'bmad.'
      print *, ' lat_file = ', lat_file
  endif

  lun = lunget()
  open(unit=lun, file='beam_def.in', STATUS ='old')
  read(lun, nml=beam_def, IOSTAT=ios)
  close(unit=lun)
  write(6,nml=beam_def)

  call bmad_parser (lat_file, lat)

  call reallocate_coord (orbit, lat%n_ele_track)
  call lat_make_mat6(lat, -1)
  call twiss_at_start(lat)
  call twiss_propagate_all(lat)
  call closed_orbit_calc(lat,orbit,6)

! write closed orbit and beta functions
  open(unit=15,file='closed_orbit_traj.dat')
  open(unit=16, file='beta.dat')
  do j=1, nturns
     call track_all(lat, orbit)
     do i=1,lat%n_ele_track
        write(15,'(i4, 3x, a,es12.4,1x,6es12.4)') j, lat%ele(i)%name, lat%ele(i)%s + (j-1)*lat%ele(lat%n_ele_track)%s , orbit(i)%vec(1:6)
        write(16,'(a,es12.4,1x,3es12.4)') lat%ele(i)%name, lat%ele(i)%s, lat%ele(i)%a%beta, lat%ele(i)%b%beta, lat%ele(i)%x%eta
     end do
     orbit(0) = orbit(lat%n_ele_track)
  end do 
  close(unit=15)
  close(unit=16)




orbit(0)%vec(1) = orbit(0)%vec(1)+.0001

!write the phase space coordinates at one point for multiple k2 values
do i=0,10
   lat%ele(1)%value(k2$) =.03*i
   write(filename, '(a,i4.4,a)') 'round_trip_traj_k',i*10,'.dat'
   open(unit=15, file=filename) 
   do j=1,nturns
      call track_all(lat,orbit)
      write(15,'(i6, 3x, a7 ,es12.4,1x,2es12.4)') j, lat%ele(0)%name, lat%ele(0)%s + (j-1)*lat%ele(lat%n_ele_track)%s , orbit(0)%vec(1:2)
      orbit(0)=orbit(lat%n_ele_track)
      do k=1,lat%n_ele_track
         if (lat%ele(k)%key == sextupole$) then
            if (lat%ele(k)%value(k2$) == .03*i) then
               !print '(i4)', k
            endif
         endif
      end do
   end do
   close(unit=15)
end do


!offset trajectory
 orbit(0)%vec(1) = orbit(0)%vec(1)  + .0001

  open(unit=15,file='x_altered_traj.dat')
  do j=1,nturns
     call track_all(lat, orbit)
     do i=1,lat%n_ele_track
        write(15,'(i4, 3x, a7,es12.4,3x,6es12.4)') j, lat%ele(i)%name, lat%ele(i)%s + (j-1)*lat%ele(lat%n_ele_track)%s, orbit(i)%vec(1:6)
     end do
     orbit(0) = orbit(lat%n_ele_track)
  end do
close(unit=15)
 ! call radiation_integrals (lat, orbit, mode, ix_cache, 0, rad_int)

!Momentum of particle as a function of s in closed orbit
 open(unit=15, file='momemtum_diff_s.dat')
 do j=1, nturns
    call track_all(lat, orbit)
    do i=1, lat%n_ele_track
       write(15, '(a7,es12.4, es12.4)') lat%ele(i)%name, lat%ele(i)%s + (j-1)*lat%ele(lat%n_ele_track)%s, orbit(i)%vec(6)
    end do
    orbit(0)=orbit(lat%n_ele_track)
 end do
 close(unit=15)

! This section of the code will be for getting the tune after all of the components
open(unit=15, file = 'tune_advance.dat')
do j=1, nturns
   call track_all(lat,orbit)
   do i = 1, lat%n_ele_track
      write(15, '(a7, 3x, es12.4, 10es12.4)') lat%ele(i)%name, lat%ele(i)%s + (j-1)*lat%ele(lat%n_ele_track)%s, lat%ele(i)%a
   end do
   orbit(0)=orbit(lat%n_ele_track)
end do
close(unit=15)

! tune_ex = ((1/(2*twopi))*24.613*8.8*.01*.145) + .63



  beam_init%n_bunch = n_bunch
  beam_init%dt_bunch = bunch_spacing
  beam_init%n_particle=n_particle
  beam_init%bunch_charge = bunch_charge

  beam_init%a_emit = 2.e-9 !beam_init%a_emit = 2.e-9
  beam_init%b_emit = mode%b%emittance !beam_init%b_emit = 10.e-12
  beam_init%sig_z = mode%sig_z        !beam_init%sig_z = 0.01
  beam_init%sig_e = mode%sigE_E       !beam_init%sig_e = 0.0008

  print '(a23,i6)', 'Number of turns', nturns
  print '(a23,i8)', 'Number of elements', lat%n_ele_track
  print '(a23,es12.4)', 'Length of loop', lat%ele(lat%n_ele_track)%s
  print *, "Tune: ", lat%a%tune/twopi, lat%b%tune/twopi


 ! print *, "Expected Tune: ", tune_ex
  print *, "beta:", lat%ele(lat%n_ele_track)%a%beta
  



  print *, "Orbit:", orbit(0)%vec
  print '(a30, 10es12.5)', "Twiss (beta,alpha,gamma,phi): " , lat%ele(lat%n_ele_track)%a

  print '(2(a23,es12.4))', 'horizontal emittance = ',beam_init%a_emit, 'vertical emittance = ', &
  beam_init%b_emit
  print '(2(a23,es12.4))', 'bunch length = ',beam_init%sig_z,'sig E/E = ', beam_init%sig_e
  print '(a23,es12.4)','bunch charge =', beam_init%bunch_charge
  print '(1(a23,i10))','number of bunches =', beam_init%n_bunch
  print '(1(a23,i10))', 'particles/bunch =', beam_init%n_particle

!  call init_beam_distribution(lat%ele(0),lat%param, beam_init, beam)
!  print *, "Beam information:"
!  print *, "     " ,"bunch"," ", "t_center", "# particles"
  
!  do i=1, size(beam%bunch)
 !   beam%bunch(i)%particle(:)%vec(3) = 0.001
!  end do

! This loops tracks the beam through nturns number of turns when the value
! of the first sextupoles k2 is varied 

!The bound_x variable specifies how far out the particles go before 
! you count them as outside

bound_x = .001

do n = 0,10
   !Every time the sextupole value changes we need to initialize a new beam
   call init_beam_distribution(lat%ele(0), lat%param, beam_init, beam)
   lat%ele(1)%value(k2$) = .15*n
   print '(es12.4)', lat%ele(1)%value(k2$)
   write(filename2, '(a,i0,a)') '/nfs/user/tdt48/beam_containmentk2_', n, '.dat'
  ! print '(a)', filename2
   open(unit=16, file=filename2)
  ! print '(a)', 'here'
   do k=1,nturns
      counter = 0 
      write(filename1, '(a,a,i0,a,i0,a)') '/nfs/user/tdt48/beam_n1000/','beam_traj_n', k, 'k2_', n, '.dat'
      open(unit=15, file=filename1)
      call track_beam(lat,beam,ele1=lat%ele(1),ele2=lat%ele(lat%n_ele_track+1), err=error)
      do i=1, size(beam%bunch)
         do j=1, n_particle
            write(15,'(2i4, 3x, 2es12.4)') i, j, beam%bunch(i)%particle(j)%vec(1:2)
            
            ! This if statement is going to tell us how many particles are inside of a certain area
            ! and write the percentage of particles into the beam_containment file
            if (abs(beam%bunch(i)%particle(j)%vec(1)) <= bound_x) then
               counter = counter + 1
            endif

         end do
      end do
      write(16, '(i0, 3x, es12.4)') k, 100*(counter/n_particle)
      close(unit=15)
   end do
   close(unit=16)
end do      






!  do i=1,nturns
!    call write_turn_data(beam_init,beam,i)
!    call track_beam(lat,beam,ele1=lat%ele(0),ele2=lat%ele(lat%n_ele_track), err=error)  
!  end do

end program