 program resonantextraction

   use bmad
   use beam_mod
   use multibunch_mod
   use multibunch_interface
   use bmadz_other_code_interface   

     implicit none

      type (lat_struct), target :: lat
      type (coord_struct), allocatable::co(:), co_fft(:), coa(:)
      type (beam_init_struct) beam_init
      type (beam_struct) beam
      type (moment_struct), allocatable :: bunch_moments(:)
      integer ix, nargs
      integer i, j, k
      integer lun
      integer lun1, lun2, lun3
      integer nturns
      integer ios
      integer n_fft/256/
      integer nparticles

      real(rp) delta
      real(rp) e_init, t0, omega, rf_phase_twopi
      real(rp) initial_offset(6)
      real(rp) Q(3), phase(3)

      character*140 lat_file
      character*120 line, lineout
      character*10 iteration

      logical err
      logical ramp

      namelist/input/e_init,ramp, t0, nturns, rf_phase_twopi, initial_offset, n_fft, nparticles

      bmad_com%auto_bookkeeper = .false.

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


  !ready an input file from the command line
      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= synch.lat) '
        read(5,'(a)') line
        call string_trim(line, lineout, ix)
        lat_file = lineout
        if(ix == 0) lat_file = '/home/jdp279/Documents/Track6610/Outgoing/darkphoton/files/synch.lat'
        print *, ' lat_file = ', lat_file
      endif
      call bmad_parser(lat_file, lat)
      call reallocate_coord(co,lat%n_ele_track)
      call reallocate_coord(co_fft, nturns)
      call reallocate_coord(coa,n_fft)
      call lattice_bookkeeper(lat)
      lat%param%ixx = 0
      call closed_orbit_calc(lat,co,6)
      call track_all(lat,co)
      call lat_make_mat6(lat,-1)        ! compute transfer matrices for all of the elements in the lattice

      call twiss_at_start(lat)          ! compute twiss parameters at the starting point
      call twiss_propagate_all(lat)     ! propagate twiss parameters to all elements in the lattice
      call writeinfo(lat)               ! write initial twiss parameters to file

      lun1 = lunget()
      lun2 = lunget()
      lun3 = lunget()
      open(unit=lun1, file = 'trajectory_ele_by_ele.dat')
      open(unit=lun2, file = 'trajectory.dat')
      open(unit=lun3, file = 'beam_trajectory.dat')
      

      do i=1,lat%n_ele_max
         if(lat%ele(i)%key == rbend$ .or. lat%ele(i)%key == quadrupole$ .or.lat%ele(i)%key == sbend$)then !  save reference values for k1, k2
           lat%ele(i)%value(custom_attribute1$) = lat%ele(i)%value(k1$)
           lat%ele(i)%value(custom_attribute2$) = lat%ele(i)%value(k2$)
         endif
         if(lat%ele(i)%key == multipole$)then    !save reference values of multipoles
           allocate(lat%ele(i)%r(1:2,0:n_pole_maxx,0:1))
           lat%ele(i)%r(1,0:n_pole_maxx,1) = lat%ele(i)%a_pole(:)
           lat%ele(i)%r(2,0:n_pole_maxx,1) = lat%ele(i)%b_pole(:)
         endif
         if(lat%ele(i)%key == rfcavity$) lat%ele(i)%value(phi0$) = rf_phase_twopi/twopi  ! set RF phase
      end do

      omega = twopi * 60
      if(ramp)then
        lat%param%ixx = 1  !if 0 no ramp.
        delta= e_init/lat%ele(0)%value(E_TOT$)
        t0 = acos(1.-2*delta)/omega
        co(0)%t = t0
        co(0)%vec(6) = -(1-delta)
      endif
!
   if(nparticles == 1)then
      co(0)%vec = co(0)%vec + initial_offset 
      print '(a,6es12.4)','co(0)%vec',co(0)%vec
      do i=1,nturns
       call track_all(lat,co)
!       do j=1,lat%n_ele_track
!        write(lun1, '(i10,1x,es12.4,1x,6es12.4,1x,es12.4)')i,lat%ele(j)%s, co(j)%vec,co(j)%t
!       end do
        co(0) = co(lat%n_ele_track)
        co_fft(i) = co(0)
        write(lun2, '(i10,1x,es12.4,1x,6es12.4,1x,es12.4)')i,i*lat%ele(lat%n_ele_track)%s, co(0)%vec,co(0)%t
      end do
   endif   ! single particle

   if(nparticles > 1)then
       beam_init%n_bunch = 1
       beam_init%dt_bunch = 0
       beam_init%n_particle=nparticles
       beam_init%bunch_charge = 1

       beam_init%a_emit = 1.e-10
       beam_init%b_emit = 1.e-10
       beam_init%sig_z = 0.001
       beam_init%sig_e = 0.001
       allocate(bunch_moments(0:beam_init%n_bunch))    
       call init_beam_distribution(lat%ele(0),lat%param, beam_init, beam)
       forall(i=1:6)beam%bunch(1)%particle(:)%vec(i) = co(0)%vec(i)+beam%bunch(1)%particle(:)%vec(i) + initial_offset(i)
       beam%bunch(1)%particle(:)%t = co(0)%t
!           do j=1,5
!            print '(i10,7es12.4)',i,beam%bunch(1)%particle(j)%vec, beam%bunch(1)%particle(j)%t
!           end do

        write(lun3,'(a10,2a72)')'turn','bunch_moments%ave(1:6)', 'bunch_moments%sigma(1:6,1:6)'
        do i=1,nturns
          do j = 0,lat%n_ele_track-1
            call track_beam(lat,beam,ele1=lat%ele(j),ele2=lat%ele(j+1), err=err)  
!           do k=1,5
!            print '(i10,7es12.4)',j,beam%bunch(1)%particle(k)%vec, beam%bunch(1)%particle(k)%t
!           end do
!           if(j==10)pause
          end do
!           do j=1,5
!            print '(i10,7es12.4)',i,beam%bunch(1)%particle(j)%vec, beam%bunch(1)%particle(j)%t
!           end do
!           pause
           call beam_moments(beam, bunch_moments)
           write(lun3,'(i10,12es12.4)')i,bunch_moments(1)%ave(1:6), bunch_moments(1)%sigma(1,1), bunch_moments(1)%sigma(2,2), &
                 bunch_moments(1)%sigma(3,3),bunch_moments(1)%sigma(4,4),bunch_moments(1)%sigma(5,5),bunch_moments(1)%sigma(6,6)
           co_fft(i)%vec = bunch_moments(1)%ave
        end do
      endif ! multiple particles

! compute fft for n_fft turns, at n_fft/2 intervals
      ix=1
      write(13,'(a10,a10,a36,a36)')'start turn','end turn','Q','Phase'
      do while(ix + n_fft <= nturns)
      coa(1:n_fft) = co_fft(ix:ix+n_fft-1)
      call fft_plane(n_fft,coa,Q,phase)
      write(13,'(i10,i10, 3es12.4,3es12.4)')ix,ix+n_fft-1,  Q,  phase
      print '(i5,a1,i5, a,3es12.4,a,3es12.4)',ix,'-',ix+n_fft-1, ' Q = ', Q, ' phase = ', phase
      ix = ix+n_fft/2
      end do

   close(unit=lun1)
   close(unit=lun2)
   close(unit=lun3)


 end


