program resonantextraction

   use bmad
    use beam_mod
     implicit none

      type (lat_struct), target :: lat
      type (ele_struct), pointer :: rex, esrefpt, magrefpt, test, mq1, mq2, mq3, mq4
      type (ele_struct), pointer :: einj, pinj, eext1, eext2, pext1, pext2, bump1, bump2, bump3
      type (ele_pointer_struct), allocatable :: eles(:), elesa(:), elesb(:), elesc(:), elesd(:), elese(:)
      type (ele_pointer_struct), allocatable :: elesf(:), elesg(:), elesh(:), elesi(:), elesj(:), elesk(:)
      type (ele_pointer_struct), allocatable :: elesl(:), elesm(:), elesn(:)
      type (coord_struct), allocatable :: orbit(:), trajectory(:,:)
      type (coord_struct) :: emptystruct

      integer ix, nargs
      integer i, j
      integer exitturn
      integer lun
      integer nturns, startpoint, printnumber
      integer n_loc
      integer bumpinitturn, bumpramplength, rexinitturn, rexramplength

      real rexinit, rexmax, kqhinit, kqh1, kqh2, kqh3, kqh4
      real initialwait, ramp1, ramp2, ramp3, ramp4, ramp5, finalwait
      real esseptinner, esseptouter, magseptinner, magseptouter, esseptkick, esouterwall, magouterwall
      real kick1, kick2, kick3, initkick1, initkick2, initkick3
      real turnreal
      real xinit, pxinit, yinit, pyinit, zinit, einit

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

      logical err

      bmad_com%auto_bookkeeper = .false.

  !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/Singletrack/darkphoton/files/synch.lat'
        print *, ' lat_file = ', lat_file
      endif

  !run parameters
      nturns = 4080       
      startpoint = 0   ! For now, depending on what parts are skipped, need to manually change the seven !** values below.
      printnumber = 30

      initialwait = 2525
      ramp1 = 250
      ramp2 = 100
      ramp3 = 750
      ramp4 = 150
      ramp5 = 300
      finalwait = 5      ! summing these should equal nturns

      kqhinit = 0.40096   ! initial main quadrupole strength
      kqh1 = 0.4470       ! after ramp 1
      kqh2 = 0.4495       ! after ramp 2
      kqh3 = 0.4575       ! after ramp 3
      kqh4 = 0.4575       ! after ramp 4

      rexinit = 0.0       !** 0
      rexinitturn = 500
      rexramplength = 1000
      rexmax = 5.0          

  !septum parameters
      esseptinner=0.0310
      esseptouter=0.0311
      esouterwall=0.040
      esseptkick=0.00017

      magseptinner=0.0247
      magseptouter=0.0257
      magouterwall=0.040

  !bump parameters

      initkick1=0.0      !** 0
      initkick2=0.0      !** 0
      initkick3=0.0      !** 0
      bumpinitturn=500
      bumpramplength=1000
      kick1= 0.00121
      kick2= 0.00057
      kick3= -0.00178

  !initial particle parameters
      xinit=0.013      !** 0.0125
      pxinit=0
      yinit=0.0025      !** 0.0025
      pyinit=0
      zinit=0
      einit=0.0020       !** 0.002

  !set up lattice parameters

      call bmad_parser (lat_file, lat)  ! parse lattice file

    !shortcut references for lattice elements.
      !extraction elements
      call lat_ele_locator('REX',lat,elesa,n_loc,err)
      rex => elesa(1)%ele

      call lat_ele_locator('ESSEPT',lat,elesb,n_loc,err)
      esrefpt => elesb(1)%ele

      call lat_ele_locator('MAGSEPT',lat,elesc,n_loc,err)
      magrefpt => elesc(1)%ele

      call lat_ele_locator('Q01V',lat,elesd,n_loc,err)
      mq1 => elesd(1)%ele
      mq2 => elesd(2)%ele

      call lat_ele_locator('Q02H',lat,elese,n_loc,err)
      mq3 => elese(1)%ele
      mq4 => elese(2)%ele

      !preexisting septa
      call lat_ele_locator('ELINJSEP',lat,elesf,n_loc,err)
      einj => elesf(1)%ele

      call lat_ele_locator('POSINJSEP',lat,elesg,n_loc,err)
      pinj => elesg(1)%ele

      call lat_ele_locator('ELEXT1SEP',lat,elesh,n_loc,err)
      eext1 => elesh(1)%ele

      call lat_ele_locator('ELEXT2SEP',lat,elesi,n_loc,err)
      eext2 => elesi(1)%ele

      call lat_ele_locator('POSEXT1SEP',lat,elesj,n_loc,err)
      pext1 => elesj(1)%ele

      call lat_ele_locator('POSEXT2SEP',lat,elesk,n_loc,err)
      pext2 => elesk(1)%ele

      !bump kickers
      call lat_ele_locator('BUMP1',lat,elesl,n_loc,err)
      bump1 => elesl(1)%ele

      call lat_ele_locator('BUMP2',lat,elesm,n_loc,err)
      bump2 => elesm(1)%ele

      call lat_ele_locator('BUMP3',lat,elesn,n_loc,err)
      bump3 => elesn(1)%ele

      !set initial element strengths
      mq1%value(k1$)=-kqhinit
      mq2%value(k1$)=-kqhinit
      mq3%value(k1$)=kqhinit
      mq4%value(k1$)=kqhinit
      call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
      call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
      call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
      call set_flags_for_changed_attribute(mq4,mq4%value(k1$))

      rex%a_pole(2)=rexinit
      call set_flags_for_changed_attribute(rex,rex%a_pole(2))

      bump1%value(kick$)=initkick1
      bump2%value(kick$)=initkick2
      bump3%value(kick$)=initkick3
      call set_flags_for_changed_attribute(bump1,bump1%value(kick$))
      call set_flags_for_changed_attribute(bump2,bump2%value(kick$))
      call set_flags_for_changed_attribute(bump3,bump3%value(kick$))

      call lattice_bookkeeper(lat)
      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

      lun = lunget()
      open(unit=lun, file='placements.dat')
      write(lun,'(a15,i4)')'rex = ', rex%ix_ele
      write(lun,'(a15,i4)')'esrefpt = ', esrefpt%ix_ele
      write(lun,'(a15,i4)')'magrefpt = ', magrefpt%ix_ele
      write(lun,'(a15,i4)')'einj = ', einj%ix_ele
      write(lun,'(a15,i4)')'pinj = ', pinj%ix_ele
      write(lun,'(a15,i4)')'eext1 = ', eext1%ix_ele
      write(lun,'(a15,i4)')'eext2 = ', eext2%ix_ele
      write(lun,'(a15,i4)')'pext1 = ', pext1%ix_ele
      write(lun,'(a15,i4)')'pext2 = ', pext2%ix_ele
      write(lun,'(a15,i4)')'bump1 = ', bump1%ix_ele
      write(lun,'(a15,i4)')'bump2 = ', bump2%ix_ele
      write(lun,'(a15,i4)')'bump3 = ', bump3%ix_ele
      close(unit=lun)

      allocate(trajectory(nturns,0:lat%n_ele_track))
      allocate(orbit(0:lat%n_ele_track))

  !track particle through lattice. (The proper functioning of this section depends on the order of elements in the lattice file!)

      exitturn = -1

      orbit(0)%vec(1)=xinit
      orbit(0)%vec(2)=pxinit
      orbit(0)%vec(3)=yinit
      orbit(0)%vec(4)=pyinit
      orbit(0)%vec(5)=zinit
      orbit(0)%vec(6)=einit

        tloop: do i=startpoint, nturns

          print *, 'turn', i, 'quad = ', mq4%value(k1$), 'bump = ', bump1%value(kick$), 'rex = ', rex%a_pole(2)

          call track_many(lat,orbit,0,esrefpt%ix_ele,1)
          !aperture violation check
              if(orbit(esrefpt%ix_ele)%state == alive$)then
               else
                print *, 'aperture violation before essept'
                exitturn = i
                exit tloop
              end if

          !electrostatic septum checks and kick
              if(orbit(esrefpt%ix_ele)%vec(1) > esseptinner)then
                if(orbit(esrefpt%ix_ele)%vec(1) > esseptouter)then
                  if(orbit(esrefpt%ix_ele)%vec(1) > esouterwall)then
                    print *, 'particle hit outer wall at essept!'
                    exitturn = i
                    exit tloop
                   else                
                    orbit(esrefpt%ix_ele)%vec(2) = orbit(esrefpt%ix_ele)%vec(2)+esseptkick
                    print *, 'particle kicked'
                  end if
                 else
                  print *, 'particle hit electrostatic septum!'
                  exitturn = i
                  exit tloop
                end if
              end if

          call track_many(lat,orbit,esrefpt%ix_ele,magrefpt%ix_ele,1)
          !aperture violation check
              if(orbit(magrefpt%ix_ele)%state == alive$)then
               else
                print *, 'aperture violation between essept and magsept'
                exitturn = i
                exit tloop
              end if

          !magnetic septum checks and extraction
              if(orbit(magrefpt%ix_ele)%vec(1) > magseptinner)then
                if(orbit(magrefpt%ix_ele)%vec(1) > magseptouter)then
                  if(orbit(magrefpt%ix_ele)%vec(1) > magouterwall)then
                    print *, 'particle hit outer wall at magsept!'
                    exitturn = i
                    exit tloop
                   else                
                    print *, 'particle successfully extracted!'
                    exitturn = i
                    exit tloop
                  end if
                 else
                  print *, 'particle hit magnetic septum!'
                  exitturn = i
                  exit tloop
                end if
              end if

          call track_many(lat,orbit,magrefpt%ix_ele,lat%n_ele_track,1)
          !aperture violation check
              if(orbit(lat%n_ele_track)%state == alive$)then
               else
                print *, 'aperture violation after magsept'
                exitturn = i
                exit tloop
              end if

          !save and reset orbit

          do j=0, lat%n_ele_track
            trajectory(i,j)=orbit(j)
          end do

          orbit(0)=orbit(lat%n_ele_track)
          do j=1, lat%n_ele_track
            orbit(j)=emptystruct
          end do

          !set element strengths for next turn

          if (i > (bumpinitturn-1) .AND. i < (bumpinitturn+bumpramplength)) then
            bump1%value(kick$)=bump1%value(kick$)+ kick1/bumpramplength
            bump2%value(kick$)=bump2%value(kick$)+ kick2/bumpramplength
            bump3%value(kick$)=bump3%value(kick$)+ kick3/bumpramplength
            call set_flags_for_changed_attribute(bump1,bump1%value(kick$))
            call set_flags_for_changed_attribute(bump2,bump2%value(kick$))
            call set_flags_for_changed_attribute(bump3,bump3%value(kick$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i > (rexinitturn-1) .AND. i < (rexinitturn+rexramplength)) then
            rex%a_pole(2)=rex%a_pole(2)+ rexmax/rexramplength
            call set_flags_for_changed_attribute(rex,rex%a_pole(2))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i > initialwait-1 .AND. i < (initialwait+ramp1)) then
            mq1%value(k1$)=mq1%value(k1$) -((kqh1-kqhinit)/ramp1)
            mq2%value(k1$)=mq2%value(k1$) -((kqh1-kqhinit)/ramp1)
            mq3%value(k1$)=mq3%value(k1$) +((kqh1-kqhinit)/ramp1)
            mq4%value(k1$)=mq4%value(k1$) +((kqh1-kqhinit)/ramp1)
            call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
            call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
            call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
            call set_flags_for_changed_attribute(mq4,mq4%value(k1$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i > (initialwait+ramp1-1) .AND. i < (initialwait+ramp1+ramp2)) then
            mq1%value(k1$)=mq1%value(k1$) -((kqh2-kqh1)/ramp2)
            mq2%value(k1$)=mq2%value(k1$) -((kqh2-kqh1)/ramp2)
            mq3%value(k1$)=mq3%value(k1$) +((kqh2-kqh1)/ramp2)
            mq4%value(k1$)=mq4%value(k1$) +((kqh2-kqh1)/ramp2)
            call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
            call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
            call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
            call set_flags_for_changed_attribute(mq4,mq4%value(k1$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i > (initialwait+ramp1+ramp2-1) .AND. i < (initialwait+ramp1+ramp2+ramp3)) then
            mq1%value(k1$)=mq1%value(k1$) -((kqh3-kqh2)/ramp3)
            mq2%value(k1$)=mq2%value(k1$) -((kqh3-kqh2)/ramp3)
            mq3%value(k1$)=mq3%value(k1$) +((kqh3-kqh2)/ramp3)
            mq4%value(k1$)=mq4%value(k1$) +((kqh3-kqh2)/ramp3)
            call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
            call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
            call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
            call set_flags_for_changed_attribute(mq4,mq4%value(k1$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i > (initialwait+ramp1+ramp2+ramp3-1) .AND. i < (initialwait+ramp1+ramp2+ramp3+ramp4)) then
            mq1%value(k1$)=mq1%value(k1$) -((kqh4-kqh3)/ramp4)
            mq2%value(k1$)=mq2%value(k1$) -((kqh4-kqh3)/ramp4)
            mq3%value(k1$)=mq3%value(k1$) +((kqh4-kqh3)/ramp4)
            mq4%value(k1$)=mq4%value(k1$) +((kqh4-kqh3)/ramp4)
            call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
            call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
            call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
            call set_flags_for_changed_attribute(mq4,mq4%value(k1$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          if (i> (initialwait+ramp1+ramp2+ramp3+ramp4-1) .AND. i < (initialwait+ramp1+ramp2+ramp3+ramp4+ramp5)) then
            mq1%value(k1$)=mq1%value(k1$) -((kqhinit-kqh4)/ramp5)
            mq2%value(k1$)=mq2%value(k1$) -((kqhinit-kqh4)/ramp5)
            mq3%value(k1$)=mq3%value(k1$) +((kqhinit-kqh4)/ramp5)
            mq4%value(k1$)=mq4%value(k1$) +((kqhinit-kqh4)/ramp5)
            call set_flags_for_changed_attribute(mq1,mq1%value(k1$))
            call set_flags_for_changed_attribute(mq2,mq2%value(k1$))
            call set_flags_for_changed_attribute(mq3,mq3%value(k1$))
            call set_flags_for_changed_attribute(mq4,mq4%value(k1$))
            call lat_make_mat6(lat,-1)
            call lattice_bookkeeper(lat)
          end if

          !adiabatic damping from energy change
          turnreal = REAL(i)
          orbit(0)%vec(1) = orbit(0)%vec(1)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))
          orbit(0)%vec(2) = orbit(0)%vec(2)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))
          orbit(0)%vec(3) = orbit(0)%vec(3)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))
          orbit(0)%vec(4) = orbit(0)%vec(4)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))
          orbit(0)%vec(5) = orbit(0)%vec(5)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))
          orbit(0)%vec(6) = orbit(0)%vec(6)*(1+((sin(turnreal/1052.0))/(2104.0*(cos(turnreal/1052.0)-1.0)-103.5)))

        end do tloop

     !save exit turn orbit
     if (exitturn > -1) then
       do j=0, lat%n_ele_track
         trajectory(exitturn,j)=orbit(j)
       end do
      else
       exitturn=nturns
     end if

     print *, 'exitturn = ', exitturn

     if (exitturn<printnumber) then
       printnumber = exitturn
     end if

     lun = lunget()
     open(unit=lun, file='trajectory.dat')
     do i=(exitturn-printnumber),exitturn
       write(lun,'(a10)'),''
       write(lun,'(a10,i5)') 'turn', i
       write(lun,'(a10,6a12)')'ele number','x','px','y','py','z','delta'
         j=897
         write(lun,'(i10,6es12.4)')j, trajectory(i,j)%vec(1:6)
     end do
     close(unit=lun)


 end