module tune_scanner_mod

use bmad_struct
use dynamic_aperture_mod
use quick_plot
!use beam_mod

implicit none

type tune_info_struct
  real(rp) :: Qa, Qb, Qz
  type(aperture_struct) :: aperture
end type

type tune_block_struct
  real(rp) :: Qa_min = 0.01_rp
  real(rp) :: Qb_min = 0.01_rp
  real(rp) :: Qa_max = 0.99_rp
  real(rp) :: Qb_max = 0.99_rp
  real(rp) :: pz_init = 0.0_rp
  integer :: nQa = 10
  integer :: nQb = 10
  integer :: n_turns = 100
  type(tune_info_struct), allocatable :: info(:)
end type

contains


!---------------------------------------------------------------------------
!+
!  
!-
!---------------------------------------------------------------------------
subroutine write_tune_info(iu, tune_info, header)
implicit none
integer :: iu
type (tune_info_struct) :: tune_info
logical, optional :: header
!
if (present(header)) then
  if (header) call write_tune_info_header(iu)
endif
write (iu, '(3f10.7, 2es15.7, i8)') tune_info%Qa, tune_info%Qb, tune_info%Qz, &
                 tune_info%aperture%x, tune_info%aperture%y, tune_info%aperture%i_turn

end subroutine

subroutine write_tune_info_header(iu)
implicit none
integer :: iu
write (iu, '(3a10, 2a15, a8)') 'Qa', 'Qb', 'Qz', &
    							   'x_aperture(m)', 'y_aperture(m)', 'i_turn'
end subroutine

!---------------------------------------------------------------------------
!+
!  
!-
!---------------------------------------------------------------------------
subroutine tune_scanner_plot_init (page_type, Qa_min, Qa_max, Qb_min, Qb_max, plot_id)

implicit none

real(rp) :: Qa_min, Qb_min, Qa_max, Qb_max
integer :: plot_id
character(*) :: page_type

select case (page_type)
case('X')
  call qp_open_page ('X', plot_id, 500.0_rp, 500.0_rp, 'POINTS')
case('PDF')
  call qp_open_page ('PDF', plot_id, 500.0_rp, 500.0_rp, 'POINTS', plot_file='tune_scanner.pdf')
case default 
  call qp_open_page ('X', plot_id, 500.0_rp, 500.0_rp, 'POINTS')
end select

call qp_set_page_border (0.1_rp, 0.1_rp, 0.1_rp, 0.1_rp, '%PAGE')
call qp_set_margin (0.1_rp, 0.1_rp, 0.1_rp, 0.1_rp, '%PAGE')
call qp_set_line_width_basic(0)
call qp_set_line_attrib ('PLOT', color = black$)
call qp_set_line_width_basic(0)

call qp_calc_and_set_axis ('X', Qa_min, Qa_max, 3, 7, 'GENERAL')
call qp_calc_and_set_axis ('Y', Qb_min, Qb_max, 3, 7, 'GENERAL')
call qp_draw_axes ('Q\da\u', 'Q\db\u' )

end subroutine
!---------------------------------------------------------------------------
!+
!  
!-
!---------------------------------------------------------------------------
subroutine plot_tune_block(tune_block, x_aperture_max)
implicit none
type (tune_block_struct) tune_block
real(rp) :: Qa_set, Qb_set, x_aperture, x_aperture_max, dQa, dQb,  color
integer :: i

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

do i=1, size(tune_block%info(:))
    Qa_set = tune_block%info(i)%Qa
    Qb_set = tune_block%info(i)%Qb
    x_aperture = tune_block%info(i)%aperture%x
    color =  1.0_rp - x_aperture/ x_aperture_max
    !print *, 'x_aperture, color: ', x_aperture, color
    !print *, Qa_set-dQa/2, Qa_set+dQa/2, Qb_set-dQb/2, Qb_set+dQb/2, qp_continuous_color(color) 
    call qp_paint_rectangle (Qa_set-dQa/2, Qa_set+dQa/2, Qb_set-dQb/2, Qb_set+dQb/2, color=qp_continuous_color(color) )
end do

end subroutine



!---------------------------------------------------------------------------
!+
! Function best_tune_from_array(tune_block_arr) result(best_tune)
!
! Picks the tune with the best info%aperture%x
! 
! Input: 
!   tune_block_arr(:)    -- tune_block_struct: tunes that have been set
!
! Output:
!   best_tune            -- tune_info_struct: tune info with the best %aperture%x
!
!-
!---------------------------------------------------------------------------
function best_tune_from_array(tune_block_arr) result(best_tune)
implicit none
type (tune_block_struct) :: tune_block_arr(:)
type (tune_info_struct) :: best_tune
real(rp) :: x_aperture_max
integer :: i, j

!

x_aperture_max = 0
do i=1, size(tune_block_arr)
  do j=1, size(tune_block_arr(i)%info)
    if (tune_block_arr(i)%info(j)%aperture%x > x_aperture_max) then
      best_tune = tune_block_arr(i)%info(j)
      x_aperture_max = best_tune%aperture%x 
    endif
  enddo
enddo

end function

!---------------------------------------------------------------------------
!+
!  
! Input:
!   block0        -- tune_block_struct: Original tune block struct
!   Na            -- integer: Number of sub-blocks in a-direction
!   Nb            -- integer: Number of sub-blocks in b-direction
!
! Output
!   block_arr(1:Na*Nb)  -- tune_block_struct: 
!
!-
!---------------------------------------------------------------------------
subroutine split_tune_block(block0, Na, Nb, block_arr )
implicit none
type (tune_block_struct) :: block0
type (tune_block_struct), allocatable :: block_arr(:)

real(rp) :: dBa, dBb, dQa, dQb
integer :: Na, Nb, ix, ia, ib, nQa, nQb
!

if (.not. allocated(block_arr)) then
  allocate(block_arr(Na*Nb))
else
  if (size(block_arr) /= Na*Nb) then
    print *, 'block array already allocated, wrong sizes:, ', size(block_arr), Na*Nb
    stop
  endif
endif

block_arr(:)%n_turns = block0%n_turns
block_arr(:)%pz_init = block0%pz_init

! numbers in each block (integer division). Make sure there are at least two
nQa = max(block0%nQa/Na, 2)
nQb = max(block0%nQb/Nb, 2)

block_arr(:)%nQa = nQa
block_arr(:)%nQb = nQb

! fine intervals
dQa = (block0%Qa_max - block0%Qa_min)/(nQa*Na-1)
dQb = (block0%Qb_max - block0%Qb_min)/(nQb*Nb-1)

! Block intervals
dBa = (nQa-1)*dQa
dBb = (nQb-1)*dQb

do ia = 1, Na
  do ib = 1, Nb
    ix = ia + (ib-1)*Na
    block_arr(ix)%Qa_min = block0%Qa_min + (ia-1)*(dBa + dQa)
    block_arr(ix)%Qa_max = block0%Qa_min + (ia)*(dBa) + (ia-1)*dQa
    block_arr(ix)%Qb_min = block0%Qb_min + (ib-1)*(dBb + dQb)
    block_arr(ix)%Qb_max = block0%Qb_min + (ib)*(dBb) + (ib-1)*dQb
  end do
end do

end subroutine




!---------------------------------------------------------------------------
!+
!  
!-
!---------------------------------------------------------------------------
subroutine scan_tune_block_parallel(lat, tune_block)

!$ use omp_lib

implicit none

type (lat_struct) :: lat
type (lat_struct), allocatable :: omp_lat(:), omp_lat0(:)
type (tune_block_struct) :: tune_block
type (aperture_struct), allocatable :: omp_aperture(:)

real(rp) :: Qa_max, Qb_max, Qa_min, Qb_min, dQa, dQb
real(rp) :: Qa_set, Qb_set, Qz_set, pz_init
integer :: nQa, nQb, n_turns
integer :: omp_n, omp_i, i, ia, ib, ix, i_dim

logical, allocatable :: omp_tune_set_ok(:)
logical :: err
!

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

! 4D or 6D calc?
if (rf_is_on(lat%branch(0))) then
  i_dim = 6
else
  i_dim = 4
endif


! Info setup
if ( allocated(tune_block%info)) then
  print *, 'tune block already allocated'  
  if (size(tune_block%info) /= nQa*nQb) then
    deallocate(tune_block%info)
    allocate(tune_block%info(nQa*nQb))
  endif
else
  allocate(tune_block%info(nQa*nQb))
endif

! Parallel memory setup
omp_n = 1
!$ omp_n = omp_get_max_threads()

allocate(omp_aperture(omp_n),  omp_tune_set_ok(omp_n))
!Set up independent lattices
allocate(omp_lat(omp_n), omp_lat0(omp_n))
do i=1, omp_n
  !omp_lat(i) = lat
  !omp_lat0(i) = lat
  call lat_equal_lat (omp_lat(i), lat)
  call lat_sanity_check( omp_lat(i), err)  
  !print *, 'Sanity check omp_lat, err:  ', i, err
  call lat_equal_lat (omp_lat0(i), lat)
  call lat_sanity_check( omp_lat0(i), err)
  !print *, 'Sanity check omp_lat0, err:  ', i, err
  omp_lat0(i)%version = i
end do

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

!$OMP parallel &
!$OMP default(private), &
!$OMP shared(nQa, nQb, dQa, dQb, Qa_min, Qb_min, i_dim, pz_init, tune_block), &
!$OMP shared(omp_lat, omp_lat0, omp_tune_set_ok, omp_aperture, n_turns)
!$OMP do schedule(dynamic)
do ia = 0, nQa-1  
  omp_i = 1
  !$ omp_i = omp_get_thread_num()+1
  !if (omp_i==1) write(*, '(i0, a, i0)') ia, '/',  nQa-1
  !print *, omp_i, ' started for ia = ', ia
  ! Reset lattice
  omp_lat(omp_i) = omp_lat0(omp_i) 
  !call lat_equal_lat (omp_lat(omp_i), omp_lat0(omp_i))
  !call transfer_lat(omp_lat0(omp_i), omp_lat(omp_i) )
  !print *, 'parallel Sanity check for ', omp_i
  call lat_sanity_check(omp_lat(omp_i), omp_tune_set_ok(omp_i))
  if (omp_tune_set_ok(omp_i)) print *, 'INSANE: ', omp_i
  do ib = 0, nQb-1
    !print *, ib
    Qa_set = ia*dQa + Qa_min
    Qb_set = ib*dQb + Qb_min
    Qz_set = -1.0_rp !dummy
 
    call set_tune_and_track (omp_lat(omp_i), Qa_set, Qb_set, Qz_set, i_dim, pz_init, n_turns, &
                             omp_tune_set_ok(omp_i), omp_aperture(omp_i))
    ! Save results to info
    ix = ia + ib*nQa + 1 !Flattened matrix index
    tune_block%info(ix)%Qa = omp_lat(omp_i)%a%tune/twopi
    tune_block%info(ix)%Qb = omp_lat(omp_i)%b%tune/twopi
    tune_block%info(ix)%Qz = omp_lat(omp_i)%z%tune/twopi
    tune_block%info(ix)%aperture = omp_aperture(omp_i)
    !write (*, '(5f10.7)') Qa_set, tune_block%info(ix)%Qa, Qb_set, tune_block%info(ix)%Qb, 1000*omp_aperture(omp_i)%x
  end do
end do
!$OMP end do
!$OMP end parallel

! Cleanup
do i=1, omp_n
  call deallocate_lat_pointers (omp_lat0(i))
  call deallocate_lat_pointers (omp_lat(i))
end do
deallocate(omp_lat, omp_lat0, omp_aperture, omp_tune_set_ok)

end subroutine


!---------------------------------------------------------------------------
!+
!  
!-
!---------------------------------------------------------------------------
subroutine set_tune_and_track (lat, Qa_set, Qb_set, Qz_set, i_dim, pz_init, n_turns, tune_set_ok, aperture)

!$ use omp_lib
use bookkeeper_mod

implicit none

type (lat_struct) :: lat, save_lat
type (coord_struct), allocatable :: closed_orb(:)
!type (normal_modes_struct) :: normal_mode
type (track_input_struct) :: track_input
type (aperture_struct) :: aperture

real(rp), allocatable :: dk1(:)
real(rp) :: theta_xy, sigma_x, sigma_y, pz_init
real(rp) :: Qa_set, Qb_set, Qz_set, Qa0, Qb0, Qz0

integer :: i_dim, n_turns, omp_i

logical tune_set_ok, verbose

!

verbose = .true.
!$ verbose = .false.

Qa0 = lat%ele(lat%n_ele_track)%a%phi/twopi
Qb0 = lat%ele(lat%n_ele_track)%b%phi/twopi
Qz0 = lat%ele(lat%n_ele_track)%z%phi/twopi

omp_i = 1
!$ omp_i = omp_get_thread_num()+1

! Setup
allocate(dk1(lat%n_ele_max))
call reallocate_coord(closed_orb, lat)

! pick k1 eles to vary
call choose_k1(lat, dk1)

if (verbose) print *, 'closed_orbit_calc(lat, closed_orb, i_dim)', omp_i
call closed_orbit_calc(lat, closed_orb, i_dim)
if (verbose) print *, 'finished closed_orbit_calc(lat, closed_orb, i_dim)', omp_i
!if (verbose) write (*, '(a, 3f12.7)') 'Tunes before: ', Qa0, Qb0, Qz0
!if (verbose) write (*, '(a, 3f12.7)') 'Trying to set: ', int(Qa0) + Qa_set, int(Qb0) + Qb_set, int(Qz0) + Qz_set

! Save lattice in case of a bad set
save_lat = lat
!call transfer_lat(lat, save_lat)
! Set Tune
call set_tune ( (int(Qa0) + Qa_set)*twopi, (int(Qb0) + Qb_set)*twopi, dk1, lat, closed_orb, tune_set_ok)
if (.not. tune_set_ok) then
  !call transfer_lat(save_lat, lat)
  lat = save_lat
  lat%a%tune=Qa0*twopi
  lat%b%tune=Qb0*twopi
  lat%z%tune=Qz0*twopi
  aperture%x = 0
  aperture%y = 0
  aperture%i_turn=1
  return
endif

if (verbose) write (*, '(a, 3f10.7)') 'Tunes after:', lat%a%tune/twopi, lat%b%tune/twopi, lat%z%tune/twopi


!call radiation_integrals (lat, closed_orb, normal_mode)
!sigma_x = sqrt(lat%ele(0)%a%beta*normal_mode%a%emittance +  (lat%ele(0)%x%eta*normal_mode%sigE_E)**2)
!sigma_y = sqrt(lat%ele(0)%b%beta*normal_mode%a%emittance +  (lat%ele(0)%y%eta*normal_mode%sigE_E)**2)

!n_turns = 5000
track_input%n_turn = n_turns
track_input%x_init = .0001_rp
track_input%y_init = .0001_rp
track_input%accuracy = .00001 
call dynamic_aperture (lat, closed_orb(0), theta_xy, track_input, aperture, pz_init)

if (verbose) write (*, '(a)') '__________________________'
!if (verbose) write (*, '(a, f15.7, a)') 'a emittance', normal_mode%a%emittance*1e9, ' nm'
!if (verbose) write (*, '(a, f15.7, a)') 'b emittance', normal_mode%b%emittance*1e9, ' nm'
if (verbose) write (*, '(a, i8, x, a)') 'i_turn', aperture%i_turn, trim(lat%ele(aperture%ix_lat)%name)
!if (verbose) write (*, '(a, f15.7, a)') 'x aperture/sigma_x',  aperture%x/sigma_x, '1'

! Cleanup
deallocate(dk1)
call deallocate_lat_pointers(save_lat)

end subroutine


!----------------------------------------------------------------------------
!+
! Subroutine choose_k1
!
! Modified version of bmadz's 'choose quads' to quads and ...
!-
!----------------------------------------------------------------------------

subroutine choose_k1(lat, dk1)

use bmad_struct
use bmad_interface

implicit none

type (lat_struct) lat
type (ele_struct). pointer :: slave

integer i, j, is

real(rp) :: dk1(:)

logical found

! find which quads to change

do i = 1, lat%n_ele_max
  if (lat%ele(i)%key == quadrupole$ .and. abs(lat%ele(i)%value(tilt$)) < 0.01 .and. &
                         attribute_free(lat%ele(i), 'K1', .false.)) then
    if (lat%ele(i)%a%beta > lat%ele(i)%b%beta) then
      dk1(i) = +1
    else
      dk1(i) = -1
    endif
  else
    dk1(i) = 0
  endif

if (lat%ele(i)%lord_status == overlay_lord$) then
  found = .false.
  do is = 1, lat%ele(i)%n_slave
    slave => pointer_to_slave(lat%ele(i), is, j)
 	  if (lat%control(j)%ix_attrib == k1$ .and. slave%key == quadrupole$ .and. slave%value(tilt$) == 0) then
	    found = .true.
	    exit
	  endif
  enddo
  if (.not. found) cycle
  if (slave%a%beta > slave%b%beta) then
	dk1(i) = +1
  else
	dk1(i) = -1
  endif
endif
enddo

end subroutine choose_k1



end module tune_scanner_mod