program tracking_master
  use bmad
  use muon_mod
  use muon_interface
  use parameters_bmad
  use runge_kutta_mod
  use materials_mod

  implicit none

  type (lat_struct), target:: lat, low_e_lat,high_e_lat
  type (lat_struct) latb
  type (branch_struct), pointer :: branch ! ring
  type (branch_struct), pointer:: from_branch ! injection line
  type (ele_struct), pointer :: fork_ele
  type (ele_struct) ele_at_s, ele_save
  type (coord_struct), allocatable :: from_orbit(:)
  type (coord_struct), allocatable, save :: co(:)
  type (coord_struct), allocatable :: coSteps(:),co_electron(:),electrons_temp(:)
  type (coord_struct), allocatable :: low_e_orb(:), high_e_orb(:) 
  type (coord_struct) to_orbit
  type (coord_struct) start_orb, orb_at_s
  type (coord_struct) opt_orb
  type (em_field_struct) field

  integer pgopen, istat1, istat2
  integer i, j, k,l,steps/1000/               !steps: number of steps for tracking the muon decay
  integer esteps/100/,tot_track/7/,n_row      !used in the track_electron_s subroutine
  integer ix, i_dim/4/
  integer ie_from

  integer lun
  integer lun32
  integer lun_inj_branch
  integer nargs, iargc, ios
  integer end, nturns, nmuons
  integer nmu
  integer bin,bins/0/
  integer n
  integer track_state
  integer unit
  integer turn
  integer ix_ele
  integer tot_inf_scatter, tot
  integer ix_ele_b, ix_ele_c, ix_ele_k
  integer status
  integer nbranch, inj_branch/0/
  integer vparam_id/0/
  integer ix_end

  real(rp) xmean_m, xsigma_m, ymean_m, ysigma_m, tmean, tsigma, pxmean, pxsigma, pymean, pysigma, pzmean, pzsigma
  real(rp) epsx, epsy, tlength, pz
  real(rp), allocatable :: NN(:)
  real(rp) circumference,ele_len
  real(rp) distance
  real(rp) kickermagnitude_0, delta_kick
  real(rp) angle
  real(rp) l_angle_b, l_angle_c, delta_angle
  real(rp) ring_theta/0./
  real(rp) delta_e/1.e-3/,chrom_x,chrom_y
  real(rp) vparam_max_tot/0./, cbo_min/100./,vparam_cbo_min/0./
  real(rp) vparam, vparam_max, delta_vparam, vparam_min
  real(rp) mat_inv(6,6)
  real(rp) s_end
  real(rp) s1,s2
  real(rp) inflector_angle
  real(rp) s_rel,t_rel
  integer  tot_max/0/

 
  character*120 line, lat_file, new_file/' '/, muon_file/' '/
  character*16 ele_name
  logical err_flag, inf_end_us/.false./, inf_end_ds/.false./, quad_plate/.false./
  logical first/.true./, twiss_file/.false./
  logical create_new_distribution/.false./
  logical loop/.false./,gnu_input_only/.false./
  logical first_loss/.true./
  logical enerloss/.false./, quad_p/.false./
  logical spin_tracking_on/.false./,muon_decay_on/.false./
  logical local_ref/.false./,calcd/.false./  !why was this removed
  logical opt_incident/.false./,inj_matrix_tracking/.false./

  type (muon_struct), allocatable :: muons(:), muons_ele(:,:), muons_temp(:), muons_ele_inj(:,:)
  type (electron_struct),allocatable::electrons_ele(:,:),electrons_s(:,:)
  type (g2twiss_struct) twiss, twiss_opt
  type (g2moment_struct) moment, moment_ele
  type (initial_offsets_struct) initial_offsets
  type (kicker_params_struct) kicker_params
  type (quad_params_struct) quad_params

  real(rp), dimension(100) :: fnal_w_integral
  character*16 epsdistr, tdistr, pzdistr, inf_aperture

  ! Input namelist structure ("/g-2/files/input.dat")
  namelist /input/ nmuons, nturns,                         &  ! GENERAL
    create_new_distribution, new_file, muon_file,           &  ! BEAM (write and/or read from file)
    tdistr, tlength, tsigma,                               &  ! BEAM (longitudinal)
    pzdistr, pz, pzsigma,                                  &  ! BEAM (energy)
    epsdistr, epsx, epsy, twiss,                           &  ! BEAM (transverse)
    inf_aperture, inf_end_us, inf_end_ds, initial_offsets, enerloss, inflector_angle, inflector_field, &  ! INFLECTOR
    ring_theta,                                            &  ! ELEMENT POSITIONING
    kickerPlates, kickerCircuit, kickerFieldType,          &  ! KICKER
    kicker_params,                                         &  ! KICKER
    kicker_tStart,                                         &  ! KICKER
    quad_plate,                                            &  ! QUAD SCATTERING
    quadPlates, quadCircuit, quadFieldType,                &  ! QUAD
    quad_params,                                           &  ! QUAD
    twiss_file,                                            &
    loop, vparam_max, vparam_id, delta_vparam, vparam_min, &
    gnu_input_only,spin_tracking_on, muon_decay_on, inflector_width, opt_incident, &
    inj_matrix_tracking

  ! Very important
!  call ran_seed_put(123)

  ! Here are the default tracking values (pp. 139-142 of the BMAD manual, version 19.7)
!   bmad_com%rel_tol_tracking           = 1.e-6
!   bmad_com%abs_tol_tracking           = 1.e-8
!   bmad_com%rel_tol_adaptive_tracking  = 1.e-8
!   bmad_com%abs_tol_adaptive_tracking  = 1.e-8 !1.e-10 
!   bmad_com%min_ds_adaptive_tracking   = 0.
!   bmad_com%fatal_ds_adaptive_tracking = 1.e-6
!    global_com%exit_on_error=.false.

  ! Read the lattice file. The default name is "bmad." The lattice file contains the layout of the
  ! ring elements, including length, locations and apertures of kicker and quads.
  ! The lattice file also defines beam energy, field of the main dipole (in terms of bending radius), 
  ! and kicker amplitude and width

!turn on spin tracking
if (spin_tracking_on) bmad_com%spin_tracking_on = .true.

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

  ! Construct the lattice from the file
  ! There are two branches, the injection line (branch=0) from somewhere upstream to the inflector exit
  !  and the ring (branch=1)
  bmad_com%auto_bookkeeper=.false.
  bmad_com%min_ds_adaptive_tracking = 0.00001
! Scattering in inflector etc. is off while parser computes transfer matrices
  call initializeMaterials()
  call bmad_parser (lat_file, lat)
  lun=lunget()
  open(unit=lun, file="log.dat",access='append')
  write(lun,'(a16,1x,4a12,a12)')'Name','x1_limit','x2_limit','y1_limit','y2_limit','limit_on'
  do nbranch=0,1
   if(nbranch==0)write(lun,'(a)')'Injection line'
   if(nbranch==1)write(lun,'(a)')'Ring'
   do j=1,lat%branch(nbranch)%n_ele_track
    write(lun,'(a16,1x,4es12.4,5x,l)')lat%branch(nbranch)%ele(j)%name, lat%branch(nbranch)%ele(j)%value(x1_limit$), lat%branch(nbranch)%ele(j)%value(x2_limit$), &
                              lat%branch(nbranch)%ele(j)%value(y1_limit$), lat%branch(nbranch)%ele(j)%value(y2_limit$), lat%param%aperture_limit_on  
   end do
  end do
  close(unit=lun)

!  lat%ele(1:lat%n_ele_track)%alias= 'No'
!  nbranch = ubound(lat%branch,1)   !refers to the ring branch
  nbranch = 1  ! ring branch
  print *,' nbranch = ', nbranch
  ! Read the input namelist ("/g-2/files/input.dat")

  OPEN (UNIT=5, FILE='input.dat', STATUS='old', IOSTAT=ios)
    READ (5, NML=input, IOSTAT=ios)
!    WRITE(6,NML=input)
  print *, 'ios=', ios
  CLOSE(5)

  call create_gnu_input(tot_max,vparam_max_tot, cbo_min, vparam_cbo_min)

  tilt = inflector_angle
! Turn off scattering until reference orbit and transfer matrices are established
  eloss = .false.
  us_scatter = .false.
  ds_scatter = .false.
  quad_p = quad_plate
  quad_plate = .false.
!  if(index(inflector_ape,'989')/=0)inflector_width=0.018

  lun=lunget()
   open(unit=lun,file='input_data.dat')
    write(lun,nml=input)
   close(lun)
  ! Rotate ring (including all elements) with respect to inflector exit.
  ! A rotation of zero places the center of the central kicker 90 degrees from the inflector exit
  if(ring_theta /= 0.)then
    bmad_com%auto_bookkeeper=.false.
!    call rotate_ring(lat%branch(nbranch),ring_theta) !this needs to be checked as to whether pass lat% branch will work
  endif

  branch => lat%branch(nbranch)  ! this is the ring branch
  from_branch => lat%branch(0)   ! injection line branch
  ie_from = branch%ix_from_ele
  fork_ele => from_branch%ele(ie_from)


print *,' Fork_ele assigned: fork_ele%name = ', fork_ele%name

  ! Set the Twiss parameters of lat%ele(0)=BEGINNING
  ! to the values in "input.dat"

  !Note we are asssigning these twiss parameters to first element in branch nbranch, namely the ring
  ! These values are set in the input file

  branch%ele(0)%a%beta  = twiss%betax 
  branch%ele(0)%b%beta  = twiss%betay
  branch%ele(0)%a%alpha = twiss%alphax
  branch%ele(0)%b%alpha = twiss%alphay
  branch%ele(0)%x%eta   = twiss%etax 
  branch%ele(0)%x%etap  = twiss%etapx 
  branch%ele(0)%y%eta   = twiss%etay 
  branch%ele(0)%y%etap  = twiss%etapy 

  ! Set variables for determining Twiss parameters, coupling, dispersion, etc. of the ring branch

  do i=1,branch%n_ele_track-1
    if (index(branch%ele(i)%name,'QUAD')/=0) then
      ! Set the quad field indices
       if (index(branch%ele(i)%name,'1')/=0)then
         branch%ele(i)%value(field_index$) = quad_params%field_index(1)
       else     
         branch%ele(i)%value(field_index$) = quad_params%field_index(2) 
       endif  
    elseif(index(branch%ele(i)%name,'KICKER')/= 0) then
      ! Momentarily set the kicker fields to zero (see below)
      branch%ele(i)%value(kicker_field$) = 0.
    endif
  end do

  ! Calculate the ring's Twiss parameters, coupling, and dispersion parameters
  call closed_orbit_calc(lat, co, i_dim=4, ix_branch=nbranch)
  call lat_make_mat6(lat,-1,co, ix_branch=nbranch)
  call twiss_at_start(lat,status, ix_branch = nbranch) !if initial twiss values are set in input file comment out this call
  if(status/=0)print '(a,a)',' Status from twiss_at_start = ', matrix_status_name(status)
  call twiss_propagate_all(lat, ix_branch=nbranch)
!  call chrom_calc(lat, delta_e,chrom_x,chrom_y, err_flag, low_e_lat, high_e_lat, low_e_orb, high_e_orb, ix_branch=nbranch)
!  write(6,'(a,4a12)')'#','Low E Qx','Low E Qy','High E Qx','High E Qy'
!  write(6,'(a,4es12.4)')'#',low_e_lat%a%tune/twopi, low_e_lat%b%tune/twopi,high_e_lat%a%tune/twopi, high_e_lat%b%tune/twopi
!  chrom_x = (high_e_lat%a%tune/twopi-low_e_lat%a%tune/twopi)/(2*delta_e)
!  chrom_y = (high_e_lat%b%tune/twopi-low_e_lat%b%tune/twopi)/(2*delta_e)
!  print '(2(a,es12.4))',' chrom_x = ', chrom_x,'chrom_y = ', chrom_y

!  if(err_flag)print *,' Error from chrom calc'
print *,' here'
  if (twiss_file) then
      call twiss_propagate_cm(lat,nbranch)
      call write_transfermatrix(lat, nbranch)
  endif
vparam = 0
do while (vparam <= (vparam_max-vparam_min))
   if(loop .and.vparam_id /= 0)then
    if(vparam_id == 1) twiss%betax = vparam +vparam_min
    if(vparam_id == 2) twiss%betay = vparam +vparam_min 
    if(vparam_id == 3) twiss%etax = vparam +vparam_min
    if(vparam_id == 4) kicker_params%kicker_field(1:3) = vparam +vparam_min
    if(vparam_id == 5) kicker_params%kick_width(1:3) = vparam +vparam_min
    if(vparam_id == 6) initial_offsets%x_mean = vparam +vparam_min
    if(vparam_id == 7) initial_offsets%pxmean = vparam +vparam_min
    if(vparam_id == 8) kicker_params%dtRise(1:3) = vparam +vparam_min
    if(vparam_id == 9) kicker_params%dtFall(1:3) = vparam +vparam_min
   endif

  print *, ' BETA x =', twiss%betax
  print *, 'kicker_field=',  kicker_params%kicker_field(1)

  ! Set kicker parameters (now that ring's Twiss parameters, etc., have been calculated)
  do i=1,branch%n_ele_track-1
    if( index(branch%ele(i)%name,'KICKER')/= 0 ) then
      if( index(branch%ele(i)%name,'1')/=0 ) then
        ! KICKER1
        branch%ele(i)%value(kicker_field$) = kicker_params%kicker_field(1)
        branch%ele(i)%value(kick_width$) = kicker_params%kick_width(1)
!        branch%ele(i)%value(custom_attribute4$)= kicker_params%dtRise(1)                              attribute4$ is used for kicker_field
        branch%ele(i)%value(dtRise_and_fall$)= kicker_params%dtRise(1)                                !dtRise_and_fall$ = custom_attribute5$
      else if( index(branch%ele(i)%name,'2')/=0 ) then
        ! KICKER2
        branch%ele(i)%value(kicker_field$) = kicker_params%kicker_field(2)
        branch%ele(i)%value(kick_width$) = kicker_params%kick_width(2)
 !       branch%ele(i)%value(custom_attribute4$)= kicker_params%dtRise(2)
        branch%ele(i)%value(dtRise_and_fall$)= kicker_params%dtRise(2)
      else
        ! KICKER3
        branch%ele(i)%value(kicker_field$) = kicker_params%kicker_field(3)
        branch%ele(i)%value(kick_width$) = kicker_params%kick_width(3)
 !       branch%ele(i)%value(custom_attribute4$)= kicker_params%dtRise(3)
        branch%ele(i)%value(dtRise_and_fall$)= kicker_params%dtRise(3)
      endif
    end if
  end do

  ! Print the values read from the namelist
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Beam::Longitudinal"
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,'(a,6a12)')   "    ", "tlength(ns)", "dP/P(%)", "What", "else", "???"
  WRITE(*,'(a)')        "----------------------------------------------------------------------------"
  WRITE(*,'(a,6f12.6)') "    ", tlength*1.e9, pz*1.e2
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Beam::Transverse(InflectorMidpoint)"
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,'(a,6a12)')   "    ", "beta(m)", "alpha ", "eta(m)", "etap  ", "phi(rad)", "gamma(1/m)"
  WRITE(*,'(a)')        "----------------------------------------------------------------------------"
  WRITE(*,'(a,6f12.6)') " X |", twiss%betax, twiss%alphax, twiss%etax, twiss%etapx, twiss%phix, twiss%gammax
  WRITE(*,'(a,6f12.6)') " Y |", twiss%betay, twiss%alphay, twiss%etay, twiss%etapy, twiss%phiy, twiss%gammay
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Offsets(InflectorExit)"
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,'(a,6a12)')   "    ","x(m)", "y(m)", "t(ns)", "x'(rad)", "y'(rad)", "dP/P"
  WRITE(*,'(a)')        "----------------------------------------------------------------------------"
  WRITE(*,'(a,6es12.4)')"    ",initial_offsets%x_mean, initial_offsets%y_mean, initial_offsets%tmean, initial_offsets%pxmean, initial_offsets%pymean, initial_offsets%pzmean
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Kickers"
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,'(a,6a12)')   "   |", "By(Gauss)", "twidth(ns)", "What", "else", "???"
  WRITE(*,'(a)')        "----------------------------------------------------------------------------"
  WRITE(*,'(a,6f12.1)') " K1|", kicker_params%kicker_field(1)*1.e4, kicker_params%kick_width(1)*1.e9
  WRITE(*,'(a,6f12.1)') " K2|", kicker_params%kicker_field(2)*1.e4, kicker_params%kick_width(2)*1.e9
  WRITE(*,'(a,6f12.1)') " K3|", kicker_params%kicker_field(3)*1.e4, kicker_params%kick_width(3)*1.e9
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Quads"
  WRITE(*,'(a)')        "============================================================================"
  WRITE(*,'(a,6a12)')   "   |", "Field index", "?", "What", "else", "???"
  WRITE(*,'(a)')        "----------------------------------------------------------------------------"
  WRITE(*,'(a,6f12.1)') " Q1|", quad_params%field_index(1)
  WRITE(*,'(a,6f12.1)') " Q2|", quad_params%field_index(2)
  WRITE(*,'(a,6f12.1)') " Q3|", quad_params%field_index(3)
  WRITE(*,'(a,6f12.1)') " Q4|", quad_params%field_index(4)
  WRITE(*,'(a)')        "============================================================================"
  do j = 0,nbranch
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a,i10)')        "@Lattice branch = ", j
  WRITE(*,'(a)')        "================================================================================"
  WRITE(*,'(7a16)')     "ELEMENT NAME", "r(m)", "s(m)", "angle(deg)", "total(deg)", "field_index"
  WRITE(*,'(a)')        "--------------------------------------------------------------------------------"
  angle = 0
  DO i=0, lat%branch(j)%n_ele_track
    angle = angle + lat%branch(j)%ele(i)%value(angle$)
    WRITE(*,'(a16,5f16.3)') lat%branch(j)%ele(i)%name, lat%branch(j)%ele(i)%value(rho$), lat%branch(j)%ele(i)%s, lat%branch(j)%ele(i)%value(angle$)*180./pi, angle*180./pi, lat%branch(j)%ele(i)%value(field_index$)
  END DO
  end do !branch
  WRITE(*,'(a)')        "--------------------------------------------------------------------------------"
  WRITE(*,'(a,7(a,f8.5,3x))') "RING PARAMS:     ", "Qx =", branch%ele(branch%n_ele_track)%a%phi/twopi, "Qy =", branch%ele(branch%n_ele_track)%b%phi/twopi, &
    "betax =", branch%ele(branch%n_ele_track)%a%beta, "betay =", branch%ele(branch%n_ele_track)%b%beta, "field_index =", branch%ele(branch%n_ele_track)%b%phi/twopi **2 , &
    " Q'x =",chrom_x," Q'y=",chrom_y
  WRITE(*,'(a)')        "================================================================================"
  WRITE(*,*)
  WRITE(*,*)


  lat%param%aperture_limit_on = .true.
  branch%param%aperture_limit_on = .true.
  circumference = branch%ele(branch%n_ele_track)%s

! branch => lat%branch(1)

  call reallocate_coord (from_orbit, from_branch%n_ele_max)
!  call init_coord(from_orbit)
  
!  allocate(co_electron(nmuons))
  call reallocate_coord (co_electron, nmuons)
  call reallocate_coord (co, branch%n_ele_max)
  call reallocate_coord (coSteps, steps)
  call init_coord(co(0))
  call init_coord(coSteps(0))

  start_orb%vec    = 0
  start_orb%t      = initial_offsets%tmean
  start_orb%s      = 0
  start_orb%vec(1) = initial_offsets%x_mean
  start_orb%vec(2) = initial_offsets%pxmean
  start_orb%vec(3) = initial_offsets%y_mean
  start_orb%vec(4) = initial_offsets%pymean
  start_orb%vec(5) = initial_offsets%tmean
  start_orb%vec(6) = initial_offsets%pzmean
  co(0) = start_orb
  print '(a,6es12.4)','start_orb',co(0)%vec

  lat%ele(1:lat%n_ele_track)%alias= 'Record'  !write magnetic field and position along trajectory to unit 12 in em_field_custom_inj.f90
  call track_all(lat,co, ix_branch = 0, track_state = track_state, err_flag = err_flag)  !compute trajectory through injection line
  lat%ele(1:lat%n_ele_track)%alias= 'No'
 
  if(opt_incident)call optimize_incident(start_orb, lat, opt_orb) !routine to find indicent offset and angle to minimize offset and angle at inflector exit

  print *,' trajectory through backleg and inflector'
  open(unit=49)
  do i=1,lat%branch(0)%n_ele_track
   print '(7es12.4,1x,a)',lat%ele(i)%s,co(i)%vec(1:6), lat%ele(i)%name 
   write(49,'(7es12.4,1x,a)')lat%ele(i)%s,co(i)%vec(1:6), lat%ele(i)%name 
  end do
  close(unit=49)
  call lat_make_mat6(lat, -1, co, ix_branch=0)  ! Create transfer matrices for injection line elements referred to trajectory
  call get_branch_inv_matrix(lat,mat_inv, inj_branch)  ! Multiply matrices for injection line together then compute inverse to pass to CREATE_PHASE_SPACE

  call create_phase_space(nmuons, create_new_distribution, new_file, muon_file, &
          epsdistr, epsx, epsy, tdistr, tlength, tsigma, pzdistr, pz, pzsigma, twiss, inf_aperture, inf_end_us, inf_end_ds,enerloss, muons, mat_inv)  !mat_inv propagates from end of inflector back to start of injection line
!  call compute_beam_params(muons, 1, 'PropagatedBackToStart')

do j = 1,nmuons
if (muons(j)%coord%state /= alive$) cycle
!print *, 'alive$',alive$,'downstream',downstream_end$
!print *, 'state,location,direction',muons(j)%coord%state,muons(j)%coord%location,muons(j)%coord%direction
end do

 !turn on scattering now that reference orbit and transfer matrices through injection line are computed

  eloss = enerloss
  us_scatter = inf_end_us 
  ds_scatter = inf_end_ds
  quad_plate = quad_p

  runge_kutta_com%check_limits = .true.
  runge_kutta_com%check_wall_aperture = quad_plate !if true then scatter in plates
print '(a,1x,l)', 'runge_kutta_com%check_wall_aperture = ', runge_kutta_com%check_wall_aperture

  runge_kutta_com%hit_when = wall_transition$
  
  n_row = nmuons/20
  allocate(muons_ele(0:nmuons, 0:branch%n_ele_track))
  allocate(muons_ele_inj(0:nmuons, 0:lat%branch(0)%n_ele_track))
  allocate(muons_temp(0:nmuons))
  allocate(electrons_ele(0:nmuons,0:branch%n_ele_track))
  allocate(electrons_temp(0:branch%n_ele_track))
  allocate(electrons_s(0:n_row,0:esteps*(tot_track+1))) 

  do j = 1,nmuons
  electrons_ele(j,0)%exist =.false.
  electrons_ele(j,0)%lost  =.false.
  end do

 from_orbit(0)%t=0.
 from_orbit(0)%s=0.
 from_orbit(0)%vec=0.

 forall (n=1:nmuons) muons(n)%coord%vec(1:6) = muons(n)%coord%vec(1:6) + from_orbit(0)%vec(1:6)
 muons(:)%coord%t = muons(:)%coord%t + from_orbit(0)%t
 muons(:)%coord%s = muons(:)%coord%s + from_orbit(0)%s

! muons(:)%lost = .false.
! unit=24
! call write_phase_space(unit,nmuons,muons,'Cut, scattered distribution after inflector -- start of tracking', tot_inf_scatter)

 call init_moment(moment)
 call init_moment(moment_ele)
 call compute_moments(muons,0.d0,moment, nmuons,nmu, 16, 'BEGINNING',0)
 call compute_moments(muons,-lat%branch(0)%ele(lat%branch(0)%n_ele_track)%s/circumference,moment, nmuons,nmu, 26,'BEGINNING',0)

end=lat%n_ele_track
    i=0
     if(nmuons == 1)then
        write(19,'(a10,1x,7a12)')' turn ','x','xp','y','yp','z','pz','t'
        write(19,'(i10,1x,7es12.4)')i,muons(1)%coord%vec, muons(1)%coord%t
        write(29,'(a10,1x,9a12)')' turn ','x','xp','y','yp','z','pz','t','s','s total'
        write(29,'(i10,1x,a16,1x,9es12.4)')i,lat%branch(0)%ele(0)%name,muons(1)%coord%vec, muons(1)%coord%t, &
                                              lat%ele(0)%s/circumference, &
                                    ((i-1)*lat%ele(lat%n_ele_track)%s +lat%ele(0)%s)/circumference
        write(71,'(a12,1x,6a12)') 's','x','xp','y','yp','z','pz'
        write(71,'(es12.4,1x,6es12.4)') lat%branch(0)%ele(0)%s,start_orb%vec(1:6)
        write(72,'(7a12,1x,a16)')'s', 'beta x','beta y','alpha x','alpha y', 'eta x','etap x', 'element'
        write(72,'(7es12.4,1x,a16)')lat%branch(0)%ele(0)%s, lat%branch(0)%ele(0)%a%beta,lat%branch(0)%ele(0)%b%beta,lat%branch(0)%ele(0)%a%alpha,lat%branch(0)%ele(0)%b%alpha, &
                                          lat%branch(0)%ele(0)%x%eta,lat%branch(0)%ele(0)%x%etap,lat%branch(0)%ele(0)%name
     endif

 branch%ele(1:branch%n_ele_track)%alias= 'trajectory'
 lat%branch(0)%ele(1:lat%branch(0)%n_ele_track)%alias = 'trajectory'

DO i=1, nturns
  IF (i>1) runge_kutta_com%check_wall_aperture = .false.
  if(i == 1 ) then !propagate through injection line to start of ring
   DO n=1, nmuons   !propagate all muons through branch 0
    muons_ele(n,:)%coord%state = lost$
    IF (muons(n)%coord%state /= alive$) CYCLE 
      from_orbit(0)%vec = muons(n)%coord%vec
      from_orbit(0)%t = muons(n)%coord%t
      
!      call track_all(lat,from_orbit, ix_branch = 0, track_state = track_state, err_flag = err_flag)
       do j=1,lat%branch(0)%n_ele_track
        ele_save%tracking_method = lat%branch(0)%ele(j)%tracking_method
        if(inj_matrix_tracking .and.lat%branch(0)%ele(j)%key /= marker$)lat%branch(0)%ele(j)%tracking_method = linear$
        call track1(from_orbit(j-1),lat%branch(0)%ele(j), lat%param, from_orbit(j))
        if(j < lat%branch(0)%n_ele_track)then
            if(lat%branch(0)%ele(j)%name == 'BACKLEG_START' )then
              from_orbit(j)%vec = from_orbit(j)%vec + start_orb%vec
            endif
        endif
        lat%branch(0)%ele(j)%tracking_method = ele_save%tracking_method

        muons_ele_inj(n,j)%coord = from_orbit(j)
       end do

      IF (nmuons==1)then  !if there is only one muon then compute its trajectory and twiss parameters along injection line
                call Bfield_along_axis
                do j=1,lat%branch(0)%n_ele_track
                    write(29,'(i10,1x,a16,1x,9es12.4,i10)')i,lat%branch(0)%ele(j)%name,from_orbit(j)%vec, co(j)%t, &
                                                          lat%branch(0)%ele(j)%s/circumference, &
                                    ((i-1)*lat%branch(0)%ele(lat%branch(0)%n_ele_track)%s +lat%branch(0)%ele(j)%s)/circumference, from_orbit(j)%state         
                call lat_make_mat6 (lat, j,from_orbit, ix_branch=0, err_flag=err_flag)
                end do

                call twiss_propagate_all(lat, ix_branch=0)

                call optimize_twiss(lat, twiss, twiss_opt) !twiss_opt are the twiss parameters at the upstream end of the injection line
                                                                                  ! that yield the desired twiss parameters in the inflector (namely twiss) 
 
                s_end=0.00001
                orb_at_s%vec=from_orbit(0)%vec
                ix_end=1
                write(71,'(es12.4,1x,6es12.4,1x,6es12.4,1x,a16,1x,i10)')s_end, orb_at_s%vec, from_orbit(1)%vec,lat%ele(1)%name, ix_end !end_orb%vec
                do while(s_end < lat%ele(lat%branch(0)%n_ele_track)%s +0.05) !the 0.05 is to make sure we overlap the very end
                   if(s_end > lat%ele(lat%branch(0)%n_ele_track)%s) s_end = lat%ele(lat%branch(0)%n_ele_track)%s !track to the very end
                   call twiss_and_track_at_s(lat, s_end, ele_at_s, from_orbit, orb_at_s,ix_branch=0)          
                   ix_end = element_at_s (lat, s_end, .true.)
                    s_rel = s_end - lat%ele(ix_end-1)%s
                   call em_field_custom(lat%ele(ix_end), lat%param, s_rel,t_rel, orb_at_s, local_ref, field,calcd, err_flag)
                   write(73,'(es12.4,1x,6es12.4,3es12.4)')s_end, orb_at_s%vec, field%B(1:3)
                   write(71,'(es12.4,1x,6es12.4,1x,6es12.4,1x,a16,1x,i10)')s_end, orb_at_s%vec, from_orbit(ix_end)%vec,lat%ele(ix_end)%name, ix_end !end_orb%vec
                   write(72,'(7es12.4,1x,a16)')ele_at_s%s, ele_at_s%a%beta, &
                   ele_at_s%b%beta, ele_at_s%a%alpha, ele_at_s%b%alpha, &
                   ele_at_s%x%eta,  ele_at_s%x%etap, ele_at_s%name
                   s_end = s_end + 0.1
                end do

      endif  !nmuons = 1

      to_orbit = from_orbit(ie_from)
      muons(n)%coord%vec = to_orbit%vec
      muons(n)%coord%t =  to_orbit%t
      muons(n)%coord = to_orbit

!     generate the same spin for every muon
    if (spin_tracking_on)      muons(n)%coord%spin = (/(1.,0.),(0.,0.)/)

    end do ! end of loop propagating all muons through branch 0

!     generate Decay time(lifetime)
      if(muon_decay_on)call compute_decayTime(muons(1:nmuons),nmuons,11)

!    call compute_beam_params(muons, 1, 'EndInfAfterScatter')
    lun=37
    call write_phase_space(lun, nmuons, muons,'Distribution at end of inflector after scattering', tot)
    unit=24
    call write_phase_space(unit,nmuons,muons,'Cut, scattered distribution after inflector -- start of tracking', tot_inf_scatter)

   endif  !for first turn only

  do n=1, nmuons  
    muons_ele(n,:)%coord%state = lost$
    IF (muons(n)%coord%state /= alive$) CYCLE

    co(0)%vec = muons(n)%coord%vec
    co(0)%t   = muons(n)%coord%t
    co(0) = muons(n)%coord

    call track_all(lat, co, nbranch, track_state, err_flag)

    DO j=1, branch%n_ele_track
      muons_ele(n,j)%coord = co(j)
      IF (nmuons == 1) write(29,'(i10,1x,a16,1x,9es12.4,i10,i10)')i,branch%ele(j)%name,co(j)%vec, co(j)%t, &
                                              branch%ele(j)%s/circumference, &
                                    ((i-1)*branch%ele(branch%n_ele_track)%s +branch%ele(j)%s)/circumference,co(j)%state,n

      if (err_flag) muons_ele(n,j)%lost = .true.
      if (j == track_state .and. track_state /= moving_forward$)then
        if(first_loss) write(33,'(a10,1x,8a12)')'muon','s','x','xp','y','yp','z','pz','element','turn'
           write(33,'(i10,1x,7es12.4,1x,a12,1x,i10)')n,branch%ele(j)%s, co(j)%vec, branch%ele(j)%name,i
           first_loss = .false.
        do k=j,branch%n_ele_track
          muons_ele(n,k)%coord%state = lost$
        enddo
      endif

    ENDDO ! track through lattice elements


    if (track_state /= moving_forward$ .or. err_flag) then
      muons(n)%coord%state = co(track_state)%state
      muons(n)%lost = .true.
      cycle
    endif


!***********************************************
   !see if muon has decayed. If so,retrack to see where it decayed and obtain the corresponding electron phase space vector

    if(muon_decay_on) then
    
     if (co(branch%n_ele_track)%t >= co(0)%path_len) then

      do j = 1, branch%n_ele_track

       if (co(j)%t >= co(0)%path_len) then
   
         coSteps(0) = co(j-1)
         ele_len = branch%ele(j)%s - branch%ele(j-1)%s
         do l = 0,steps-1
           s1 = branch%ele(j-1)%s + l*ele_len/steps
           s2 = s1 + ele_len/steps
    
           if (ele_len /= 0) s1 = s1 + 0.00001 
           if (ele_len /= 0) s2 = s2 - 0.00001

           call track_from_s_to_s(lat,s1,s2,coSteps(l),coSteps(l+1),ix_branch = nbranch)
           ! print *, 's1',s1,'s2',s2, 't',coSteps(l)%t, 'lifetime=',co(0)%path_len    

           if (coSteps(l)%t >= co(0)%path_len) then 
             call decayElectron(coSteps(l),co_electron(n))
             distance = float(i-1) + branch%ele(j-1)%s/circumference + l/steps*ele_len/circumference 
             call decay_info(41,distance,n,co_electron(n),coSteps(l))    !record decay information
            ! call track_electron_ele(lat,nbranch,s1,co_electron(n),j,electrons_ele,n,nmuons,branch%n_ele_track,.false.)
            call track_electron_s(lat,nbranch,s1,co_electron(n),j,electrons_s(0:n_row,0:esteps*(1+tot_track)),n,nmuons,esteps,tot_track,n_row,esteps*(1+tot_track),branch%n_ele_track,.false.)

            exit       !exit steps
           end if     !if costeps(l)%t >= co(0)%path_len
    
     
           if (l == steps-1) then    !if decayed at the end
             call decayElectron(coSteps(steps),co_electron(n))           
             distance = float(i-1)+branch%ele(j)%s/circumference
             call decay_info(41,distance,n,co_electron(n),coSteps(steps))  
             !call track_electron_ele(lat,nbranch,s1,co_electron(n),j,electrons_ele,n,nmuons,branch%n_ele_track,.true.)
             call track_electron_s(lat,nbranch,s1,co_electron(n),j,electrons_s(0:n_row,0:esteps*(1+tot_track)),n,nmuons,esteps,tot_track,n_row,esteps*(1+tot_track),branch%n_ele_track,.true.)

           end if     !if l == steps-1

         end do     !  steps

         exit       ! exit element search
       end if     ! if decay in element j

      end do      ! element search j=1,branch%n_ele_track

      muons(n)%coord%state =  lost$
      muons(n)%lost = .true.
      cycle
     end if      ! if decay at the end      if (co(branch%n_ele_track)%t >= co(0)%path_len) then
    endif !muon_decay_on
!**********************************************
    
    muons(n)%coord%vec = co(branch%n_ele_track)%vec
    muons(n)%coord%t   = co(branch%n_ele_track)%t
    muons(n)%coord   = co(branch%n_ele_track)
  ENDDO ! n muons

 if(i == 1)then !first turn
  
  do j=1,lat%branch(0)%n_ele_track
     distance = float(i-1)-(lat%branch(0)%ele(lat%branch(0)%n_ele_track)%s-lat%branch(0)%ele(j)%s)/circumference
     forall(n=1:nmuons)muons_temp(n) = muons_ele_inj(n,j)
     call compute_moments(muons_temp,distance,moment_ele, nmuons,nmu, 26, lat%branch(0)%ele(j)%name, j)
     lun_inj_branch = lunget()
     open(unit = lun_inj_branch, file = trim(lat%branch(0)%ele(j)%name)//'phase_space.dat')
     call write_phase_space(lun_inj_branch,nmuons,muons_temp, lat%branch(0)%ele(j)%name, tot) !write phase space after each element in injection line
     close(unit=lun_inj_branch)

     if (spin_tracking_on) call compute_spins(muons_temp(1:nmuons),distance,nmuons,38, branch%ele(j)%name, j)                              !compute and write average spins for element
     call moment_range(moment_ele)
  end do
 endif
 
  do j=1,branch%n_ele_track
     distance = float(i-1)+branch%ele(j)%s/circumference
     forall(n=1:nmuons)muons_temp(n) = muons_ele(n,j)
     call compute_moments(muons_temp,distance,moment_ele, nmuons,nmu, 26, branch%ele(j)%name, j)
      if (spin_tracking_on) call compute_spins(muons_temp(1:nmuons),distance,nmuons,38, branch%ele(j)%name, j)                              !compute and write average spins for elements
     if(nturns - i < 100)call moment_range(moment_ele)
  end do

                                                                                       
  distance = float(i)
  call compute_moments(muons,distance,moment, nmuons,nmu, 16, branch%ele(branch%n_ele_track)%name, j-1)                !j somehow becomes n_ele_track+1 at the end of the loop
   if (spin_tracking_on) call compute_spins(muons(1:nmuons),distance,nmuons,39, branch%ele(branch%n_ele_track)%name, j-1)                     !compute and write average spins

  if(nturns - i < 100)call moment_range(moment)  !compute range of moments only for last 100 turns
  call collect_times(muons, nturns, NN, bins)

  if (nmuons==1) write(19,'(i10,1x,7es12.4)') i, muons(1)%coord%vec, muons(1)%coord%t

ENDDO ! nturns


if(muon_decay_on)call electron_info_s(electrons_s,n_row,esteps*(1 + tot_track),42)

  lun = lunget()
  open(unit = lun, file = 'muons_vs_time.dat')
 print '(2a)',' write binned time data to ','muons_vs_time.dat'
 print *,' bins = ', bins
 write(lun,'(a10,a12)')'Bin','NN(bin)'
 write(lun,'(i10,es12.4)')(bin, NN(bin), bin=1, bins)
 close(lun) 

 unit = 25
 call write_phase_space(unit,nmuons,muons,'End of tracking', tot)

print '(/,5a,/)',' Use plotting script ',"'",'sigma.gnu',"'",' to plot turn-by-turn average position and size'

if (first)then
  lun32 = lunget()
  open(unit = lun, file = 'VparamDependence.dat')
   write(lun32,'(16a16,9a12)')  &
  'K1[Guass]-1', 'K2[Gauss]-2', 'K3[Gauss]-3', &
  'depth<x>-4',' depth<xp>-5','depth<y>-6','depth<yp>-7','depth<t>-8','depth<pz>-9', &
  'depth_SIG(1,1)-10','depth_SIG(3,3)-11','depth_SIG(5,5)-12','depth_SIG(6,6)-13', &
  'n muons-14',' after_inf-15','remaining-16','beta x-17','eta x-18','kicker_field-19','beta y-20',' kickWid-21','dtRise-22','dtFall-23','x init','px init'
first=.false.
endif !open lun32

write(lun32,'(3es16.3,10es16.6,3i16,es12.4,es12.4,5es12.4, 2es12.4)') &
  kicker_params%kicker_field(1) * 1.e4, & ! Gauss
  kicker_params%kicker_field(2) * 1.e4, & ! Gauss
  kicker_params%kicker_field(3) * 1.e4, & ! Gauss
  (moment%max_ave-moment%min_ave)                          *1.e3, & ! mm or mrad
  (sqrt(moment%max_sigma(1,1))-sqrt(moment%min_sigma(1,1)))*1.e3, & ! mm or mrad
  (sqrt(moment%max_sigma(3,3))-sqrt(moment%min_sigma(3,3)))*1.e3, & ! mm or mrad
  (sqrt(moment%max_sigma(5,5))-sqrt(moment%min_sigma(5,5)))*1.e3, & ! mm or mrad
  (sqrt(moment%max_sigma(6,6))-sqrt(moment%min_sigma(6,6)))*1.e3, & ! mm or mrad
  nmuons, &
  tot_inf_scatter, &
  tot, twiss%betax, twiss%etax, maxval(kicker_params%kicker_field(1:3)), twiss%betay, kicker_params%kick_width(1),&
  kicker_params%dtRise(1), kicker_params%dtFall(1), &
  initial_offsets%x_mean, initial_offsets%pxmean

  if(tot_max < tot)then
    tot_max = tot
    vparam_max_tot = vparam + vparam_min
  endif
  if(cbo_min >( moment%max_ave(1)-moment%min_ave(1))*1.e3)then
    cbo_min = (moment%max_ave(1)-moment%min_ave(1))*1.e3
    vparam_cbo_min = vparam + vparam_min
  endif
  call create_gnu_input(tot_max,vparam_max_tot, cbo_min,vparam_cbo_min)
  deallocate(muons)
  deallocate(muons_ele)
  deallocate(muons_temp)
  deallocate(muons_ele_inj)
  deallocate(electrons_ele)
  deallocate(electrons_temp)
  deallocate(electrons_s)

close(16)
close(26)
! twiss%etax = twiss%etax + 0.2
! twiss%betax = twiss%betax + 0.2
! kicker_params%kicker_field(2) = kicker_params%kicker_field(2) + 7.5e-4
  vparam = vparam + delta_vparam
 if(.not. loop .or. vparam_id <= 0) exit
end do
close(unit=lun32)
end