program ring_emittance_and_aperture

use bmad_struct
use beam_mod
use dynamic_aperture_mod
use ptc_layout_mod

type (lat_struct), target :: lat
type (branch_struct), pointer :: branch
type (ele_struct), pointer :: ele0
type (coord_struct) :: orb0
type (coord_struct), allocatable :: closed_orb(:)
type (normal_modes_struct) :: normal_mode
type (aperture_struct) :: aperture
type (track_input_struct) :: track_input
type (ele_pointer_struct), allocatable :: eles(:)
type (bmad_normal_form_struct) :: normal_form
type (all_pointer_struct) a_ptr

real(rp) :: emit_a, emit_b
real(rp) :: delta_e, chrom_x, chrom_y
real(rp) :: da_angle, da_pz
real(rp) :: value

integer :: i, ix, i_dim, nlines
integer :: da_n_turns
integer :: namelist_file, outfile, constraintfile, n_char, iu, ios, ix_word, n_loc

character(100) :: lat_name, lat2_name, lat_path, base_name
character(400) ::  in_file, outfile_name, constraintfile_name
character(1000) :: line
character(100) :: ele_value_filename, dummy, ele_and_property, constraint
character(400), allocatable :: locator(:), property(:), constraints(:)
character(1) :: delim
character(100) :: outstr
character(30), parameter :: r_name = 'ring_emittance_and_aperture'

logical :: err, err_flag, delim_found, map_calc_on, rf_on
logical :: verbose, write_constraints

namelist / ring_emittance_and_aperture_params / &
    lat_name, lat2_name, outfile_name, da_n_turns, da_pz, da_angle, &
    ele_value_filename, verbose, map_calc_on, rf_on, write_constraints, constraintfile_name

!------------------------------------------
!Defaults for namelist
lat_name =  'lat.bmad'
lat2_name = ''
ele_value_filename=''
outfile_name = 'merit.txt'
constraintfile_name = 'constraints.txt'
da_n_turns = 100
da_pz = 0
da_angle = 0
verbose = .false.
map_calc_on = .false.
rf_on = .true.
write_constraints = .false.

!Read namelist
in_file = 'ring_emittance_and_aperture.in'
if (command_argument_count() > 0) call get_command_argument(1, in_file)

namelist_file = lunget()
if (verbose) print *, 'Opening: ', trim(in_file)
open (namelist_file, file = in_file, status = "old")
read (namelist_file, nml = ring_emittance_and_aperture_params)
close (namelist_file)

! Read lat2_name
if (command_argument_count() > 1) call get_command_argument(2, lat2_name)

! If a lat2_name is given on the command line, use that for the merit and constraint filenames
if(lat2_name /= '') then
  outfile_name = trim(lat2_name)//'.merit'
  constraintfile_name = trim(lat2_name)//'.constraints'
endif
!Trim filename
n_char= SplitFileName(lat_name, lat_path, base_name) 

! Output file
outfile = lunget()
open(outfile, file = outfile_name)

! Constraint file 
if (write_constraints) then
  constraintfile = lunget()
  open(constraintfile, file = constraintfile_name)
endif

!Parse Lattice
call bmad_parser (lat_name, lat)
!branch => lat%branch(0)

!Parse additional settings
if (lat2_name /= '') then
  if (verbose)  print *, 'Parsing: '//trim(lat2_name)
  call bmad_parser2 (lat2_name, lat)
endif

! Get ele value list: 
! ele_locator_string value string constraints(optional, separated by commas)
! Example: 
! quad::* k1
! sex::* k2 LESS_THAN 40, GREATER_THAN -40
if (ele_value_filename /= '') then
  if (verbose) print *, 'Opening: ', trim(ele_value_filename)
  iu = lunget()
  open(iu, file = ele_value_filename, status ="old")
  nlines = 0
  ! Count number of lines
  do
    read(iu, *, iostat=ios) dummy
    !print *, nlines, dummy, ios
    if (ios == -1) exit
    nlines = nlines + 1
  enddo
  rewind(iu)
  allocate(locator(nlines), property(nlines), constraints(nlines))
  ! Get words.
  delim = ' '
  do i=1, nlines
    read(iu, '(a)') line
    call word_read (line,  ' ', locator(i), ix_word, delim, delim_found, line)
    call word_read (line,  ' ', property(i), ix_word, delim, delim_found, line)
    ! The constraint will be the last part of the line
    constraints(i) = line
  enddo
  !do i=1, nlines
  !  write (*, '(i8, 3a)') i, trim(locator(i)), trim(property(i)), trim(constraints(i))
  !enddo
  close(iu)
endif


!--------------------
! RF on/off
if (rf_on) then
  call set_on_off (rfcavity$, lat, on$)
  if (verbose) print *, 'RF is ON' 
  i_dim=6
else
  call set_on_off (rfcavity$, lat, off$)
  if (verbose) print *, 'RF is OFF'
  i_dim=4
endif

!--------------------
! Closed orbit calc

call reallocate_coord(closed_orb, lat)
call closed_orbit_calc(lat, closed_orb, i_dim)

! orb and element at 0 for convenience
orb0=closed_orb(0)
ele0 => lat%ele(0)


!---------------------
! Emittance calc

! Prepare for radiation integrals
call twiss_at_start(lat)
call twiss_propagate_all(lat)
!print *, 'radiation_integrals'
call radiation_integrals (lat, closed_orb, normal_mode)
emit_a = normal_mode%a%emittance
emit_b = normal_mode%b%emittance

!---------------------
! Chromaticity calc
call chrom_calc (lat, delta_e, chrom_x, chrom_y, err_flag)

!---------------------
! Dynamic Aperture (DA) calc
track_input%n_turn = da_n_turns
track_input%x_init = 0.005 
track_input%y_init = .001 
track_input%accuracy = .00001 
call dynamic_aperture1(lat, orb0, da_angle, track_input, aperture, da_pz)


!---------------------
! Map calc (PTC)
if (map_calc_on) then
  call lat_to_ptc_layout (lat)
  call ptc_one_turn_map_at_ele (lat%ele(0), normal_form%M, pz = 0.0_rp, order = bmad_com%taylor_order)
  call normal_form_taylors(normal_form%M, rf_on, dhdj = normal_form%dhdj, A_t = normal_form%A, A_t_inverse = normal_form%A_inv)
  !TODO: Use these in the merit
endif

!---------------------
! Write to file(s)

! Special merits
write (outfile, '(a, es15.7)') 'emit_a', emit_a
write (outfile, '(a, es15.7)') 'chrom_x', chrom_x
write (outfile, '(a, es15.7)') 'chrom_y', chrom_y
write (outfile, '(a, es15.7)') 'aperture_x', aperture%x
write (outfile, '(a, es15.7)') 'aperture_y', aperture%y

! Special constraints, hard-coded
if (write_constraints) then
  write (constraintfile, '(a)') 'chrom_x GREATER_THAN -1'
  write (constraintfile, '(a)') 'chrom_x LESS_THAN 1'
  write (constraintfile, '(a)') 'chrom_y GREATER_THAN -1'
  write (constraintfile, '(a)') 'chrom_y LESS_THAN 1'
endif

! Element values, with optional constraints
if (ele_value_filename /= '') then
do i=1, nlines
  call lat_ele_locator (locator(i), lat, eles, n_loc, err)
  if (n_loc == 0) then
    if (verbose) print *, 'Warning: No elements found for: ', trim(locator(i))
    cycle
  endif
  if (err) then
    if (verbose) print *, 'Error parsing locator string: ', trim(locator(i))
    cycle
  endif
  do ix=1, n_loc
  ! Get attributes
    call  pointer_to_attribute (eles(ix)%ele, upcase(property(i)), .false., a_ptr, err)
    if (err) then
      if (verbose) print *, 'Problem with property: ', trim(eles(ix)%ele%name), ' : ', trim(property(i))
      cycle
    endif
    value = value_of_all_ptr(a_ptr)
    !print *, trim(eles(ix)%ele%name), ' : ', trim(property(i)), ' = ', value
    
    ! Final output
    ele_and_property = trim(eles(ix)%ele%name)//'['//trim(property(i))//']'
    write (outfile, '(a, es15.7)') trim(ele_and_property) , value
    
    ! Constraints
    if (write_constraints) then
      line = constraints(i)
      do  
        call word_read (line,  ',', constraint, ix_word, delim, delim_found, line)
        write (constraintfile, '(3a)') trim(ele_and_property), ' ', trim(constraint)
        if (.not. delim_found) exit
      enddo
    endif
    
  enddo
enddo
endif


!---------------------
! Cleanup
close(outfile)
if (write_constraints) then
  close(constraintfile)
  if (verbose) print *, 'Wrote: ', trim(constraintfile_name)
endif

if (verbose) print *, 'Wrote: ', trim(outfile_name)

end program