program local_chromaticity

  use bmad_struct
  use bmad_interface
  use bmadz_interface
 
  implicit none

  type (lat_struct) lat
  type (coord_struct), allocatable :: orb(:)
  type (ele_struct) ave

  character*40 lattice
  character*120 lat_file
  character*120 line

!Local variables for scan_params removed
  integer i, ix
  integer nargs, iargc
  real(rp) tot_chrom_x, tot_chrom_y, chrom_x, chrom_y, ks

  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

  call bmad_parser (lat_file, lat)

  call twiss_at_start(lat)
  call twiss_propagate_all(lat)
 ! compute contribution to chromaticity from each element
  write(6, '(a7,4x,a16,1x,8a12)') 'index', 'ele_name','s','ks', 'ave beta x', 'ave beta y', 'chrom_x', 'chrom_y', 'tot_chrom_x', 'tot_chrom_y'
  write(11, '(a7,4x,a16,1x,8a12)') 'index', 'ele_name','s','ks', 'ave beta x', 'ave beta y', 'chrom_x', 'chrom_y', 'tot_chrom_x', 'tot_chrom_y'
  tot_chrom_x=0.
  tot_chrom_y = 0.
  do i=1,lat%n_ele_track
   if(lat%ele(i)%key /= quadrupole$ .and. lat%ele(i)%key /= sol_quad$ .and. lat%ele(i)%key /= sbend$)cycle  
   if(lat%ele(i)%value(k1$) == 0)cycle  
    call twiss_at_element (lat%ele(i), average = ave)
    ks =  lat%ele(i)%value(k1$) * lat%ele(i)%value(l$)
    chrom_x = ave%a%beta*ks/2/twopi       
    chrom_y = ave%b%beta*ks/2/twopi
    tot_chrom_x = tot_chrom_x + chrom_x
    tot_chrom_y = tot_chrom_y + chrom_y
    write(6, '(i10,1x,a16,1x,8es12.4)') i, lat%ele(i)%name,lat%ele(i)%s, ks, ave%a%beta, ave%b%beta, chrom_x, chrom_y, tot_chrom_x, tot_chrom_y
    write(11, '(i10,1x,a16,1x,8es12.4)') i, lat%ele(i)%name,lat%ele(i)%s, ks, ave%a%beta, ave%b%beta, chrom_x, chrom_y, tot_chrom_x, tot_chrom_y
  end do

end program