use bmad
  use muon_mod
  use muon_interface
  use parameters_bmad
  use runge_kutta_mod

  implicit none

  type (lat_struct) lat, low_e_lat,high_e_lat
  type (coord_struct), allocatable, save :: co(:)
  type (coord_struct), allocatable :: low_e_orb(:), high_e_orb(:) 

  integer pgopen, istat1, istat2
  integer i, j, k
  integer ix, i_dim/4/

  integer lun
  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

  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
  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
  character*120 line, lat_file, new_file/' '/, muon_file/' '/
  character*16 ele_name
  logical err_flag, inf_end_us/.true./, inf_end_ds/.true./, quad_plate/.true./
  logical first/.true./, twiss_file/.false./
  logical create_new_distribution/.false./

  type (muon_struct), allocatable :: muons(:), muons_ele(:,:), muons_temp(:)
  type (g2twiss_struct) twiss
  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, &  ! 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

  ! 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-8

  ! Here are some more appropriate values, I think.  (Q: What is the absolute
  ! smallest bmad_com%min_ds_adaptive_tracking we can get away with?)
!  bmad_com%rel_tol_tracking           = 1.e-8
!  bmad_com%abs_tol_tracking           = 1.e-10
!  bmad_com%rel_tol_adaptive_tracking  = 1.e-10
!  bmad_com%abs_tol_adaptive_tracking  = 1.e-12
!  bmad_com%min_ds_adaptive_tracking   = 1.e-14
!  bmad_com%fatal_ds_adaptive_tracking = 0.

  ! 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
  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
  call bmad_parser (lat_file, lat)

  ! 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)
  CLOSE(5)

  ! 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,ring_theta)
  endif

  ! Set the Twiss parameters of lat%ele(0)=BEGINNING
  ! to the values in "input.dat"
  lat%ele(0)%a%beta  = twiss%betax 
  lat%ele(0)%b%beta  = twiss%betay
  lat%ele(0)%a%alpha = twiss%alphax
  lat%ele(0)%b%alpha = twiss%alphay
  lat%ele(0)%x%eta   = twiss%etax 
  lat%ele(0)%x%etap  = twiss%etapx 
  lat%ele(0)%y%eta   = twiss%etay 
  lat%ele(0)%y%etap  = twiss%etapy 

  ! Set variables for determining Twiss parameters, coupling, dispersion, etc. of the ring
  do i=1,lat%n_ele_track-1
    if (index(lat%ele(i)%name,'QUAD')/=0) then
      ! Set the quad field indices
       if (index(lat%ele(i)%name,'1')/=0)then
         lat%ele(i)%value(field_index$) = quad_params%field_index(1)
       else     
         lat%ele(i)%value(field_index$) = quad_params%field_index(2) 
       endif  
    elseif(index(lat%ele(i)%name,'KICKER')/= 0) then
      ! Momentarily set the kicker fields to zero (see below)
      lat%ele(i)%value(hkick$) = 0.
    endif
  end do

  ! Calculate the ring's Twiss parameters, coupling, and dispersion parameters
  call lat_make_mat6(lat,-1)
  call twiss_at_start(lat,status)
!  if(status/=0)print '(a,a)',' Status from twiss_at_start = ', status_name(status)
  call twiss_propagate_all(lat)
  call chrom_calc(lat, delta_e,chrom_x,chrom_y, err_flag, low_e_lat, high_e_lat, low_e_orb, high_e_orb)
!  write(91,'(a,4a12)')'#','Low E Qx','Low E Qy','High E Qx','High E Qy'
!  write(91,'(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
!  write(91,'(5es12.4)')(low_e_lat%ele(i)%s,low_e_orb(i)%vec(1),low_e_orb(i)%vec(3),high_e_orb(i)%vec(1),high_e_orb(i)%vec(3), i=1,lat%n_ele_track) 
  if(err_flag)print *,' Error from chrom calc'
  if (twiss_file) then
      call twiss_propagate_cm(lat)
      call write_transfermatrix(lat)
  endif

  ! Set kicker parameters (now that ring's Twiss parameters, etc., have been calculated)
  do i=1,lat%n_ele_track-1
    if( index(lat%ele(i)%name,'KICKER')/= 0 ) then
      if( index(lat%ele(i)%name,'1')/=0 ) then
        ! KICKER1
        lat%ele(i)%value(hkick$) = kicker_params%hkick(1)
        lat%ele(i)%value(kick_width$) = kicker_params%kick_width(1)
      else if( index(lat%ele(i)%name,'2')/=0 ) then
        ! KICKER2
        lat%ele(i)%value(hkick$) = kicker_params%hkick(2)
        lat%ele(i)%value(kick_width$) = kicker_params%kick_width(2)
      else
        ! KICKER3
        lat%ele(i)%value(hkick$) = kicker_params%hkick(3)
        lat%ele(i)%value(kick_width$) = kicker_params%kick_width(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%hkick(1)*1.e4, kicker_params%kick_width(1)*1.e9
  WRITE(*,'(a,6f12.1)') " K2|", kicker_params%hkick(2)*1.e4, kicker_params%kick_width(2)*1.e9
  WRITE(*,'(a,6f12.1)') " K3|", kicker_params%hkick(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)')        "============================================================================"
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,*)
  WRITE(*,'(a)')        "@Lattice"
  WRITE(*,'(a)')        "================================================================================"
  WRITE(*,'(7a16)')     "ELEMENT NAME", "r(m)", "s(m)", "angle(deg)", "total(deg)", "field_index"
  WRITE(*,'(a)')        "--------------------------------------------------------------------------------"
  angle = 0
  DO i=0, lat%n_ele_track
    angle = angle + lat%ele(i)%value(angle$)
    WRITE(*,'(a16,5f16.3)') lat%ele(i)%name, lat%ele(i)%value(rho$), lat%ele(i)%s, lat%ele(i)%value(angle$)*180./pi, angle*180./pi, lat%ele(i)%value(field_index$)
  END DO
  WRITE(*,'(a)')        "--------------------------------------------------------------------------------"
  WRITE(*,'(a,7(a,f8.5,3x))') "RING PARAMS:     ", "Qx =", lat%a%tune/twopi, "Qy =", lat%b%tune/twopi, &
    "betax =", lat%ele(lat%n_ele_track)%a%beta, "betay =", lat%ele(lat%n_ele_track)%b%beta, "field_index =", (lat%b%tune/twopi)**2, &
    " Q'x =",chrom_x," Q'y=",chrom_y
  WRITE(*,'(a)')        "================================================================================"
  WRITE(*,*)
  WRITE(*,*)


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


  call reallocate_coord (co, lat%n_ele_max)
  call init_coord(co(0))
  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, muons)
! NSF: Moved to create_phase_space()
! if(inf_end) call inflector_end_scatter(nmuons, muons)
  runge_kutta_com%check_wall_aperture = quad_plate !if true then scatter in plates
  runge_kutta_com%hit_when = wall_transition$


  allocate(muons_ele(0:nmuons, 0:lat%n_ele_track))
  allocate(muons_temp(0:nmuons))
  
   
  co(0)%vec    = 0
  co(0)%t      = initial_offsets%tmean
  co(0)%s      = 0
  co(0)%vec(1) = initial_offsets%x_mean
  co(0)%vec(3) = initial_offsets%y_mean
  co(0)%vec(5) = initial_offsets%tmean
  co(0)%vec(6) = initial_offsets%pzmean


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

 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)

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,9es12.4)')i,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
     endif

DO i=1, nturns
  IF (i>1) runge_kutta_com%check_wall_aperture = .false.
  DO n=1, nmuons
    IF (muons(n)%coord%state /= alive$) CYCLE
    co(0)%vec = muons(n)%coord%vec
    co(0)%t   = muons(n)%coord%t

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

    DO j=1, lat%n_ele_track
      muons_ele(n,j)%coord%state = alive$
      IF (nmuons==1) write(29,'(i10,1x,9es12.4)')i,co(j)%vec, co(j)%t, &
                                              lat%ele(j)%s/circumference, &
                                    ((i-1)*lat%ele(lat%n_ele_track)%s +lat%ele(j)%s)/circumference         
         
      if (track_state /= moving_forward$ .or. err_flag) then
        muons_ele(n,j)%coord%state = co(track_state)%state
        muons_ele(n,j)%lost = .true.
      endif

      muons_ele(n,j)%coord%vec = co(j)%vec
      muons_ele(n,j)%coord%t   = co(j)%t

      ! Write kicker phase-space info (first turn only)
      IF (i==1) THEN
        IF (index(lat%ele(j)%name,'BFREE')/=0) THEN
          ! START OF KICKERS
          WRITE(111,'(i10,5x,6es15.5,5x,1es15.5)') n, co(j)%vec, co(j)%t
        ELSEIF (index(lat%ele(j)%name,'MARK_CENTERLINE')/=0) THEN
          ! MIDDLE OF KICKERS
          WRITE(112,'(i10,5x,6es15.5,5x,1es15.5)') n, co(j)%vec, co(j)%t
        ELSEIF (index(lat%ele(j)%name,'KICKER3')/=0) THEN
          ! END OF KICKERS
          WRITE(113,'(i10,5x,6es15.5,5x,1es15.5)') n, co(j)%vec, co(j)%t
        ENDIF
      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
 
    muons(n)%coord%vec = co(lat%n_ele_track)%vec
    muons(n)%coord%t   = co(lat%n_ele_track)%t
  ENDDO ! n muons

  do j=1,lat%n_ele_track
     distance = float(i-1)+lat%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)
     if(nturns - i < 100)call moment_range(moment_ele)
  end do

  distance = float(i)
  call compute_moments(muons,distance,moment, nmuons,nmu, 16)
  if(nturns - i < 100)call moment_range(moment)
  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


  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) write(32,'(16a16)')  &
  'K1[Guass]', 'K2[Gauss]', 'K3[Gauss]', &
  'depth<x>',' depth<xp>','depth<y>','depth<yp>','depth<t>','depth<pz>', &
  'depth_SIG(1,1)','depth_SIG(3,3)','depth_SIG(5,5)','depth_SIG(6,6)', &
  'n muons',' after_inf','remaining'
first=.false.

write(32,'(3es16.3,10es16.6,3i16)') &
  kicker_params%hkick(1) * 1.e4, & ! Gauss
  kicker_params%hkick(2) * 1.e4, & ! Gauss
  kicker_params%hkick(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

  deallocate(muons)
  deallocate(muons_ele)
  deallocate(muons_temp)
close(16)
close(26)
!end do
end