program tune_scanner

use bmad_struct
use beam_mod
use tune_scanner_mod
use quick_plot
use sim_utils_interface
!$ use omp_lib

type (lat_struct), target :: lat
type (lat_struct), allocatable :: omp_lat(:), omp_lat0(:)
type (lat_struct) :: save_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 (tune_block_struct) :: tune_block
type (tune_block_struct), allocatable :: tune_block_arr(:)
type (tune_info_struct) :: best_tune

real(rp) :: Qa_set, Qb_set, Qz_set, dQa, dQb, dQz
real(rp) :: Qa_min, Qb_min, Qz_min, Qa_max, Qb_max, Qz_max, pz_init
real(rp) :: x_aperture, x_aperture_max, color, time
real(rp), allocatable :: x_aperture_mat(:,:)
integer :: ia, ib, iz, nQa, nQb, nQz, nBlocka, nBlockb, n_turns
integer :: i_dim
integer :: namelist_file, n_char, outfile
integer :: omp_n, omp_i, i, j
integer :: plot_id, plot_every_n
integer, allocatable :: omp_counter(:)


character(100) :: lat_name, lat_path, base_name, in_file, outfile_name
character(30), parameter :: r_name = 'tune_scanner'

logical :: tune_set_ok, omp_on, plot_on, write_pdf
logical :: radiation_damping_on, radiation_fluctuations_on, rf_on, parallel

namelist / tune_scanner_params / &
    lat_name, parallel, radiation_damping_on, radiation_fluctuations_on, rf_on, &
    Qa_min, Qb_min, Qz_min, Qa_max, Qb_max, Qz_max, &
    nQa, nQb, nQz,  plot_on, plot_every_n, nBlocka, nBlockb, n_turns, pz_init, write_pdf

!------------------------------------------
!Defaults for namelist
lat_name = 'lat.bmad'
outfile_name='tune_scanner.out'
parallel = .false.
radiation_damping_on = .true. 
radiation_fluctuations_on = .true. 
rf_on = .true. 
nQa = 10
nQb = 10
nQz = 1
Qa_min = 0.01
Qa_max = 0.99
Qb_min = 0.01
Qb_max = 0.99
Qz_min = -1
Qz_max = -1
plot_on = .false.
plot_every_n = 1
nBlocka = 4
nBlockb = 4
n_turns = 100
pz_init = 0.0_rp
write_pdf=.true.


write(*,*) ''
write(*,*) '___            ___     __   __                  ___  __  '
write(*,*) ' |  |  | |\ | |__     /__` /  `  /\  |\ | |\ | |__  |__) '
write(*,*) ' |  \__/ | \| |___    .__/ \__, /~~\ | \| | \| |___ |  \ '
write(*,*) ''                                            


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

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

!Trim filename
n_char= SplitFileName(lat_name, lat_path, base_name) 

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

! Prepare Outfile
outfile = lunget()
open(outfile, file = outfile_name)

!-------------------
! OpenMP
!$ write(*, '(a)')  '________________________________________'
!$ write(*, '(a)')  'OpenMP enabled'
!$ write(*, '(a, i8)') 'max threads: ', omp_get_max_threads()

!--------------------
! Radiation 
write(*, '(a)')  '________________________________________'
write(*, '(a)')  'Lattice'
write(*, '(a)') trim(lat%input_file_name)
write(*, '(a)') trim(lat%lattice)
bmad_com%radiation_damping_on = radiation_damping_on
bmad_com%radiation_fluctuations_on = radiation_fluctuations_on
print *, 'bmad_com%radiation_damping_on: ', bmad_com%radiation_damping_on
print *, 'bmad_com%radiation_fluctuations_on: ', bmad_com%radiation_fluctuations_on

!--------------------
! RF on/off
if (rf_on) then
  call set_on_off (rfcavity$, lat, on$)
  print *, 'RF is ON' 
  i_dim=6
  if(.not. rf_is_on(lat%branch(0))) print *, 'Error with RF on/off'
else
  call set_on_off (rfcavity$, lat, off$)
  print *, 'RF is OFF'
  i_dim=4
  if(rf_is_on(lat%branch(0))) print *, 'Error with RF on/off'
endif


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

call reallocate_coord(closed_orb, lat)
!allocate(orb0(0:lat%n_ele_track))

call closed_orbit_calc(lat, closed_orb, i_dim)

! Prepare for radiation integrals
call twiss_at_start(lat)
call twiss_propagate_all(lat)

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

! Radiation integral info
call radiation_integrals (lat, closed_orb, normal_mode)
!write (*, '(a)') '__________________________'
write (*, '(a, f15.7, a)') 'a emittance', normal_mode%a%emittance*1e9, ' nm'
write (*, '(a, f15.7, a)') 'b emittance', normal_mode%b%emittance*1e9, ' nm'
write (*, '(a, 3es15.7)') 'J damp     a, b, z:', normal_mode%a%j_damp, normal_mode%b%j_damp, normal_mode%z%j_damp
write (*, '(a, 3es15.7)') 'alpha damp a, b, z:', normal_mode%a%alpha_damp, normal_mode%b%alpha_damp, normal_mode%z%alpha_damp
write (*, '(a, f10.5, a)') 'Qa', lat%a%tune/(2*pi), ' (lat%a%tune/(2*pi))'
write (*, '(a, f10.5, a)') 'Qb', lat%b%tune/(2*pi), ' (lat%b%tune/(2*pi))'

Qa_set = .7_rp
Qb_set = .123_rp
!Qa_set = .70_rp
!Qb_set = .123_rp

if (parallel) then

tune_block%Qa_min = Qa_min
tune_block%Qb_min = Qb_min
tune_block%Qa_max = Qa_max
tune_block%Qb_max = Qb_max
tune_block%nQa = nQa
tune_block%nQb = nQb
tune_block%n_turns = n_turns
tune_block%pz_init = pz_init

dQa = (Qa_max - Qa_min)/(nQa-1)
dQb = (Qb_max - Qb_min)/(nQb-1)

if (plot_on) then
  call tune_scanner_plot_init ('X', Qa_min-dQa/2, Qa_max+dQa/2, Qb_min-dQb/2, Qb_max+dQb/2, plot_id)
endif

! Split tune_block into job blocks
call split_tune_block(tune_block, nBlocka, nBlockb, tune_block_arr )
write (*, '(a)') '__________________________'
print *, 'Tune blocks: ', nBlocka, nBlockb
print *, 'block sizes: ', tune_block_arr(1)%nQa, tune_block_arr(1)%nQb


! Headers 

write (*, '(a)') '__________________________'
write (*, '(a)') 'Best tunes'
call write_tune_info_header(6)
call write_tune_info_header(outfile)

x_aperture_max = 0

! Start timer
call run_timer ('START') 
do i=1, size(tune_block_arr)
  !----------------------
  ! Main parallel routine
  call scan_tune_block_parallel(lat, tune_block_arr(i))
  !----------------------
  call run_timer ('READ', time)
  write(*, '(a, f10.2)') 'Estimated time remaining (min): ', time/60*(size(tune_block_arr)/(1.0_rp*i) - 1.0_rp)
  ! Write to file
  do j=1, size(tune_block_arr(i)%info)
    call write_tune_info(outfile, tune_block_arr(i)%info(j))
  enddo

  best_tune = best_tune_from_array(tune_block_arr(1:i))
  call write_tune_info(6, best_tune)
  x_aperture_max = best_tune%aperture%x 
  !write (*, '(4f10.7)') best_tune%Qa, best_tune%Qb, best_tune%Qz, 1000*best_tune%aperture%x
  if (plot_on) then
    ! Re-plot all tunes
    do j=1,i
      call plot_tune_block(tune_block_arr(j), x_aperture_max)
      call qp_draw_circle (best_tune%Qa, best_tune%Qb, dQa, color=black$)
    enddo
  endif
enddo
call run_timer ('READ', time)
write(*, '(a, f10.2)') 'Total time (min): ', time/60


else
	! Serial version
	do ia = 0, nQa
	  do ib = 0, nQb
		! Save lattice in case of a problem
		!call transfer_lat(lat, save_lat)
		Qa_set = ia*dQa + Qa_min
		Qb_set = ib*dQb + Qb_min
		Qz_set = -1.0_rp !dummy
		call set_tune_and_track (lat, Qa_set, Qb_set, Qz_set, i_dim, pz_init, n_turns, tune_set_ok, aperture)
		!if (.not. tune_set_ok) call transfer_lat(save_lat, lat)
	  end do
	end do
endif


! Cleanup
close(outfile)

! PDF 
if (write_pdf) then
  call tune_scanner_plot_init ('PDF', Qa_min-dQa/2, Qa_max+dQa/2, Qb_min-dQb/2, Qb_max+dQb/2, plot_id)
  do i=1, size(tune_block_arr)
    do j=1,i
     call plot_tune_block(tune_block_arr(j), x_aperture_max)
     call qp_draw_circle (best_tune%Qa, best_tune%Qb, dQa, color=black$)
    enddo
  enddo
endif


!--------------------
! Plotting control
if (plot_on) then
   write (*, '(a)', advance = 'NO') ' Hit any key to end program: '
   read (*, '(a)') lat_name
  call qp_close_page 
endif

end program