! JSh, 2012.06.11
! Module for applying a Fourier spectrum of noise to magnets
! Note: fft_data_struct%a are maxima, not RMS.

module ele_noise_mod

  use bmad
  use bookkeeper_mod ! for intelligent bookkeeping
  
  implicit none

  ! n_ele_noise_max = max. # of unique kinds of noise (k1$, k2$, y_offset$, ...), 
  ! with unique spectra for each one. May need to update this limit later.
  integer, parameter :: n_ele_noise_max = 1
  integer, parameter :: fft_max_size = 1
  
  type fft_data_struct
     real(rp) :: a   = 0.
     real(rp) :: phi = 0.
  end type fft_data_struct

  type ele_noise_struct
     character(40) :: mask = ''  ! Select magnets whose name matches the mask
     character(40) :: key_name = ''
     real(rp) :: s_min = 0.0 ! ignore magents inside [s_min, s_max]
     real(rp) :: s_max = 0.0 ! s_max = s_min will ignore the s options
     logical :: inverse = .false. ! if true, select magnets inside [s_min, s_max] 
     integer :: key = 0
     character(40) :: attrib_name = '' ! use g_err for bend field change
     character(40) :: depend_attrib_name = '' ! Option for the g_err changes depend on 'g'
     !integer :: param = 0
     character(40) :: error_type = 'additive' ! must be additive or multiplicative or custom
     integer :: n_turns_recalc = 0 ! how often to update magnets
     type(fft_data_struct) :: fft_data(1:fft_max_size)
     real(rp) :: start_freq = 0. ! initial frequency
     real(rp) :: dfreq = 0. ! step size between frequencies. a(1) is at frequency = dfreq
     logical :: random_phase = .false. ! if .false., all elements have phase = 0.
     logical :: initialized = .false.
     logical :: random_jitter = .false.
     real(rp) :: random_jitter_sigma = 0.
     logical :: individual_ele_noise = .true.
     real(rp) :: const_amp = 1E-5  ! used in "custom" error type
  end type ele_noise_struct
  
    
contains
  
  ! subroutine init_ele_noise(ring, ele_noise)
  ! NOTE: this assumes ran_seed_put has already been called!
  subroutine init_ele_noise(ring, ele_noise)
    use precision_def
    use random_mod
    
    implicit none 
    
    type(lat_struct), target :: ring
    type(ele_struct), pointer :: ele
    type(ele_noise_struct) :: ele_noise(:)
    type(all_pointer_struct) a_ptr
    integer :: ix, jx, kx,  n_ele_match = 0
    real(rp) :: harvest
    logical err, match

    ! check if already initialized
    if (all(ele_noise(:)%initialized) .eqv. .true.) then
       write(*,*) 'init_ele_noise: ele_noise already initialized'
       return
    endif

    call set_custom_attribute_name('ELE_PHASE', err, 1)
    
    ! populate ele_noise(:)%key:
    do ix = 1, size(ele_noise)
       if (len(trim(ele_noise(ix)%attrib_name)) .eq. 0) cycle
       call upcase_string(ele_noise(ix)%mask)
       call upcase_string(ele_noise(ix)%attrib_name)
       call upcase_string(ele_noise(ix)%key_name)
       call downcase_string(ele_noise(ix)%error_type)
       ele_noise(ix)%key = key_name_to_key_index(trim(ele_noise(ix)%key_name), .true.)
    enddo
    
    ! initialize all ele_phase$ to zero
    do ix = 1, ring%n_ele_max
       ele => ring%ele(ix)

       match = .false.
       do jx = 1, size(ele_noise)
          if ( ((ele%key .eq. ele_noise(jx)%key) .and. .not.(ele_noise(jx)%key .eq. -1))) match = .true.
       end do
       if (match) then
          call pointer_to_attribute(ele, 'ELE_PHASE', .true., a_ptr,  err)
          a_ptr%r = 0.0_rp
          call attribute_set_bookkeeping(ele, 'ELE_PHASE', err, a_ptr)
       end if
    enddo
    call lattice_bookkeeper(ring) ! probably unnecessary for ele_phase$, but good to be consistent. 

    if (all(len_trim(ele_noise(:)%attrib_name) .eq. 0)) then
       write(*,*) "All ele_noise(:)%attrib_name are zero. Skipping ele_noise_mod routines."
       return
    endif

    ele_noise_loop: do jx = 1, size(ele_noise)

       if (len_trim(ele_noise(jx)%attrib_name) .eq. 0) cycle ele_noise_loop

       write(*,*) ''
       write(*,*) ''
       write(*,'(6a)') 'key_name = ', trim(ele_noise(jx)%key_name), ' | attrib_name = ', trim(ele_noise(jx)%attrib_name), & 
            ' | mask = ', trim(ele_noise(jx)%mask)

       ! now cycle through elements to find all matches for print-out; apply random phase at each element if applicable:       
       n_ele_match = 0   
       ring_ele_loop: do ix = 1, ring%n_ele_max

          ele => ring%ele(ix)
          
          ! conditionals: if a slave, or not free to vary, or not the element key we're 
          ! looking for, cycle the loop
          if (ele_noise(jx)%inverse .and. ele_noise(jx)%s_min .ne. ele_noise(jx)%s_max) then
             if ((ele%s .lt. ele_noise(jx)%s_min) .or. (ele%s .gt. ele_noise(jx)%s_max)) cycle ring_ele_loop
          else
             if ((ele%s .gt. ele_noise(jx)%s_min) .and. (ele%s .lt. ele_noise(jx)%s_max)) cycle ring_ele_loop
          end if

          if (ele_noise(jx)%mask .ne. "") then
             if (ele_noise(jx)%inverse) then
                if (match_reg(ele%name, ele_noise(jx)%mask)) cycle ring_ele_loop
             else
                if (.not. match_reg(ele%name, ele_noise(jx)%mask)) cycle ring_ele_loop
             end if
          end if
          
          if (.not. ((ele%key .eq. ele_noise(jx)%key) .or. (ele_noise(jx)%key .eq. -1))) cycle ring_ele_loop
          if (ele%slave_status .eq. super_slave$) cycle ring_ele_loop
          if (.not. (has_orientation_attributes(ele) .or. & 
               attribute_free(ele, ele_noise(jx)%attrib_name, err_print_flag = .false.))) cycle ring_ele_loop

          ! now we have only the elements we want 
          write(*,'(a,x,i4,a8,f10.3,x,a)') 'Found matching ele: ', ix, trim(ele%name), ele%s, ele_noise(jx)%mask
          n_ele_match = n_ele_match+1

          ! apply the random phase if applicable:
          if(ele_noise(jx)%random_phase .eqv. .true.) then
             call ran_uniform(harvest) ! generates flat distribution on [0,1]
             call pointer_to_attribute(ele, 'ELE_PHASE', .true., a_ptr,  err)
             a_ptr%r = harvest * 2.*pi
             call attribute_set_bookkeeping(ele,'ELE_PHASE', err, a_ptr)
          endif
          
       enddo  ring_ele_loop ! ix=1, ring%n_ele_max

       write(*,*) ''
       write(*,*) ''

       if (n_ele_match .eq. 0) then
          write(*,*) 'init_ele_noise: No elements matched! Check attrib_name. Stopping here...'
          call err_exit()
       endif

    enddo ele_noise_loop ! jx=1, size(ele_noise)
    call lattice_bookkeeper(ring) ! probably unnecessary for ele_phase$, but good to be consistent. 

    ele_noise(:)%initialized = .true.

  end subroutine init_ele_noise

  !---------------------------------------------------------
  ! subroutine apply_ele_noise(ring, ring_ideal, ele_noise, turns)
  ! subroutine to apply periodic noise to a lattice
  !---------------------------------------------------------

  subroutine apply_ele_noise(ring, ring_ideal, ele_noise, turns)

    use precision_def
    use random_mod
    
    implicit none 

    type(lat_struct), target :: ring
    type(lat_struct) :: ring_ideal
    type(ele_struct), pointer :: ele, ele_ideal
    type(ele_noise_struct) :: ele_noise(:)
    type (all_pointer_struct) :: a_ptr, a_ptr_ideal
    integer, intent(in) :: turns
    integer :: ix, jx, kx
    real(rp) :: net_kick, harvest
    real(rp) :: circ_time = 0., dfreq = 0.
    logical :: err = .false., recalc = .false.
   

    if (all(len_trim(ele_noise(:)%attrib_name) .eq. 0)) return

    ! determine if any noise spectra are ready to be updated; 
    ! if not, just return to main program
    if (any(ele_noise(:)%n_turns_recalc .ne. 0)) then
       do ix = 1, size(ele_noise)
          if (ele_noise(ix)%n_turns_recalc .eq. 0) cycle
          if (mod(turns, ele_noise(ix)%n_turns_recalc) .eq. 0) recalc = .true.
       enddo
    endif

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

!    if (recalc .eqv. .false.) return ! go back to main program if no spectra need updating
    

    ele_noise_loop: do jx = 1, size(ele_noise)
       
!       write (6, *), 'test size=', size(ele_noise)
       ! conditionals: if param_name is null, or if not ready to update the kick from 
       ! that spectrum, cycle.
       if (len(trim(ele_noise(jx)%attrib_name)) .eq. 0) cycle ele_noise_loop
!       if (mod(turns, ele_noise(jx)%n_turns_recalc) .ne. 0) cycle ele_noise_loop

       !-------------------------------
       if ( ele_noise(jx)%individual_ele_noise .eqv. .false. ) call ran_gauss(harvest)

       ! now cycle through elements to apply the updated kicks      
       ring_ele_loop: do ix = 1, ring%n_ele_max

          
          ele => ring%ele(ix)
          ele_ideal => ring_ideal%ele(ix)
          
          ! conditionals: if a slave, or not free to vary, or not the element key we're 
          ! looking for, cycle the loop
          if (ele_noise(jx)%inverse .and. ele_noise(jx)%s_min .ne. ele_noise(jx)%s_max) then
             if ((ele%s .lt. ele_noise(jx)%s_min) .or. (ele%s .gt. ele_noise(jx)%s_max)) cycle ring_ele_loop
          else
             if ((ele%s .gt. ele_noise(jx)%s_min) .and. (ele%s .lt. ele_noise(jx)%s_max)) cycle ring_ele_loop
          end if
          
          if (ele_noise(jx)%mask .ne. "") then
             if (ele_noise(jx)%inverse) then
                if (match_reg(ele%name, ele_noise(jx)%mask)) cycle ring_ele_loop
             else
                if (.not. match_reg(ele%name, ele_noise(jx)%mask)) cycle ring_ele_loop
             end if
          end if
          
          if (.not. ((ele%key .eq. ele_noise(jx)%key) .or. (ele_noise(jx)%key .eq. -1))) cycle ring_ele_loop
          if (ele%slave_status .eq. super_slave$) cycle ring_ele_loop
          if (.not. (has_orientation_attributes(ele) .or. & 
               attribute_free(ele, ele_noise(jx)%attrib_name, err_print_flag = .false.))) cycle ring_ele_loop

          ! now we have only the elements we want. apply the noise:
          call calc_noise_kick(ele_noise(jx),ring%ele(ix), turns, circ_time, net_kick)

          ! now apply this kick to the element:
          call pointer_to_attribute(ele, ele_noise(jx)%attrib_name, .true., a_ptr, err)
          if (len(trim(ele_noise(jx)%depend_attrib_name)) .eq. 0) then
             call pointer_to_attribute(ele_ideal, ele_noise(jx)%attrib_name, .true., a_ptr_ideal, err)
          else
             call pointer_to_attribute(ele_ideal, ele_noise(jx)%depend_attrib_name, .true., a_ptr_ideal, err)
          end if

          if (err .or. .not. associated(a_ptr%r)) then ! failed to find attribute
             write(*,*) "Applying noise FAILED! Check attrib_name. Stopping here..."
             call err_exit()
          endif

          if (match_reg('additive',ele_noise(jx)%error_type)) then
             a_ptr%r = a_ptr_ideal%r + net_kick
          elseif (match_reg('multiplicative',ele_noise(jx)%error_type)) then
             if (ele_noise(jx)%random_jitter) then
                if (ele_noise(jx)%individual_ele_noise) call ran_gauss(harvest)
                if (len(trim(ele_noise(jx)%depend_attrib_name)) .eq. 0) then
                   a_ptr%r = a_ptr_ideal%r * (1.0 + net_kick + harvest*ele_noise(jx)%random_jitter_sigma)
                else
                   a_ptr%r = a_ptr_ideal%r * (net_kick + harvest*ele_noise(jx)%random_jitter_sigma)
                end if
             else
                if (len(trim(ele_noise(jx)%depend_attrib_name)) .eq. 0) then
                   a_ptr%r = a_ptr_ideal%r * (1.0 + net_kick)
                else
                   a_ptr%r = a_ptr_ideal%r * net_kick
                end if
             end if
          elseif (match_reg('custom',ele_noise(jx)%error_type)) then
             if (len(trim(ele_noise(jx)%depend_attrib_name)) .eq. 0) then
                a_ptr%r = a_ptr_ideal%r * (1.0 + net_kick)
             else
                a_ptr%r = a_ptr_ideal%r * net_kick
             end if
          endif
          
!          write(458,'(2i10,5a,2es20.10)') turns, ix, '   ', trim(ele%name), '   ', &
!                                trim(ele_noise(jx)%attrib_name), '   ', a_ptr%r, ring%ele(ix)%value(43)

          call attribute_set_bookkeeping(ring%ele(ix), ele_noise(jx)%attrib_name, err, a_ptr)
          
       enddo  ring_ele_loop ! ix=1, ring%n_ele_max
    enddo ele_noise_loop ! jx=1, size(ele_noise)

    ! update the lattice bookkeeping:
    call lattice_bookkeeper(ring)

  end subroutine apply_ele_noise
  
  !###################################################
  subroutine calc_noise_kick(ele_noise, ele, turns, circ_time, net_kick)
    
    use precision_def
    
    implicit none 
    
    type(ele_noise_struct) :: ele_noise
    real(rp) :: net_kick
    type(ele_struct) :: ele
    real(rp) :: circ_time, dfreq, freq
    integer :: turns, kx
    logical err
    
    
    ! start with initializing net_kick
    if (match_reg('additive',ele_noise%error_type)) then
       net_kick = 0.
    elseif (match_reg('multiplicative',ele_noise%error_type)) then
       net_kick = 1.
    elseif (match_reg('custom',ele_noise%error_type)) then
       net_kick = 1.
    endif

    dfreq = ele_noise%dfreq

    do kx = 1, size(ele_noise%fft_data)
       if (ele_noise%fft_data(kx)%a .eq. 0) cycle
       freq = ele_noise%start_freq + kx*dfreq
       if (match_reg('additive',ele_noise%error_type)) then
          net_kick = net_kick + ele_noise%fft_data(kx)%a*sin(twopi*freq*(circ_time*&
               turns) + ele_noise%fft_data(kx)%phi + value_of_attribute(ele,'ELE_PHASE',err))
       elseif (match_reg('multiplicative',ele_noise%error_type)) then
          net_kick = net_kick * ele_noise%fft_data(kx)%a*sin(twopi*freq*(circ_time*&
               turns) + ele_noise%fft_data(kx)%phi + value_of_attribute(ele,'ELE_PHASE',err))
       elseif (match_reg('custom', ele_noise%error_type)) then
          net_kick = net_kick * ele_noise%const_amp
       else ! error_type not identified
          write(*,*) "ele_noise%error_type ", trim(ele_noise%error_type), &
               " not valid-- must be 'additive', 'multiplicative',or 'custom'. Stopping here."
          stop
       endif
    enddo
    

  end subroutine calc_noise_kick
  
end module ele_noise_mod