module sim_eta_mod

use bmad
use sim_bpm_mod
use ele_noise_mod
use bsim_cesr_interface
use radiation_mod

contains

!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine sim_eta_data(ring, bpm, delta, eta)
!
! Subroutine to simulate a basic dispersion measurement, with lattice and BPM
! errors taken into account. 
!
! Input:
!  ring         -- lat_struct
!  bpm(:)       -- det_struct: holds information about misalignments specific to that BPM.
!
! Output:
!  eta(:,2)     -- First index runs over all BPMs. Second index (1:2) = (x,y)
!                  dispersion measurement.
!--------------------------------------------------------------------------
   subroutine sim_eta_data(ring, bpm, delta, eta)
     use sim_bpm_mod

     implicit none

     type(lat_struct) :: ring
     type(det_struct) :: bpm(:)
     type(coord_struct), allocatable :: co(:)
     type (coord_struct), allocatable :: eta1(:), eta2(:), eta(:)
     
     real(rp) :: harvest, current = 1000000
     real(rp), allocatable :: on_off_save(:)
     real(rp), intent(in) :: delta
     integer :: jx, n_bpms
     character(40) lat_type

     lat_type = ring%use_name


     n_bpms = ubound(bpm,1)
     
     if (.not. allocated(co)) allocate(co(0:ring%n_ele_track))
     if (.not. allocated(eta1)) allocate(eta1(0:ring%n_ele_track), eta2(0:ring%n_ele_track))
     if (.not. allocated(eta)) allocate(eta(0:ring%n_ele_track))

     call set_on_off(rfcavity$, ring, save_state$, saved_values = on_off_save)
     call set_on_off(rfcavity$, ring, on$)
     call closed_orbit_calc(ring, co, 6)
     call twiss_at_start(ring)
     call twiss_propagate_all(ring)
     call set_on_off(rfcavity$, ring, restore_state$, saved_values = on_off_save)

     ! dispersion (orbit 1):
     bpm(:)%vec(1) = co(bpm(:)%ix_lat)%vec(1) + (ring%ele(bpm(:)%ix_lat)%x%eta * delta)
     bpm(:)%vec(2) = 0.
     bpm(:)%vec(3) = co(bpm(:)%ix_lat)%vec(3) + (ring%ele(bpm(:)%ix_lat)%y%eta * delta)
     bpm(:)%vec(4) = 0.
     bpm(:)%vec(5) = 0. ! necessary for apply_bpm_errors
     bpm(:)%vec(6) = 0. 

     do jx = 1, n_bpms
        call apply_bpm_errors(bpm(jx), current, lat_type)  ! note: also transfers vec(6)-->amp(4) button signals
        eta1(bpm(jx)%ix_lat)%vec = bpm(jx)%vec
     enddo
     
     
     ! dispersion (orbit 2):
     bpm(:)%vec(1) = co(bpm(:)%ix_lat)%vec(1) - (ring%ele(bpm(:)%ix_lat)%x%eta * delta)
     bpm(:)%vec(2) = 0.
     bpm(:)%vec(3) = co(bpm(:)%ix_lat)%vec(3) - (ring%ele(bpm(:)%ix_lat)%y%eta * delta)
     bpm(:)%vec(4) = 0.
     bpm(:)%vec(5) = 0.
     bpm(:)%vec(6) = 0.

     do jx = 1, n_bpms
        call apply_bpm_errors(bpm(jx), current, lat_type)  ! note: also transfers vec(6)-->amp(4) button signals
        eta2(bpm(jx)%ix_lat)%vec = bpm(jx)%vec
        eta(bpm(jx)%ix_lat)%vec = (eta1(bpm(jx)%ix_lat)%vec - eta2(bpm(jx)%ix_lat)%vec) / (2.*delta)

     enddo

  
end subroutine sim_eta_data


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine sim_dc_eta_data(ring, bpm, delta, eta, ele_noise)
!
! More advanced subroutine to simulate a full DC dispersion measurement, with lattice and BPM
! errors taken into account on a TBT basis. 
!
! Input:
!  ring         -- lat_struct
!  bpm(:)       -- det_struct: holds information about misalignments specific to that BPM
!  n_turns      -- integer: number of turns to track for each orbit (typically 1024)
!  delta        -- real_rp: Desired change in energy
!  ele_noise(:) -- ele_noise_struct(:), optional. Noise to apply to elements when tracking.
!
! Output:
!  eta(:)       -- coord_struct: Dispersion function in both planes. 
!--------------------------------------------------------------------------
   subroutine sim_dc_eta_data(ring, bpm, n_turns, delta, eta, ele_noise)
     use bmad
     use sim_bpm_mod
     implicit none

     type(lat_struct) :: ring
     type(ele_struct), pointer :: ele
     type(det_struct) :: bpm(:)
     type(coord_struct), allocatable :: co(:)
     type (coord_struct), allocatable :: eta1(:), eta2(:), eta(:)
     type(det_data_struct), allocatable :: bpm_save(:,:)
     type(normal_modes_struct) modes
     type(ele_noise_struct), optional :: ele_noise(:)
     
     real(rp) :: harvest, current = 1000000, dfreq
     real(rp), intent(in) :: delta
     real(rp), allocatable :: on_off_save(:)
     integer :: ix, jx, n_bpms, n_turns
     integer :: rad_cache = 0

     integer :: ix_lat
     real(rp) :: alpha_p, rf_freq, gamma
     real(rp) :: d_eta_x, d_eta_y
     real(rp) :: init_vec(6) = 0.
     logical :: err, rad_state

     integer, allocatable :: cavity_ix(:)
     
     if (.not. allocated(bpm_save)) allocate(bpm_save(n_turns,size(bpm)))

     ! initialize on each call:
     d_eta_x = 0.
     d_eta_y = 0.

     ! set absolute time tracking:
     bmad_com%absolute_time_tracking = .true.

     rad_state = bmad_com%radiation_fluctuations_on
     bmad_com%radiation_fluctuations_on = .false.
     call set_on_off(rfcavity$, ring, save_state$, saved_values = on_off_save)
     call set_on_off(rfcavity$, ring, on$)
     call locate_rf(ring, cavity_ix, rf_freq)
     call set_phi0(ring, cavity_ix)
     

     ! set the RF frequency such that delta_x is minimized:
     !call rffm_to_center(ring, rf_freq, cavity_ix)
     
     n_bpms = ubound(bpm,1)
     
     if (.not. allocated(co)) allocate(co(0:ring%n_ele_track))
     if (.not. allocated(eta1)) allocate(eta1(0:ring%n_ele_track), eta2(0:ring%n_ele_track))
     if (.not. allocated(eta)) allocate(eta(0:ring%n_ele_track))

     ! compute dfreq for user-defined delta:
     call de_to_df(ring, rf_freq, delta, dfreq)

     ! dispersion (orbit 1):

     write(*,*) ""
     write(*,'(a,f0.0,a)') "Varying RF by -", dfreq, "Hz"
     call rf_fm(ring, cavity_ix, rf_freq - dfreq, 'target')
     !call set_phi0(ring, cavity_ix)
     
     call closed_orbit_calc(ring, co, 6)

     ! remove this functionality for now, 2013.02.12 JSh
     init_vec(:) = 0.

     if (present(ele_noise)) then
        call sim_tbt_data(ring, bpm, n_turns, .true., n_damping = 0, init_vec=init_vec, bpm_save = bpm_save, ele_noise = ele_noise) 
     else
        call sim_tbt_data(ring, bpm, n_turns, .true., n_damping = 0, init_vec=init_vec, bpm_save = bpm_save) 
     end if

     do ix = 1, size(bpm)
        ix_lat = bpm(ix)%ix_lat
        eta1(ix_lat)%vec(1) = sum(bpm_save(:,ix)%vec(1))/size(bpm_save(:,ix)) ! average orbit over several turns
        eta1(ix_lat)%vec(3) = sum(bpm_save(:,ix)%vec(3))/size(bpm_save(:,ix))
     enddo
     
     ! dispersion (orbit 2):

     write(*,*) ""
     write(*,'(a,f0.0,a)') "Varying RF by +", dfreq, "Hz"
     call rf_fm(ring, cavity_ix, rf_freq + dfreq, 'target')
     !call set_phi0(ring, cavity_ix)
     
     call closed_orbit_calc(ring, co, 6)
    
     ! remove this functionality for now, 2013.02.12 JSh
     !init_vec(:) = init_vec_in(:) ! sim_tbt_data offsets onto the closed orbit
     init_vec(:) = 0.

     if (present(ele_noise)) then
        call sim_tbt_data(ring, bpm, n_turns, .true., n_damping = 0, init_vec=init_vec, bpm_save = bpm_save, ele_noise = ele_noise)
     else
        call sim_tbt_data(ring, bpm, n_turns, .true., n_damping = 0, init_vec=init_vec, bpm_save = bpm_save)
     end if
 
     do ix = 1,  size(bpm)
        ix_lat = bpm(ix)%ix_lat
        eta2(ix_lat)%vec(1) = sum(bpm_save(:,ix)%vec(1))/size(bpm_save(:,ix)) ! average orbit over several turns
        eta2(ix_lat)%vec(3) = sum(bpm_save(:,ix)%vec(3))/size(bpm_save(:,ix))

        ! compute dispersion:
        eta(ix_lat)%vec = (eta1(ix_lat)%vec - eta2(ix_lat)%vec) / (2.*delta)

        ! compute agreement w/ Bmad-reported values:
        d_eta_x = d_eta_x + (eta(ix_lat)%vec(1) - ring%ele(ix_lat)%x%eta)**2
        d_eta_y = d_eta_y + (eta(ix_lat)%vec(3) - ring%ele(ix_lat)%y%eta)**2

     enddo
     
     d_eta_x = sqrt(d_eta_x / size(bpm))
     d_eta_y = sqrt(d_eta_y / size(bpm))

     ! reset things:
     write(*,*) "Resetting RF to nominal frequency..."
     call rf_fm(ring, cavity_ix, rf_freq, 'target')
     bmad_com%radiation_fluctuations_on = rad_state
     call set_on_off(rfcavity$, ring, restore_state$, saved_values = on_off_save)
     bmad_com%absolute_time_tracking = .false.
     call set_phi0(ring, cavity_ix) ! resets phi0$ = 0 at all cavities 

     write(*,*) "========================================================================"
     write(*,*) "RMS (measured - Bmad reported) DC dispersion:"
     write(*,'(a,es10.3,a,es10.3)') " RMS delta(dc_eta.x)   = ", d_eta_x, "    RMS delta(dc_eta.y)   = ", d_eta_y
     write(*,*) "========================================================================"


   end subroutine sim_dc_eta_data

   ! ===========================================
   ! subroutine locate_rf
   ! returns list of RF cavity indices, and RF frequency in Hz
   subroutine locate_rf(ring, cavity_ix, freq)
     implicit none
     type(lat_struct), target :: ring
     type(ele_struct), pointer :: ele
     integer, allocatable :: cavity_ix(:)
     integer cav_tot, ix
     real(rp), intent(out) :: freq
     logical err

     ! determine RF frequency
     cav_tot = 0
     do ix = 1, ring%n_ele_track
        ele => ring%ele(ix)
        if (ele%key .eq. rfcavity$) then
           freq = value_of_attribute(ele, 'RF_FREQUENCY', err)
           cav_tot = cav_tot + 1
        endif
     enddo
     
     ! generate list of indices for RF cavities:
     allocate(cavity_ix(1:cav_tot))
     cav_tot = 0
     do ix = 1, ring%n_ele_track
        ele => ring%ele(ix)
        if (ele%key .eq. rfcavity$) then
           cav_tot = cav_tot + 1
           cavity_ix(cav_tot) = ix
        endif
     enddo
     
   end subroutine locate_rf

   ! ===================================================
   subroutine de_to_df(ring, rf_freq, delta, dfreq)
     implicit none
     type(lat_struct) ring
     type(coord_struct), allocatable :: co(:)
     type(normal_modes_struct) modes
     real(rp) :: alpha_p, gamma, dfreq, rf_freq, delta
     integer :: rad_cache = 0
     logical err

     call closed_orbit_calc(ring, co, 6)
     call twiss_at_start(ring) 
     call twiss_propagate_all(ring)
     call radiation_integrals(ring, co, modes, rad_cache) ! required for determining alpha_p
     call release_rad_int_cache(rad_cache)
     
     ! compute frequency change corresponding to requested energy change
     alpha_p = modes%synch_int(1)/ring%param%total_length
     gamma = value_of_attribute(ring%ele(0),'E_TOT',err)/(m_electron)
     dfreq = abs(rf_freq * (1./(1. + alpha_p * delta) - 1.)) ! don't care about sign, only magnitude

   end subroutine de_to_df


   ! ===========================================
   ! simple RMS computation for horizontal orbit
   subroutine rms_x(co, delta_x)
    implicit none
    
    real(rp), intent(out) :: delta_x
    type(coord_struct), allocatable :: co(:)
    integer ix

    delta_x = 0. ! reset on each call
    
    do ix = 0, size(co)
       delta_x = delta_x + (co(ix)%vec(1))**2
    enddo
    delta_x = sqrt(delta_x / size(co)) 
  end subroutine rms_x
  
  
  ! ==========================================
  ! change rf frequency:
  subroutine rf_fm(ring, cavity_ix, dfreq, method)
    implicit none 
    
    type(lat_struct) :: ring
    real(rp) :: dfreq
    type(ele_struct), pointer :: ele
    type(all_pointer_struct) rf_freq_ptr
    integer ix
    integer, allocatable :: cavity_ix(:)
    character(6), optional :: method ! either 'target' or 'delta'
    character(6) :: method_used = 'delta'
    logical err
    
    if (present(method)) method_used = method
    
    do ix = 1, size(cavity_ix)
       ele => ring%ele(cavity_ix(ix))
       call pointer_to_attribute(ele, 'RF_FREQUENCY', .true., rf_freq_ptr, err)
       if (trim(method_used) .eq. 'target') then 
          rf_freq_ptr%r =  dfreq
       else ! either 'delta' or not recognized; default to 'delta'
          rf_freq_ptr%r =  rf_freq_ptr%r + dfreq
       endif
       call set_flags_for_changed_attribute(ele, rf_freq_ptr)
    enddo
    
    call lattice_bookkeeper(ring)
    
  end subroutine rf_fm

  ! ==========================================
  subroutine rffm_to_center(ring, rf_freq, cavity_ix)
    implicit none
    type(lat_struct), target :: ring
    type(ele_struct), pointer :: ele
    type(coord_struct), allocatable :: co(:)

    real(rp) :: rf_freq

    integer, allocatable :: cavity_ix(:)
    integer :: ix, max_ix, sum_tot

    real(rp) :: max_x, de, dfreq
    real(rp) :: de_crit = 5.e-4 ! maximum allowed energy change

    max_x = 0.
    max_ix = 0

    ! we know the closed orbit, which is on a dispersive trajectory. 
    ! therefore, we can determine the energy change (and hence the 
    ! frequency change) required to set the orbit on-axis. 
    call closed_orbit_calc(ring, co, 6)
    call twiss_at_start(ring) 
    call twiss_propagate_all(ring)

    ! wigglers cause problems, so only check at quads
    de = 0.
    sum_tot = 0
    do ix = 1, ring%n_ele_track
       if ((ring%ele(ix)%key .eq. quadrupole$) .and. &
            (abs(co(ix)%vec(1) / ring%ele(ix)%x%eta) .lt. de_crit)) then
          write(23,'(i8,2es18.8)') ix, co(ix)%vec(1), ring%ele(ix)%x%eta
          de = de + co(ix)%vec(1) / ring%ele(ix)%x%eta
          sum_tot = sum_tot + 1
       endif
    enddo
    de = de / sum_tot
    
    
    ! translate to a frequency change:
    call de_to_df(ring, rf_freq, de, dfreq)
    if (de .gt. 0) dfreq = -dfreq
    rf_freq = rf_freq - dfreq

    write(*,'(a,es14.7,a,f0.3,a)') "Changing energy by delta = ", de, " (dfreq = ", -dfreq, "Hz)"
    
    ! set RF:
    call rf_fm(ring, cavity_ix, rf_freq, 'target')

  end subroutine rffm_to_center
  
  ! ==========================================
  ! set rf cavity phases so bunch passage corresponds to phi = 0

  subroutine set_phi0(ring, cavity_ix)
    implicit none 
    
    type(lat_struct) ring
    type(ele_struct), pointer :: ele
    type(all_pointer_struct) :: a_ptr
    type(coord_struct), allocatable :: co(:), orb(:)
    integer, allocatable :: cavity_ix(:) ! list of indices for rf cavities
    integer ix, ix_0, ix_1, jx
    real(rp) phi0
    logical err

    if (.not. allocated(co)) allocate(co(0:ring%n_ele_track), orb(0:ring%n_ele_track))
    !co(:)%t = 0.
    
    do ix = 1, ubound(cavity_ix,1)
       
       ix_0 = 0
       ix_1 = cavity_ix(ix)
       ele => ring%ele(ix_1)

       if (bmad_com%absolute_time_tracking .eqv. .true.) then
          ! note: bmad evaluates at the END of elements by default
          call track_many(ring, orb, ix_0, ix_1, 1)
          phi0 = mod(value_of_attribute(ele,'RF_FREQUENCY',err)*orb(ix_1-1)%t, 1.)
          if (phi0 .gt. pi) phi0 = phi0 - twopi
       else ! not using absolute_time_tracking
          phi0 = 0.
       endif

       call pointer_to_attribute(ele, 'PHI0', .true., a_ptr, err)
       a_ptr%r = phi0
       call set_flags_for_changed_attribute(ele, a_ptr)
    enddo ! size(cavity_ix)
    
    call lattice_bookkeeper(ring)
    
  end subroutine set_phi0
   
 end module sim_eta_mod