! counts quads, dipoles, sextupoles
! writes beta and eta every centimeter for plotting
! writes outline of magnets for plotting - takes out curve
! use beta_eta_magnets.gnu to plot - fort.21 ->cell_arc_quad.dat, fort.23 -> cell_arc_bend.dat
! fort.24 -> beta_eta.dat, fort.25 -> cell_arc_wiggler.dat
! fort.26 -> rf
! no need to copy. 
! Plot with elements using beta_eta_elements.gnu
! Plot with no elements beta_eta_nomags.gnu

 program find_quads

   use bmad

      type (lat_struct),target:: lat
      type (branch_struct),pointer:: branch
      type (ele_struct) ele 
      type (coord_struct), allocatable :: co(:)
      type (coord_struct) orb_at_s 

      real(rp) w,s1,s2, half
      real(rp) offset
      real(rp) dist      

      integer ix, i, j, m, ii, n_turns, index, ixx, k
      integer nargs,iargc
      integer iu,iu1
      integer nbend/0/, nquad/0/, nsext/0/, nwig/0/, nrf/0/
      integer nq60cm/0/, nq30cm/0/, nq02h/0/
      integer nskewquad/0/, nkicker/0/
      integer n
      integer nvkicker/0/
      integer lun, nhkicker/0/

      character*140 lat_file
      character*120 line, last_line, changes(1:3)/3*' '/
      character*3 ans
      character*5 width
      logical transfer_line/.false./

      nargs = cesr_iargc()
      if(nargs == 1)then
         call cesr_getarg(1,lat_file)
         print *, 'Using ', trim(lat_file)
       else
      lat_file = 'bmad.'
      print '(a37,$)',' 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

       print *, ' lat_file = ', lat_file
       call bmad_and_xsif_parser (lat_file, lat)
       print *, ' lat_file = ', lat_file
          w=3.0
       print *, 'Width of elements ',w
       read(5,*)width
       if(len(width) >0 )then
         read(width,*)w
         print *,' width = ', w
        endif  
          offset = -5
       print *,' element offset = ', offset
       read(5,*)offset
       print *,' element offset = ', offset
       call reallocate_coord (co, lat%n_ele_track)
       call closed_orbit_calc(lat, co, i_dim=4)
       call twiss_at_start(lat)
     print '(a,5f12.4)','Closed ring:', lat%ele(0)%a%beta, lat%ele(0)%b%beta, &
                         & lat%ele(0)%a%alpha,lat%ele(0)%b%alpha,lat%ele(0)%a%eta
     print '(a,$)',' Initialize ? '
     read(5,*)ans
     if(index(ans,'y') /= 0 .or. index(ans,'Y')/= 0)then
      transfer_line=.true.
      print '(a)','Type betax, betay, alphax, alphay, etax'
      read(5,*)lat%ele(0)%a%beta, lat%ele(0)%b%beta, & 
                         & lat%ele(0)%a%alpha,lat%ele(0)%b%alpha,lat%ele(0)%a%eta
     endif
       call twiss_propagate_all(lat)
       dist = 0
       write(24, '(a16,4a12,2a14)')'element','s','beta x','beta y','eta','x','y'
       do
         dist = dist + 0.01
         if(dist >lat%ele(lat%n_ele_track)%s)exit
         call twiss_and_track_at_s(lat,dist,ele, orb=co,orb_at_s=orb_at_s) !, ix_branch=nbranch, use_last=.false.)
         write(24,'(a16,4f12.4,2e14.6)')ele%name,dist,ele%a%beta, ele%b%beta, ele%x%eta, orb_at_s%vec(1), orb_at_s%vec(3)
       end do

      call dbeta_ddelta(lat,transfer_line)

       print '(a,f12.4)', ' The length is ',lat%ele(lat%n_ele_track)%s

       open(unit=11, file = 'quad_strengh_length.dat')
       write(11,'(a16,2a12)')'Element','k1[/m^2]',&
           'length[m]'
       open(unit=33, file = 'bend_properties.dat')
       write(33,'(a16,3a12)')'Element','rho[m]','k1[/m^2]', &
           'length[m]'


       open(unit=21, file = 'cell_arc_quad.dat')
       open(unit=22, file = 'cell_arc_sext.dat')
       open(unit=23, file = 'cell_arc_bend.dat')
       open(unit=25, file = 'cell_arc_wiggler.dat')
       open(unit=29, file = 'hkicker.dat')
       open(unit=48, file = 'vkicker.dat')
       open(unit=31, file = 'branch_quad.dat')
       open(unit=32, file = 'branch_sext.dat')
       open(unit=33, file = 'branch_bend.dat')
       open(unit=39, file = 'branch_hkicker.dat')
       open(unit=58, file = 'branch_vkicker.dat')

       do n= 0,size(lat%branch)-1
       branch => lat%branch(n)
       do i=1,branch%n_ele_track
        half = branch%param%total_length/2
        s1= branch%ele(i-1)%s
         if(s1 < -half) s1=s1+2*half
        s2= branch%ele(i)%s
         if(s2 < -half) s2 = s2 +2*half

        if(branch%ele(i)%key == quadrupole$ .and. branch%ele(i)%value(tilt$)==0)then
         nquad=nquad+1
!         print '(2a,f10.2)', branch%ele(i)%name, '   quad length ', branch%ele(i)%value(l$)
         if(branch%ele(i)%value(l$) > 0.55 .and. branch%ele(i)%value(l$) < 0.65)nq60cm=nq60cm+1
         if(branch%ele(i)%value(l$) > 0.25 .and. branch%ele(i)%value(l$) < 0.35)nq30cm=nq30cm+1
         if(index(branch%ele(i)%name,'Q02_H')/= 0)nq02h=nq02h+1 
         write(11,'(a16,2f12.4)')trim(branch%ele(i)%name),branch%ele(i)%value(k1$),&
           branch%ele(i)%value(l$)
          lun=21+n*10
          write(lun,'(/,a,/)')'#'
          write(lun,'(2f10.3)')s1, -w+offset
          write(lun,'(2f10.3)')s1,  w+offset
          write(lun,'(2f10.3)')s2, w+offset
          write(lun,'(2f10.3)')s2, -w+offset
          write(lun,'(2f10.3)')s1, -w+offset
        endif

        if(branch%ele(i)%key == quadrupole$ .and. branch%ele(i)%value(tilt$)/=0)then
         nskewquad=nskewquad+1
!         print '(2a,f10.2)', branch%ele(i)%name, '   quad length ', branch%ele(i)%value(l$)
         write(11,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(k1$),&
           branch%ele(i)%value(l$)
          write(28,'(/,a,/)')'#'
          write(28,'(2f10.3)')s1, -w+offset
          write(28,'(2f10.3)')s1,  w+offset
          write(28,'(2f10.3)')s2, w+offset
          write(28,'(2f10.3)')s2, -w+offset
          write(28,'(2f10.3)')s1, -w+offset
        endif
        if(branch%ele(i)%key == sextupole$)then
         nsext=nsext+1
         write(12,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(k2$),&
           branch%ele(i)%value(l$)
          lun = 22+n*10
          write(lun,'(/,a,/)')'#'
          write(lun,'(2f10.3)')s1, -w+offset
          write(lun,'(2f10.3)')s1,  w+offset
          write(lun,'(2f10.3)')s2, w+offset
          write(lun,'(2f10.3)')s2, -w+offset
          write(lun,'(2f10.3)')s1, -w+offset
        endif
        if(branch%ele(i)%key == sbend$)then
         nbend=nbend+1
         write(33,'(a16,3f12.4)')trim(branch%ele(i)%name),branch%ele(i)%value(rho$),branch%ele(i)%value(k1$), &
           branch%ele(i)%value(l$)
          lun = 23+n*10
          write(lun,'(/,a,/)')'#'
          write(lun,'(2f10.3)')s1, -w+offset
          write(lun,'(2f10.3)')s1,  w+offset
          write(lun,'(2f10.3)')s2, w+offset
          write(lun,'(2f10.3)')s2, -w+offset
          write(lun,'(2f10.3)')s1, -w+offset
        endif

        if(branch%ele(i)%key == wiggler$)then
         nwig=nwig+1
         write(13,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(rho$),&
           branch%ele(i)%value(l$)
          write(25,'(/,a,/)')'#'
          write(25,'(2f10.3)')s1, -w+offset
          write(25,'(2f10.3)')s1,  w+offset
          write(25,'(2f10.3)')s2, w+offset
          write(25,'(2f10.3)')s2, -w+offset
          write(25,'(2f10.3)')s1, -w+offset
        endif

     if(branch%ele(i)%key == rfcavity$)then
         nrf=nrf+1
         write(13,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(rho$),&
           branch%ele(i)%value(l$)
          write(26,'(/,a,/)')'#'
          write(26,'(2f10.3)')s1, -w+offset
          write(26,'(2f10.3)')s1,  w+offset
          write(26,'(2f10.3)')s2, w+offset
          write(26,'(2f10.3)')s2, -w+offset
          write(26,'(2f10.3)')s1, -w+offset
        endif

     if(branch%ele(i)%key == kicker$)then
         nkicker = nkicker+1
         write(13,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(rho$),&
           branch%ele(i)%value(l$)
          write(27,'(/,a,/)')'#'
          write(27,'(2f10.3)')s1, -w+offset
          write(27,'(2f10.3)')s1,  w+offset
          write(27,'(2f10.3)')s2, w+offset
          write(27,'(2f10.3)')s2, -w+offset
          write(27,'(2f10.3)')s1, -w+offset
        endif

     if(branch%ele(i)%key == hkicker$)then
         nhkicker = nhkicker+1
         write(13,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(kick$),&
           branch%ele(i)%value(l$)
          lun = 29 + 10*n
          write(lun,'(/,a,/)')'#'
          write(lun,'(2f10.3)')s1, -w+offset
          write(lun,'(2f10.3)')s1,  w+offset
          write(lun,'(2f10.3)')s2, w+offset
          write(lun,'(2f10.3)')s2, -w+offset
          write(lun,'(2f10.3)')s1, -w+offset
        endif

     if(branch%ele(i)%key == vkicker$)then
         nvkicker = nvkicker+1
         write(13,'(a,2f12.4)')branch%ele(i)%name,branch%ele(i)%value(kick$),&
           branch%ele(i)%value(l$)
          lun = 48 + 10*n
          write(lun,'(/,a,/)')'#'
          write(lun,'(2f10.3)')s1, -w+offset
          write(lun,'(2f10.3)')s1,  w+offset
          write(lun,'(2f10.3)')s2, w+offset
          write(lun,'(2f10.3)')s2, -w+offset
          write(lun,'(2f10.3)')s1, -w+offset
        endif

       end do

      print *
      print '(a,1x,i10)',' Branch = ', n
      print '(i4,a)',nbend,' bends'
      print '(i4,a,i4,a,i4,a)',nquad,' quads   ',nq60cm,' 60cm quads  ', nq30cm,' 30cm quads '
      print '(i4,a)',nq02h ,' split arc quads'
      print '(i4,a)',nsext,' sextupoles'
      print '(i4,a)',nwig,' wigglers'
      print '(i4,a)',nrf,' rfcavities'
      print '(i4,a)',nskewquad,' nskewquad'
      print '(i4,a)',nkicker,' nkicker'
      print '(i4,a)',nhkicker,' nhkicker'
      print '(i4,a)',nvkicker,' nvkicker'
      
      close(unit=21+n*10)
      close(unit=22+n*10)
      close(unit=23+n*10)
      close(unit=29+n*10)
      close(unit=48+n*10)
  enddo !branch
end

subroutine dbeta_ddelta(lat, transfer_line)
 
use bmad
implicit none
type(lat_struct) lat, ring_two(-1:1)
type(coord_struct), allocatable :: co_off(:)
type(ele_struct) ele(-1:1)

real(rp) de/1.e-4/,dist/0.01/
logical transfer_line
integer i_dim/4/
integer i,j, lun

      lun= lunget()
      open(unit=lun,file='dbeta_denergy.dat')
      write(lun,'(a16,7a12)')'Element','s','betax high','betax low','betay high','betay low','dbetax/de','dbetay/de'
      call reallocate_coord(co_off,lat%n_ele_track)
      ring_two(1) = lat
      ring_two(-1) = lat
      do i = -1,1,2
       co_off(0)%vec(6) = de *i
       if(transfer_line)then
         co_off(0)%vec(1) = ring_two(i)%ele(0)%a%eta * co_off(0)%vec(6)
         co_off(0)%vec(2) = ring_two(i)%ele(0)%a%etap * co_off(0)%vec(6)
         co_off(0)%vec(3) = ring_two(i)%ele(0)%b%eta * co_off(0)%vec(6)
         co_off(0)%vec(4) = ring_two(i)%ele(0)%b%etap * co_off(0)%vec(6)

         call track_all(ring_two(i), co_off)
       endif
       if(.not. transfer_line) call closed_orbit_calc(ring_two(i), co_off,i_dim)
       call lat_make_mat6(ring_two(i), -1, co_off)
       if(.not. transfer_line) call twiss_at_start(ring_two(i))
       call twiss_propagate_all(ring_two(i))
      end do

       dist = 0
       do
         dist = dist + 0.01
         if(dist >lat%ele(lat%n_ele_track)%s)exit
         call twiss_and_track_at_s(ring_two(1),dist,ele(1))
         call twiss_and_track_at_s(ring_two(-1),dist,ele(-1))
         write(lun,'(a16,7f12.4)')ele(1)%name,dist,ele(1)%a%beta,ele(-1)%a%beta,&
                         ele(1)%b%beta,ele(-1)%b%beta,(ele(1)%a%beta - ele(-1)%a%beta)/2/de, &
                           (ele(1)%b%beta-ele(-1)%b%beta)/2/de
       end do
       close(unit=lun)
return
end