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

   use bmad
   use muon_mod
   use muon_interface
   use parameters_bmad
   use quad_scrape_parameters   
      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
      real(rp) f_rev
      real(rp) delta_e, chrom_x, chrom_y
      real(rp) delta_B/0./

      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/
      integer ix_branch, status, direction/1/
      integer datetime_values(8)
      integer system
      integer ios
      integer ie
      
            
      character*140 lat_file
      character*120 line, last_line, changes(1:3)/3*' '/, filename
      character*3 ans
      character*5 width
      character*1 num_ix_branch
      character*16 date, time, zone
      logical transfer_line/.false./
      logical plot_details/.false./

      namelist/quad_input/quad_params
      
    call date_and_time(date,time,zone,datetime_values)
    print '(a10,a10,a1,a10)',' date = ',date,' time = ', time
    directory = trim(date)//'_'//trim(time(1:6))
    status= system('mkdir ' // directory)
    if(status == 0)print '(a20,a)',' create directory = ', trim(directory)

      OPEN (UNIT=5, FILE='quad_input.dat', STATUS='old', IOSTAT=ios)
    READ (5, NML=quad_input, IOSTAT=ios)

  print *, 'ios=', ios
   rewind(unit=5)

  CLOSE(5)
  write(6,nml = quad_input)

!      quad_fringe_energy_change = .false.
!      quad_fringe_energy_change = .true.
      bmad_com%spin_tracking_on = .true.
      nargs = command_argument_count()
      filename = 'test'
      ix_branch = 0
      
      if(nargs == 1)then
         call get_command_argument(1,lat_file)
         print *, 'Using ', trim(lat_file)
      elseif(nargs == 2)then
         call get_command_argument(1,lat_file)
         print *, 'Using ',trim(lat_file)
         call get_command_argument(2,filename)
         print *,'Using ', trim(filename)
      elseif(nargs == 3)then
         call get_command_argument(1,lat_file)
         print *, 'Using ',trim(lat_file)
         call get_command_argument(2,filename)
         print *,'Using ', trim(filename)
         call get_command_argument(3, num_ix_branch)
         read(num_ix_branch,'(i)')ix_branch
         print *,'Using branch ', ix_branch
      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 '(2a)', ' lat_file = ', lat_file
       print '(a,$)',' Output file name = ?'
       read(5,'(a)')filename
       print '(a,$)',' Branch = ?'
       read(5,'(i)')ix_branch
    endif

       print '(2a)', ' lat_file = ', lat_file
       call bmad_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 set_dipole_params(lat, ix_branch, Delta_B)
       call set_quad_params(lat,ix_branch)
       call reallocate_coord (co, lat%branch(ix_branch)%n_ele_track)
       do ie=1,lat%branch(ix_branch)%n_ele_track
          print *,ie,lat%branch(ix_branch)%ele(ie)%name
       end do
       
       print '(a,$)',' Fractional momentum offset ? '
       read(5,*)co(0)%vec(6)
       print '(a,es12.4)',' Fractional momentum offset = ', co(0)%vec(6)
       scraping_on = .false.       
       call closed_orbit_calc(lat, co, i_dim=4,direction=direction, ix_branch=ix_branch)
       print '(a,6es12.4)','closed orbit : ',co(lat%branch(ix_branch)%n_ele_track)%vec
       do ie = 1,lat%branch(ix_branch)%n_ele_track
          print *,ie,lat%branch(ix_branch)%ele(ie)%name, lat%branch(ix_branch)%ele(ie)%s
          call lat_make_mat6(lat, ix_ele=i, ref_orb =co, ix_branch=ix_branch)
       end do
       
       call twiss_at_start(lat, status=status, ix_branch=ix_branch)
       delta_e = 1.e-4
       call chrom_calc(lat, delta_e, chrom_x, chrom_y, ix_branch = ix_branch)
       lun= lunget()
       open(unit=lun,file = filename)
       print '(a,es12.4)', 'Energy = ',lat%ele(0)%value(E_TOT$)
       print '(10a12)',' ','Q_x','Q_y','beta_x','beta_y','alpha_x','alpha_y','eta_x', 'chrom_x','chrom_y'
       print '(a,9f12.4)','Closed ring:',lat%a%tune/twopi, lat%b%tune/twopi, 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, chrom_x, chrom_y
       print '(a,es16.8)',' Spin tune = ',lat%param%spin_tune/twopi
       f_rev=c_light*co(0)%beta/(lat%ele(lat%n_ele_track)%s-co(lat%n_ele_track)%vec(5))
       print '(a,es16.8, a, es16.8)',' f_rev = v/l ', f_rev,' f_rev = 1/t ',1./co(lat%n_ele_track)%t
       print '(a,es16.8,a,es16.8)',' f_a = ', f_rev*lat%param%spin_tune/twopi,' or ', lat%param%spin_tune/twopi/co(lat%n_ele_track)%t

       write(lun,'(a)')lat_file
       write(lun, '(a,es12.4)') 'Energy = ',lat%ele(0)%value(E_TOT$)
       write(lun, '(10a12)')' ','Q_x','Q_y','beta_x','beta_y','alpha_x','alpha_y','eta_x', 'chrom_x','chrom_y'
       write(lun, '(a,9f12.4)')'Closed ring:',lat%a%tune/twopi, lat%b%tune/twopi, 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, chrom_x, chrom_y
       write(lun, '(a,es16.8)')' Spin tune = ',lat%param%spin_tune/twopi
       f_rev=c_light*co(0)%beta/(lat%ele(lat%n_ele_track)%s+co(lat%n_ele_track)%vec(5))
       write(lun, '(a,es16.8)')' f_rev = ', f_rev
       write(lun, '(a,es16.8)')' f_a = ', f_rev*lat%param%spin_tune/twopi
       close(lun)
       
       print '(a,$)',' Initialize twiss ,(or use closed ring values ) (y/n)  ? '
     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 (separate values with comma)'
      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)%x%eta
     endif
     call twiss_propagate_all(lat, ix_branch=ix_branch)
     do ie = 1, lat%branch(ix_branch)%n_ele_track
        print '(a,a,2es12.4)',' beta x, beta_y', lat%branch(ix_branch)%ele(ie)%a%beta,lat%branch(ix_branch)%ele(ie)%b%beta
     end do
     
     print '(a,6es12.4)',' Closed orbit:', co(0)%vec
     print '(a,$)',' Initialize orbit, (or use closed orbit values) (y/n) ? '
     read(5,*)ans
     if(index(ans,'y') /= 0 .or. index(ans,'Y')/= 0)then
        print '(a)','x, px, y,py (separate values with comma)'
        read(5,*)co(0)%vec(1:4)
        print '(a,6es12.4)','Initial coordinates = ',co(0)%vec
        call track_all(lat,co, ix_branch)
    endif
   
       dist = 0
       lun = lunget()
       open(unit=lun,file = trim(filename)//'.beta_eta_x_y')
       write(lun,'(a)')lat_file
       write(lun, '(a16,4a12,3a14)')'element','s','beta x','beta y','eta','x','y',' delta l'
       do
         dist = dist + 0.1
         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 = ix_branch) !, ix_branch=nbranch, use_last=.false.)
         write(lun,'(a16,4f12.4,3e14.6)')ele%name,dist,ele%a%beta, ele%b%beta, ele%x%eta, orb_at_s%vec(1), orb_at_s%vec(3), orb_at_s%vec(5)
       end do
      close(unit=lun)
      call dbeta_ddelta(lat,transfer_line)

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

If(plot_details)then
       
       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 If
end program first_turn

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