module tt_tools_mod

  use bmad
  use bmadz_interface
  use mode3_mod
  use tune_tracker_mod
  
contains

!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine tt_init(ring, orb, tt_params, n_TTs, id)
!
! Subroutine to initialize the tune trackers
!
! Input:
!   ring       -- lat_struct; Lattice information
!   orb        -- coord_struct; orbit array
!   tt_params  -- tt_param_struct;  tune-tracker parameters
!   n_TTs      -- integer;  number of tune trackers, maxm=3
! Output:
!   id(max_tt) -- the ID number for tune trackers
!--------------------------------------------------------------------------
subroutine tt_init(ring, orb, tt_params, n_TTs, id)

  implicit none

  type (lat_struct) ring
  type (coord_struct), allocatable :: orb(:)
  TYPE(tt_param_struct) :: tt_params(max_tt) !max_tt=3 set on tune_tracker module
  integer n_TTs
  integer id(max_tt)
  integer i
  real(rp) Tring, Deltaphi
  if (.not. allocated(orb)) allocate(orb(0:ring%n_ele_track))

  Tring = ring%ele(ring%n_ele_track)%s / c_light

 DO i=1,n_TTs
     !calculate initial frequency of modulator
     WRITE(*,'(A14,I3,A21,F8.5,A22)') &
          "Tune tracker #",i, " VCO base frequency: ", tt_params(i)%modTfrac0, " oscillations per turn"
     tt_params(i)%modw0 = 2.0_rp*pi/(Tring/tt_params(i)%modTfrac0)   ! 2pi/(initial period of modulator)

     ! Determine phase between BPM and kicker
     IF(tt_params(i)%orientation == 'h') THEN
        !The following calculation works because the PLL leads the BPM data by pi/2
        Deltaphi = ring%ele(ring%n_ele_track)%mode3%a%phi - &
             ring%ele(tt_params(i)%bpm_loc)%mode3%a%phi + ring%ele(tt_params(i)%kck_loc)%mode3%a%phi
        tt_params(i)%Onum = 1  !element of coord_struct for BPM to observe
        !Normalize kick amplitude by beta at kicker
        tt_params(i)%kickAmplitude = tt_params(i)%kickAmplitude / SQRT(ring%ele(tt_params(i)%kck_loc)%mode3%a%beta)
     ELSEIF(tt_params(i)%orientation == 'v') THEN
        !The following calculation works because the PLL leads the BPM data by pi/2
        Deltaphi = ring%ele(ring%n_ele_track)%mode3%b%phi - &
             ring%ele(tt_params(i)%bpm_loc)%mode3%b%phi + ring%ele(tt_params(i)%kck_loc)%mode3%b%phi
        tt_params(i)%Onum = 3  !element of coord_struct for BPM to observe
        !Normalize kick amplitude by beta at kicker
        tt_params(i)%kickAmplitude = tt_params(i)%kickAmplitude / SQRT(ring%ele(tt_params(i)%kck_loc)%mode3%b%beta)
     ELSEIF(tt_params(i)%orientation == 'l') THEN
        Deltaphi = -pi/2.0_rp + ring%ele(ring%n_ele_track)%mode3%c%phi - &
             ring%ele(tt_params(i)%bpm_loc)%mode3%c%phi + ring%ele(tt_params(i)%kck_loc)%mode3%c%phi
        tt_params(i)%Onum = 1  !element of coord_struct for BPM to observe
     ELSE
        write(*,*) "Tune tracker orientation improperly defined! Must be one of (h,v,l). Stopping here..."
        stop
     ENDIF
     tt_params(i)%phi_to_kicker = Deltaphi
     tt_params(i)%phi_to_kicker = MOD(tt_params(i)%phi_to_kicker,2.0*pi)

     tt_params(i)%Dt = Tring
     tt_params(i)%offset = orb(tt_params(i)%bpm_loc)%vec(tt_params(i)%Onum)

     id(i) = init_dTT(tt_params(i), orb(0))
  ENDDO

! Print information about the location of each TT's kicker and BPM
  DO i=1,n_TTs
     WRITE(*,'(A,I2,A,A,A,A)') "Kicker #", i, " at element: ", ring%ele(tt_params(i)%kck_loc)%name, &
          " Type: ", key_name(ring%ele(tt_params(i)%kck_loc)%key)
     WRITE(*,'(A,I2,A,A,A,A)') "   BPM #", i, " at element: ", ring%ele(tt_params(i)%bpm_loc)%name, &
          " Type: ", key_name(ring%ele(tt_params(i)%bpm_loc)%key)
  ENDDO

end subroutine tt_init


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine check_tt_kickers(ring, tt_params, n_TTs)
!
! Subroutine to check whether kickers for the tune trackers are set correctly.
! If necessary, also translates user-supplied kicker and BPM names into ele
! locations.
!
! Input:
!   ring       -- lat_struct; Lattice information
!   tt_params  -- tt_param_struct;  tune-tracker parameters
!   n_TTs      -- integer;  number of tune trackers, maxm=3
!--------------------------------------------------------------------------
subroutine check_tt_kickers(ring, tt_params, n_TTs)

  implicit none

  type (lat_struct) ring
  TYPE(tt_param_struct) :: tt_params(max_tt) !max_tt=3 set on tune_tracker module
  integer n_TTs
  integer i, j
 

 DO i=1,n_TTs
    ! check if both bpm_loc and bpm_name are specified (and same for 
    ! kck_loc and kck_name), make sure they agree, otherwise stop
    if ((tt_params(i)%kck_loc .gt. -1) .and. (len(trim(tt_params(i)%kck_name)) .gt. 0)) then
      if (trim(ring%ele(tt_params(i)%kck_loc)%name) .ne. trim(tt_params(i)%kck_name)) then
        write(*,'(a,i0,a,i0,2a)') "ERROR: kck_loc ", tt_params(i)%kck_loc, " for tune tracker ", i, " does not match provided kck_name ", &
                                 tt_params(i)%kck_name
        write(*,*) "TERMINATING"
        stop
      endif
    endif
    if ((tt_params(i)%bpm_loc .gt. -1) .and. (len(trim(tt_params(i)%bpm_name)) .gt. 0)) then
      if (trim(ring%ele(tt_params(i)%bpm_loc)%name) .ne. trim(tt_params(i)%bpm_name)) then
        write(*,'(a,i0,a,i0,2a)') "ERROR: bpm_loc i", tt_params(i)%bpm_loc, " for tune tracker ", i, " does not match provided bpm_name ",&
                                 tt_params(i)%bpm_name
        write(*,*) "TERMINATING"
        stop
      endif
    endif

    ! if bpm_loc or kck_loc not specified, but bpm_name or kck_name are, locate the proper ele indices
    if ((tt_params(i)%kck_loc .eq. -1) .and. (len(trim(tt_params(i)%kck_name)) .gt. 0)) then
      do j = 1, ring%n_ele_track
        if (trim(ring%ele(j)%name) .eq. trim(tt_params(i)%kck_name)) then 
          tt_params(i)%kck_loc = j
          exit
        endif
      enddo
      if (tt_params(i)%kck_loc .eq. -1) then ! couldn't match ele name
        write(*,'(3a,i2,a)') "ERROR: kck_name ", trim(tt_params(i)%kck_name), " for tune tracker ", i, " does not match any element!"
        write(*,*) "TERMINATING"
        stop
      endif
    endif
    if ((tt_params(i)%bpm_loc .eq. -1) .and. (len(trim(tt_params(i)%bpm_name)) .gt. 0)) then
      do j = 1, ring%n_ele_track
        if (trim(ring%ele(j)%name) .eq. trim(tt_params(i)%bpm_name)) then 
          tt_params(i)%bpm_loc = j
          exit
        endif
      enddo
      if (tt_params(i)%bpm_loc .eq. -1) then ! couldn't match ele name
        write(*,'(a,i2,a)') "ERROR: bpm_name for tune tracker ", i, " does not match any element!"
        write(*,*) "TERMINATING"
        stop
      endif
    endif
    
    IF( ring%ele(tt_params(i)%kck_loc)%key == hkicker$ ) THEN
      IF( tt_params(i)%orientation /= 'h' ) THEN
        WRITE(*,'(A,I2,A)') "ERROR: Orientation of tune tracker ", i," does not match kicker orientation"
        WRITE(*,'(A)') "TERMINATING"
        STOP
      ENDIF
      IF( .not. attribute_free(tt_params(i)%kck_loc,'KICK',ring,.false.) ) THEN
        WRITE(*,'(A,I2,A)') "ERROR: kick attribute for TT ", i, " kicker is not free."
        WRITE(*,'(A)') "TERMINATING"
        STOP
      ENDIF
    ELSEIF( ring%ele(tt_params(i)%kck_loc)%key == vkicker$ ) THEN
      IF( tt_params(i)%orientation /= 'v' ) THEN
        WRITE(*,'(A,I2,A)') "ERROR: Orientation of tune tracker ", i," does not match kicker orientation"
        WRITE(*,'(A)') "TERMINATING"
        STOP
      ENDIF
      IF( .not. attribute_free(tt_params(i)%kck_loc,'KICK',ring,.false.) ) THEN
        WRITE(*,'(A,I2,A)') "ERROR: kick attribute for TT ", i, " kicker is not free."
        WRITE(*,'(A)') "TERMINATING"
        STOP
      ENDIF
    ELSE
      IF( .not. has_hkick_attributes(ring%ele(tt_params(i)%kck_loc)%key) ) THEN
        WRITE(*,'(A,I2,A)') "ERROR: Element for TT ", i, " kicker does not have kick attributes."
        WRITE(*,'(A)') "TERMINATING"
        STOP
      ENDIF
      IF( tt_params(i)%orientation == 'h' ) THEN
        IF( .not. attribute_free(tt_params(i)%kck_loc,'HKICK',ring,.false.) ) THEN
          WRITE(*,'(A,I2,A)') "ERROR: hkick attribute for TT ", i, " kicker is not free."
          WRITE(*,'(A)') "TERMINATING"
          STOP
        ENDIF
      ELSEIF( tt_params(i)%orientation == 'v' ) THEN
        IF( .not. attribute_free(tt_params(i)%kck_loc,'VKICK',ring,.false.) ) THEN
          WRITE(*,'(A,I2,A)') "ERROR: vkick attribute for TT ", i, " kicker is not free."
          WRITE(*,'(A)') "TERMINATING"
          STOP
        ENDIF
      ELSEIF( tt_params(i)%orientation == 'l' ) THEN
        IF( .not. attribute_free(tt_params(i)%kck_loc,'PHI0',ring,.false.) ) THEN
          WRITE(*,'(A,I2,A)') "ERROR: phi0 attribute for TT ", i, " kicker is not free."
          WRITE(*,'(A)') "TERMINATING"
          STOP
        ENDIF
      ENDIF
    ENDIF
  ENDDO

end subroutine check_tt_kickers


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine track_tunes(ring, tt_params, n_TTs, id, tt_tunes, kick0)
!
! Subroutine to track the tunes
!
! Input:
!   ring         -- lat_struct; Lattice information
!   tt_params    -- tt_param_struct;  tune-tracker parameters
!   n_TTs        -- integer;  number of tune trackers, n_TTs<=3
!   id(max_tt)   -- the ID number for tune trackers
!   kick0(n_TTs) -- Optional; initial kick for elements used in tune trackers
! Output:
!   tt_tunes(n_TTs) -- the output tunes from the tune tracker
!--------------------------------------------------------------------------
subroutine track_tunes (ring, orb, tt_params, n_TTs, id, tt_tunes, kick0)

  implicit none

  type (lat_struct) ring
  type (ele_struct), pointer :: ele
  type (coord_struct)  orb(:)
  type (all_pointer_struct) a_ptr
  TYPE(tt_param_struct) :: tt_params(max_tt) !max_tt=3 set on tune_tracker module
  integer n_TTs
  integer id(max_tt)
  real(rp), optional :: kick0(1:max_tt)
  integer bpm_loc, Onum

  real(rp), allocatable :: tt_tunes(:)
  real(rp) bpmdata, sinphi, kick, TTw, Tring
  logical err
  integer j

  if ((.not. present(kick0)) .and. n_TTs .gt. 0) then
     kick0(:) = 0.
  endif

  Tring = ring%ele(ring%n_ele_track)%s / c_light

  if(.not. allocated(tt_tunes)) allocate(tt_tunes(n_TTs))

    DO j=1,n_TTs
       bpmdata = orb(tt_params(j)%bpm_loc)%vec(tt_params(j)%Onum)
       ele => ring%ele(tt_params(j)%kck_loc)
       sinphi = TT_update(bpmdata,id(j))
       kick = sinphi*tt_params(j)%kickAmplitude
       IF( ele%key == hkicker$ ) THEN
          call pointer_to_attribute(ele, 'KICK', .true., a_ptr, err)
       ELSEIF( ele%key == vkicker$ ) THEN
          call pointer_to_attribute(ele, 'KICK', .true., a_ptr, err)
       ELSE
          IF(tt_params(j)%orientation == 'h') THEN
             call pointer_to_attribute(ele, 'HKICK', .true., a_ptr, err)
          ELSEIF(tt_params(j)%orientation == 'v') THEN
             call pointer_to_attribute(ele, 'VKICK', .true., a_ptr, err)
          ELSEIF(tt_params(j)%orientation == 'l') THEN
             call pointer_to_attribute(ele, 'PHI0', .true., a_ptr, err)
          ENDIF
       ENDIF

       a_ptr%r = kick0(j) + kick
       call set_flags_for_changed_attribute(ele, a_ptr%r)

       TTw = get_dTT('wf',id(j))
       tt_tunes(j)=TTw*Tring
    ENDDO

    call lattice_bookkeeper(ring)

end subroutine track_tunes

end module tt_tools_mod