!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Module sim_bpm_mod
!
! This module contains subroutines to simulate: position resolution in terms 
! of additive button signal; physical misalignment and gain errors on BPMs; 
! and applying random resolution, gain errors, tilt, shear, etc. to a 
! measurement.
!--------------------------------------------------------------------------
module sim_bpm_mod

  use bmad
  use sim_utils
  use random_mod
  use cesr_basic_mod
  use nonlin_bpm_mod

  implicit none

  type det_data_struct
     real(rp) :: vec(6) ! 6-vector of positions and momenta
     real(rp) :: amp(4) = 0. ! button signals
     integer :: ix_db = 0
     integer :: ix_lat = 0
  end type det_data_struct

 
  type det_struct
     real(rp) :: vec(6) ! 6-vector of positions and momenta
     real(rp) :: amp(4) = 0. ! button signals
     real(rp) :: gain(4) = 1.
     real(rp) :: timing(4) = 0.
     integer :: ix_db = 0
     integer :: ix_lat = 0
     integer :: ix_quad = 0 ! for relative BPM-to-quadrupole centering
     real(rp) :: shear_x = 0.
     real(rp) :: butn_res_sigma = 0.
     real(rp) :: button_offset(4) = 0.
     real(rp) :: tilt = 0.
     real(rp) :: x_offset = 0.
     real(rp) :: y_offset = 0.
     real(rp)    wt, param ! compatability with ring_ma
  end type det_struct
  
  type det_error_struct
     real(rp) :: bpm_offset = 0.
     real(rp) :: bpm_rotation = 0.
     real(rp) :: shear_x = 0.
     real(rp) :: gain_sigma = 0.
     real(rp) :: timing_sigma = 0.
  end type det_error_struct


  type use_bpm_struct
     character(40) :: name = ''
     real(rp) :: s = 0.
     logical :: is_on = .true.
  end type use_bpm_struct
  

!---------------------------------------------

contains

  
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine find_bpms(ring, bpm_mask, bpm)
!
! From lat_struct, locate BPMs of interest. 
!
! Input:
!      lat - lat_struct
! bpm_mask - char(40); regex mask for locating BPMs.
!
! Output:
!   bpm(1:n_bpms) -- det_struct(:); Holds information about BPMs of interest.
!--------------------------------------------------------------------------
 subroutine find_bpms(ring, bpm_mask, bpm)

  use cesr_basic_mod
   
   implicit none
   type(db_struct) cesr
   type(lat_struct) :: ring
   type(ele_struct), pointer :: ele
   type(det_struct), allocatable :: bpm(:)
   integer :: i, ix, j, jx, bpm_index, n_bpms = 0
   character(40) bpm_mask
   real(rp) :: ds1, ds2
   
   do i=1, ring%n_ele_max
      if (match_reg(ring%ele(i)%name, trim(bpm_mask))) then
         n_bpms = n_bpms + 1 
      endif
   enddo
   
   allocate(bpm(1:n_bpms))

   bpm_index = 1
   bpm(:)%ix_db = -1

   do i=1, ring%n_ele_max
      if (match_reg(ring%ele(i)%name, trim(bpm_mask))) then
         bpm(bpm_index)%ix_lat = ring%ele(i)%ix_ele
         bpm_index = bpm_index + 1
      endif
   enddo  

  ! if CESR-type lattice, assign bpm(:)%ix_db indices
!  if (match_reg(ring%use_name, "CESR")) then
!     call bmad_to_cesr(ring, cesr)
!     jx = 1
!     do i = 1, ubound(cesr%detector,1)
!        if (match_reg(cesr%detector(i)%bmad_name, 'DUMMY')  .or. match_reg(cesr%detector(i)%bmad_name,'[aA]')) cycle
!        bpm(jx)%ix_db = cesr%detector(i)%ix_db
!        jx = jx + 1
!        if (jx .gt. n_bpms) exit
!     enddo
!  else ! non-CESR-type lattice
     do ix = 1, n_bpms
        bpm(ix)%ix_db = ix
     enddo
!  endif
 
  !-----------------------------------------------------------
  ! For BPM offsets: must be relative to nearest quadrupole.
  ! Search for element index of nearest quadrupole, and apply
  ! BPM offset on top of that quad's existing offset.

  do ix = 1, n_bpms
     ! first search forward:
     do j = bpm(ix)%ix_lat, ring%n_ele_track + bpm(ix)%ix_lat
        if (j .gt. ring%n_ele_track) then
           jx = j - ring%n_ele_track
        else
           jx = j
        endif
        ele => ring%ele(jx)
        if (ele%key .eq. quadrupole$) then
           bpm(ix)%ix_quad = jx
           ds1 = abs(ele%s - ring%ele(bpm(ix)%ix_lat)%s)
           exit
        endif
     enddo

     ! now search backward:
     do j = bpm(ix)%ix_lat, -ring%n_ele_track + bpm(ix)%ix_lat
        if (j .lt. 0) then
           jx = j + ring%n_ele_track
        else
           jx = j
        endif
        ele => ring%ele(jx)
        if (ele%key .eq. quadrupole$) then
           ds2 = abs(ele%s - ring%ele(bpm(ix)%ix_lat)%s)
           if (ds2 .lt. ds1) then
              ! this quad is closer than the nearest quad 
              ! when searching forward; use this quad.
              bpm(ix)%ix_quad = jx
           endif
           exit
        endif
     enddo
  enddo
  

 end subroutine find_bpms
 
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine toggle_bpms(bpm, lat, use_bpm_file)
!
! Toggle BPMs "is_on" in order to mask some BPMs from being used in 
! corrections. If use_bpm_file is provided, all BPMs default to 
! is_on = .false. unless otherwise declared in the input file.
!
! Input:
!   bpm -- det_struct, holds information about a single BPM
!   lat -- lat_struct
!   use_bpm_file -- character(200), holds list of BPMs, their s-locations,
!                   and use_it logical. Does not have to be a complete
!                   list of BPMs.
!
! Output:
!   lat
!      %ele(bpm(:)%ix_lat)%is_on -- toggled is_on
!--------------------------------------------------------------------------
 subroutine toggle_bpms(bpm, lat, use_bpm_file, frac_failed, seed)
   
   implicit none
   
   type(lat_struct) :: lat
   type(ele_struct), pointer :: ele
   type(det_struct), allocatable :: bpm(:)
   type(use_bpm_struct), allocatable :: use_bpm_list(:)
   character(200) :: use_bpm_file
   real(rp) x, y, b1(4), b2(4), d(4,0:2,0:2)
   real(rp) :: butn_res_sigma = 0
   real(rp), optional :: frac_failed
   integer ix, jx, jx_old, eof, ix_lat

   ! for statistical failure of BPMs
   integer, optional :: seed
   integer :: seed_used, num_valid_bpms
   real(rp) :: harvest

   character(40) :: name = ''
   real(rp) :: s = 0., tolerance = 1.e-6
   logical :: is_on = .true.
     
   if ((trim(use_bpm_file) .eq. '') .and. (.not. present(frac_failed))) then
      write(*,*)
      write(*,*) "use_bpm_file or frac_failed  not provided; enable all BPMs within regex match"
      write(*,*)
      return
   endif

   allocate(use_bpm_list(size(bpm)))
   
   ix = 0
   
   ! read in file:
   if ((trim(use_bpm_file) .ne. '')) then
      open(1, file=use_bpm_file)
      read(1,*) ! dispose of comment line
      
      do
         ix = ix + 1
         read(1,*, iostat = eof) name, s, is_on
         if (eof .eq. 0)  then
            use_bpm_list(ix)%name  = name
            use_bpm_list(ix)%s     = s
            use_bpm_list(ix)%is_on = is_on
         elseif (eof .lt. 0) then
            exit
         else
            write(*,*) "Problem reading file ", trim(use_bpm_file), ". Stopping here..."
            stop
         endif
      enddo
      
      close(1)
      
      ! initially mask all BPMs
      do ix = 1, size(bpm)
         ix_lat = bpm(ix)%ix_lat
         lat%ele(ix_lat)%is_on = .false.
      enddo
      
      ! now go through and unmask only what we want to use
      jx_old = 1
      use_bpm: do ix = 1, size(use_bpm_list)
         bpm_list: do jx = jx_old, size(bpm)
            ix_lat = bpm(jx)%ix_lat
            if (abs(lat%ele(ix_lat)%s -  use_bpm_list(ix)%s) .lt. tolerance) then
               lat%ele(ix_lat)%is_on = use_bpm_list(ix)%is_on
               jx_old = jx
               cycle use_bpm
            endif
         enddo bpm_list
      enddo use_bpm
      
   endif
   
   ! mask BPMs which have statistically failed:
   if (present(frac_failed)) then
      num_valid_bpms = count(lat%ele(bpm(:)%ix_lat)%is_on)
      write(*,'(a,i0,a,i0)') "Total # of BPMs: ", size(bpm), "; # valid BPMs: ", num_valid_bpms
      write(*,'(a,i0,a)') "Disabling ", int(num_valid_bpms*frac_failed), " BPMs due to random failures."
      call ran_seed_put(seed)
      do jx = 1, int(num_valid_bpms*frac_failed)
         is_on = .false.
         do while (is_on .eqv. .false.) ! ignore BPMs already disabled
            ! choose a random number from 1-size(bpm)
            call ran_uniform(harvest)
            ix = int(harvest * (size(bpm)-1)+1)
            ix_lat = bpm(ix)%ix_lat
            is_on = lat%ele(ix_lat)%is_on
         enddo
         lat%ele(ix_lat)%is_on = .false. ! if BPM is selected, toggle it
      enddo
   endif
      

!   write(*,*) '----------------------------------------------------' 
!   write(*,'(a10,a50,a10)') "BPM", "s-loc", "is_on"
!   do ix = 1, size(bpm)
!      ix_lat = bpm(ix)%ix_lat
!      ele => lat%ele(ix_lat)
!      write(*,*) ele%name, ele%s, 'is_on = ', ele%is_on
!   enddo
!   write(*,*) '----------------------------------------------------' 
   
 end subroutine toggle_bpms
 
!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine resolution_to_button_error(bpm, current, res)
!
! Determine the desired resolution in terms of an additive error on button 
! signal amplitudes.
!
! Input:
!   bpm -- det_struct, holds information about a single BPM
!   current -- real(rp), in "nonlin_bpm" units
!   res -- real(rp). resolution, in meters
!
! Output:
!   bpm
!      %butn_res_sigma -- resolution in terms of additive button signal error
!--------------------------------------------------------------------------
  subroutine resolution_to_button_error(bpm, current, res)

    implicit none

    type(det_struct) :: bpm
    real(rp), intent(in) :: res, current
    real(rp) x, y, b1(4), b2(4), d(4,0:2,0:2)
    real(rp) :: butn_res_sigma = 0
    integer i


    call nonlin_bpm_set_pointers(25) !use BPM #25 as an example
    x = 1.e-3 !1mm orbit in x and y
    y = 1.e-3
    call nonlin_bpm_interpolate((/x,y,current/),d)
    b1(1:4) = d(1:4,0,0)
    
    ! Now add 1\sigma to the x and y positions, see what the effect is on the button signal amplitudes
    
    x = 1.e-3 + res
    y = 1.e-3 + res
    call nonlin_bpm_interpolate((/x,y,current/),d) 
    b2(1:4) = d(1:4,0,0)
    
    ! butn_res_sigma initialized to zero already
    
    do i=1,4
       butn_res_sigma = butn_res_sigma + (b2(i) - b1(i))**2/4
    enddo
    butn_res_sigma = sqrt(butn_res_sigma)
    bpm%butn_res_sigma = butn_res_sigma
    
  end subroutine resolution_to_button_error
  
  !------------------------------------------


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine bpm_errors(this_bpm, bpm_error_sigmas)
!
! Subroutine to assign random Gaussian errors to a BPM. All supplied sigmas are real(rp).
!
! Input:
!  this_bpm         -- det_struct: holds information about misalignments specific to that BPM.
!  bpm_error_sigmas -- structure holding all BPM misalignment sigmas
!     %bpm_offset   -- Sigma for random horizontal/vertical offsets (meters)
!     %bpm_rotation -- Sigma for BPM tilts (radians)
!     %shear_x      -- Sigma for a "shear" effect where top and bottom BPM button blocks are 
!                         transversely misaligned (in x). (meters)
!     %gain_sigma   -- Sigma for single-button gain errors. (unitless: 0.01 = 1%)
!     %timing_sigma -- Sigma for timing errors (seconds: 10.0e-12)
!  quad             -- ele_struct, optional: nearest quadrupole to BPM (for offsets)
!  sigma_cutoff_in  -- integer, optional: if present, veto any errors larger than this
!
! Output:
!   this_bpm         -- Now with a unique, Gaussian set of misalignments/errors:
!      %x_offset     -- 
!      %y_offset     -- 
!      %tilt         -- 
!      %shear_x      --
!--------------------------------------------------------------------------
  subroutine bpm_errors(this_bpm, bpm_error_sigmas, quad, sigma_cutoff_in)
        
    implicit none

    real(rp), optional, intent(in) :: sigma_cutoff_in
    type(ele_struct), optional :: quad
    real(rp) :: sigma_cutoff = 3.
    type(det_struct) :: this_bpm
    type(det_error_struct) :: bpm_error_sigmas
    real(rp) :: harvest, quad_x, quad_y
    integer jx
    logical err
    
    if (present(sigma_cutoff_in)) sigma_cutoff = sigma_cutoff_in

    ! apply tilt first
    call ran_gauss(harvest)
    do while (abs(harvest) > sigma_cutoff)
       call ran_gauss(harvest)
    enddo
    this_bpm%tilt = harvest *  bpm_error_sigmas%bpm_rotation

    ! offsets are relative to nearest quadrupole:
    call ran_gauss(harvest)
    do while (abs(harvest) > sigma_cutoff)
       call ran_gauss(harvest)
    enddo
    if (present(quad)) then
       quad_x = value_of_attribute(quad, 'X_OFFSET',err)
    else 
       quad_x = 0.
    endif
    this_bpm%x_offset = quad_x + harvest * bpm_error_sigmas%bpm_offset

    call ran_gauss(harvest)
    do while (abs(harvest) > sigma_cutoff)
       call ran_gauss(harvest)
    enddo
    if (present(quad)) then
       quad_y = value_of_attribute(quad,'Y_OFFSET',err)
    else 
       quad_y = 0.
    endif
    this_bpm%y_offset = quad_y + harvest *  bpm_error_sigmas%bpm_offset

    
    ! shear is in the rotated coordinate system, so this is ok
    call ran_gauss(harvest)
    do while (abs(harvest) > sigma_cutoff)
       call ran_gauss(harvest)
    enddo
    this_bpm%shear_x = harvest *  bpm_error_sigmas%shear_x  
 
    do jx = 1,4 ! loop over buttons at one BPM
       call ran_gauss(harvest)
       do while (abs(harvest) > sigma_cutoff)
          call ran_gauss(harvest)
       enddo
       this_bpm%gain(jx) = 1. + harvest * bpm_error_sigmas%gain_sigma

       call ran_gauss(harvest)  
       do while (abs(harvest) > sigma_cutoff)
          call ran_gauss(harvest)
       enddo
       this_bpm%timing(jx) = harvest * bpm_error_sigmas%timing_sigma
    enddo
    
    
  end subroutine bpm_errors

  !------------------------------------------

!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine rotate_meas(x,y,angle)
!
! Subroutine to rotate a measurement.
!
! Input:
!   x,y   -- Horizontal and vertical position
!   angle -- angle of rotation
!
! Output:
!   x,y   -- Rotated measurements
!--------------------------------------------------------------------------
  subroutine rotate_meas(x, y, angle)
    implicit none
    real(rp), intent(inout) :: x, y
    real(rp), intent(in) :: angle
    real(rp) :: x0, y0
    
    x0 = x
    y0 = y
    x =  cos(angle)*x0 + sin(angle)*y0
    y = -sin(angle)*x0 + cos(angle)*y0
  end subroutine rotate_meas

  
!---------------------------------------------


!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine apply_bpm_errors(this_bpm, current, lat_type)
!
! Subroutine to assign previously-determined random Gaussian errors to 
! BPM data. 
!
! Input:
!   this_bpm -- det_struct; holds misalignment information
!   current  -- real(rp); in "nonlin_bpm" units
!   lat_type -- character(40) = ring%use_name
!
! Output:
!   this_bpm  -- det_struct, now with applied errors.
!     %vec(1)  -- horizontal orbit, after BPM errors are applied
!     %vec(3)  -- vertical orbit, after BPM errors have been applied
!--------------------------------------------------------------------------
  subroutine apply_bpm_errors(this_bpm, current, lat_type, absolute_time_tracking)
    implicit none
    real(rp) :: harvest
    real(rp), intent(in) :: current
    logical, optional, intent(in) :: absolute_time_tracking
    real(rp) :: x_offset, y_offset
    real(rp) :: g(4), b(4), d(4,0:2,0:2), x2(3),x3(3), &
         x, y, butn_res_sigma, shear_x, temp
    integer bpm_ix,kx
    type(det_struct) :: this_bpm
    character(40) lat_type
    
    real(rp) timing_factor(4)

    butn_res_sigma = this_bpm%butn_res_sigma
    shear_x = this_bpm%shear_x
    g(1:4) = this_bpm%gain(1:4)
    b(1:4) = this_bpm%amp(1:4)

    if (trim(lat_type) .ne. 'CESR') then 
       bpm_ix = 25 ! set to something in the arc, with standard arc map
    else
       bpm_ix = this_bpm%ix_db
    endif

    ! offset BPMs FIRST:
    this_bpm%vec(1) = this_bpm%vec(1) - this_bpm%x_offset
    this_bpm%vec(3) = this_bpm%vec(3) - this_bpm%y_offset

    ! corrected 2013.04.01 after discussion w/ DCS. -JSh
    call rotate_meas(this_bpm%vec(1),this_bpm%vec(3), this_bpm%tilt)

    call bpm_time_to_gain(this_bpm, timing_factor, absolute_time_tracking)

! add shear in 2 steps-- buttons 1 and 2, simulate shifting to the right (therefore, the BEAM
! would shift to the LEFT)
!NOTE: this is not entirely equivalent to a button shift! this is actually shifting the BEAM,
!therefore the image charges will not be the same!
!buttons 3/4, simulate shifting to the left: (ie, beam shifts to the right)

    
    call nonlin_bpm_set_pointers(bpm_ix)

    x2(1)=this_bpm%vec(1)-shear_x
    x2(2)=this_bpm%vec(3)
    x2(3)=current

    call nonlin_bpm_interpolate(x2,d)
    b(1) = g(1)*d(1,0,0)*timing_factor(1)
    b(2) = g(2)*d(2,0,0)*timing_factor(2)

    x2(1)=this_bpm%vec(1)+shear_x
    x2(2)=this_bpm%vec(3)
    x2(3)=current

    call nonlin_bpm_interpolate(x2,d)
    b(3) = g(3)*d(3,0,0)*timing_factor(3)
    b(4) = g(4)*d(4,0,0)*timing_factor(4)
    

    do kx=1,4
       call ran_gauss(harvest)
       b(kx) = b(kx) + harvest*butn_res_sigma
    enddo

    call nonlin_orbit(-1,b,x2) ! disable reading BPM offsets
    call nonlin_orbit(bpm_ix,b,x2)

    this_bpm%vec(1) = x2(1)
    this_bpm%vec(3) = x2(2)
    this_bpm%amp(1:4) = b(1:4)
  
  end subroutine apply_bpm_errors

!++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
! Subroutine bpm_time_to_gain(this_bpm, timing_factor)
!
! Subroutine to convert the timing errors to a gain factor
!
! Input:
!   this_bpm -- det_struct; holds misalignment information
!
! Output:
!   timing_factor -- real(rp); holds the gain factors for 4 buttons
!--------------------------------------------------------------------------
subroutine bpm_time_to_gain(this_bpm, timing_factor, absolute_time_tracking)

    type(det_struct) this_bpm
    real(rp) timing_factor(4)
    integer j
    logical, optional, intent(in) :: absolute_time_tracking
    real(rp) z_offset, t_offset, timing_err(4)
    real(rp) no_err_signal, err_signal

!   quadratic fit parameters from an actual fit of the button pulse
!   ( f(t)=a0*t**2+a1*t+a2 )
    real(rp) :: a0= -14.1676, a1= 7586.69, a2= -993141.0

    z_offset=this_bpm%vec(5)
    t_offset=z_offset / c_light

    if (present(absolute_time_tracking) .and. (absolute_time_tracking .eqv. .true.)) then
       timing_err(:) = this_bpm%timing(:) ! neglect longitudinal bunch position if using absolute_time_tracking
    else
       timing_err(:)=this_bpm%timing(:)+t_offset
    endif

    no_err_signal= a2-a1**2/a0/4

    do j=1,4
       err_signal=a0*(timing_err(j)/10.0e-12)**2+a2-a1**2/a0/4
       timing_factor(j)=err_signal/no_err_signal
    enddo

end subroutine bpm_time_to_gain


end module sim_bpm_mod