program optical_stochastic_cooling

use bmad
use super_recipes_mod

implicit none

type lens_grid1_struct
  real(rp) x, y, z
  real(rp) omega
  complex(rp) Ex
  complex(rp) Ey
end type

type lens_grid_struct
  type (lens_grid1_struct), allocatable :: pt(:,:)
end type

type (lat_struct) lat
type (ele_struct), pointer :: source_ele, kick_ele, lens_ele, ele1, ele2
type (track_struct) track
type (coord_struct), allocatable :: closed_orb(:)
type (coord_struct) start_orb, end_orb
type (lens_grid_struct), target :: grid
type (lens_grid1_struct), pointer :: g1
type (lens_grid1_struct) g1_max

real(rp) grid_r_max, dr, radius, x_ave, omega_f, t0, t1, dist, ddist, phase_tmp, foc_dist, coef_tmp, b_max_helical
real(rp) z_dist_tmp, e_transfer
real(rp) gamma, len_period, k_wik, k_u, v_bar, omega_u, k_wig, theta, phi, l_tmp, d_omega, omega0, omega_t, y_min
real(rp) r(3), vel(3), acc(3) 
real(rp), allocatable :: Ex(:), Ey(:), time(:)
real(rp), allocatable :: freq_ratio(:)

complex(rp)  exp_integral, field_x_tmp, field_y_tmp, field_x_obs, field_y_obs
complex(rp), allocatable :: Ex_lens_tmp(:), Ey_lens_tmp(:)
complex(rp), allocatable :: Ex_field_3D(:, :, :)
complex(rp), allocatable :: Ey_field_3D(:, :, :)
real(rp) lens_n_index, lens_thickness, phase_delay, light_amplification_ratio

integer grid_n_r, grid_n_theta, n_track_pts, num_freq, use_helical_und, side_from_peak
integer i, ir, it, ifreq

character(40) source_wig_name, kick_wig_name, lens_ele_name
character(200) lat_file, init_file

namelist / params / lat_file, n_track_pts, source_wig_name, kick_wig_name, lens_ele_name, & 
                    grid_n_r, grid_n_theta, grid_r_max, lens_n_index, lens_thickness, side_from_peak, & 
                    phase_delay, light_amplification_ratio, & 
                    use_helical_und, b_max_helical

!&
! Read parameters

init_file = 'osc.init'
if (cesr_iargc() == 1) call cesr_getarg(1, init_file)

open (1, file = init_file)
read (1, nml = params)
close (1)


lens_thickness=lens_thickness/1000.

! Some init

call bmad_parser (lat_file, lat)

source_ele => pointer_to_ele(lat, source_wig_name)
kick_ele => pointer_to_ele(lat, kick_wig_name)
lens_ele => pointer_to_ele (lat, lens_ele_name)

!if (source_ele%tracking_method == bmad_standard$) source_ele%tracking_method = runge_kutta$
!if (kick_ele%tracking_method == bmad_standard$) kick_ele%tracking_method = runge_kutta$

call find_element_ends(source_ele, ele1, ele2)

call closed_orbit_calc (lat, closed_orb, 4)

!do i=source_ele%ix_ele - 1, kick_ele%ix_ele
!  write(6,*) i, closed_orb(i)%t
!end do

!start_orb = closed_orb(source_ele%ix_ele-1)
start_orb = closed_orb(ele1%ix_ele)

gamma = (1 + start_orb%vec(6)) * start_orb%p0c / (mass_of(start_orb%species) * start_orb%beta)
len_period =  source_ele%value(l$) / source_ele%value(n_period$)
k_wig = charge_of(lat%param%particle) * source_ele%value(b_max$) * len_period / & 
                          (twopi * mass_of(lat%param%particle) / c_light)
if (use_helical_und == 1 ) then
  k_wig = charge_of(lat%param%particle) * b_max_helical * len_period / & 
                          (twopi * mass_of(lat%param%particle) / c_light)
endif

k_u = twopi / len_period

! For helical undulator use (1 + k_wig**2) factor

v_bar = c_light * (1 - (1 + k_wig**2/2) / (2 * gamma**2))
if (use_helical_und == 1 )    v_bar = c_light * (1 - (1 + k_wig**2) / (2 * gamma**2))

omega_u = k_u * v_bar
omega0 = omega_u * 2 * gamma**2 / (1 + k_wig**2/2)
if (use_helical_und == 1 ) omega0 = omega_u * 2 * gamma**2 / (1 + k_wig**2)
l_tmp=source_ele%value(l$)/2 + lens_ele%s
theta = atan(grid_r_max /(source_ele%value(l$)/2 + lens_ele%s))
d_omega = omega0 - omega_u * 2 * gamma**2 / (1 + k_wig**2/2 + (theta * gamma)**2)
if (use_helical_und == 1 )  d_omega = omega0 - omega_u * 2 * gamma**2 / (1 + k_wig**2 + (theta * gamma)**2)

num_freq=30
field_x_obs = cmplx(0, 0)
field_y_obs = cmplx(0, 0)
! Track through source wiggler

track%n_pt = -1
track%ds_save = source_ele%value(l$) / n_track_pts
call track1 (start_orb, source_ele, lat%param, end_orb, track)

allocate (Ex(0:track%n_pt), Ey(0:track%n_pt), time(0:track%n_pt))
allocate (grid%pt(0:grid_n_r, 1:grid_n_theta))
allocate (Ex_field_3D(0:grid_n_r, 0:grid_n_theta, 0:num_freq))
allocate (Ex_lens_tmp(0:num_freq), Ey_lens_tmp(0:num_freq))
allocate (Ey_field_3D(0:grid_n_r, 0:grid_n_theta, 0:num_freq))
allocate (freq_ratio(0:num_freq))

do ifreq = 0 , num_freq-1
  freq_ratio(ifreq) = 0.5 + ifreq * 1./num_freq
!  freq_ratio(ifreq) = 0. + ifreq * 2./num_freq
enddo

x_ave = sum(track%pt(0:track%n_pt)%orb%vec(1)) / (track%n_pt + 1)
!x_ave = 0
!t0 = 0
foc_dist = (lens_ele%s + source_ele%value(l$) / 2.) /2.
e_transfer = 0

do ir = 0, grid_n_r
  radius = ir * grid_r_max / grid_n_r
  do it = 1, grid_n_theta
    if (ir == 0 .and. it /= 1) exit
    theta = it * twopi / grid_n_theta
    phi = it * twopi / grid_n_theta

    g1 => grid%pt(ir, it)
    g1%x = radius * cos(theta)
    g1%y = radius * sin(theta)
    g1%z = source_ele%value(l$) + lens_ele%s

    ! Field in time domain

    do i = 0, track%n_pt
      r(1) = g1%x - (track%pt(i)%orb%vec(1) - x_ave)
!      r(1) = g1%x - track%pt(i)%orb%vec(1)
      r(2) = g1%y - track%pt(i)%orb%vec(3)
      r(3) = g1%z - (track%pt(i)%orb%s - source_ele%s_start)
      dr = norm2(r)
      time(i) = track%pt(i)%orb%t + dr / c_light
      if ( i == 0 .and. ir == 0 ) t0 = r(3) / c_light + track%pt(i)%orb%t

      vel(1) = track%pt(i)%orb%vec(2) * track%pt(i)%orb%beta / (1+track%pt(i)%orb%vec(6))
      vel(2) = track%pt(i)%orb%vec(4) * track%pt(i)%orb%beta / (1+track%pt(i)%orb%vec(6))
      vel(3) = sqrt(track%pt(i)%orb%beta**2 - vel(1)**2 - vel(2)**2)

      acc(1) = (vel(2) * track%pt(i)%field%b(3) - vel(3) *track%pt(i)%field%b(2))& 
               * c_light / ((1 + start_orb%vec(6)) * start_orb%p0c)
      acc(2) = (vel(3) * track%pt(i)%field%b(1) - vel(1) *track%pt(i)%field%b(3))& 
               * c_light / ((1 + start_orb%vec(6)) * start_orb%p0c)
      acc(3) = (vel(1) * track%pt(i)%field%b(2) - vel(2) *track%pt(i)%field%b(1))& 
               * c_light / ((1 + start_orb%vec(6)) * start_orb%p0c)

      Ex(i) = (-dot_product(r, acc) * (r(1) - dr * vel(1)) + & 
                dr * acc(1) * (dr - dot_product(r, vel))) / (dr - dot_product(r, vel))**3
      Ex(i) = Ex(i) * e_charge / (fourpi * eps_0_vac)


      Ey(i) = (-dot_product(r, acc) * (r(2) - dr * vel(2)) + & 
                dr * acc(2) * (dr - dot_product(r, vel))) / (dr - dot_product(r, vel))**3
      Ey(i) = Ey(i) * e_charge / (fourpi * eps_0_vac)

    enddo

    theta = atan(radius/(source_ele%value(l$)/2 + lens_ele%s))
    g1%omega = omega_u * 2 * gamma**2 / (1 + k_wig**2/2 + (theta * gamma)**2)
    if (use_helical_und == 1 )   g1%omega = omega_u * 2 * gamma**2 / (1 + k_wig**2 + (theta * gamma)**2)

    g1%Ex = cmplx(0, 0)
    g1%Ey = cmplx(0, 0)
    do ifreq = 0, num_freq - 1
      omega_f = g1%omega * freq_ratio(ifreq);
      call field_at_lens(omega_f, t0, Ex_lens_tmp(ifreq), Ey_lens_tmp(ifreq))

!     get field behind lens
!     phase advance such that light in center is delayed by exact amount that it matches delay generated by 
!     longer path length for light hitting edge of lens, then going to focal point

      ddist = 2.*sqrt(grid_r_max**2+(2.*foc_dist)**2)-2.*sqrt(radius**2+(2.*foc_dist)**2)
      phase_tmp=-1.*omega_f/c_light*ddist

      exp_integral = cmplx(cos(phase_tmp), sin(phase_tmp))

      Ex_field_3D(ir, it, ifreq) = Ex_lens_tmp(ifreq)*exp_integral
      Ey_field_3D(ir, it, ifreq) = Ey_lens_tmp(ifreq)*exp_integral

    enddo

    call fft_field (g1%omega, g1%Ex, g1%Ey)

    y_min = super_brent (g1%omega - d_omega, g1%omega, g1%omega + d_omega, brent_fft, 1d-8, 0.0_rp, g1_max%omega)
    call fft_field (g1_max%omega, g1_max%Ex, g1_max%Ey)

    print '(2i6, 4es12.4, 5x, 4es12.4)', ir, it, radius, phi, g1%omega, & 
             abs(g1%Ex), abs(g1%Ey), g1_max%omega, abs(g1_max%Ex), abs(g1_max%Ey)
  enddo
enddo

call find_element_ends(kick_ele, ele1, ele2)
start_orb = closed_orb(ele1%ix_ele)
track%n_pt = -1
call track1 (start_orb, kick_ele, lat%param, end_orb, track)


do i = 0, track%n_pt
  field_x_obs = cmplx(0, 0)
  field_y_obs = cmplx(0, 0)
  !time when light from the center of the lens hits the center of the 2nd undulator
!  t1 = t0 + lens_ele%s/c_light + (2.*sqrt(grid_r_max**2+(2.*foc_dist)**2)-2.*sqrt((2.*foc_dist)**2))/c_light + (track%pt(i)%orb%s - kick_ele%s_start)/v_bar;
  t1 = track%pt(i)%orb%t

  ! get field at observation point

  do ir = 0, grid_n_r
    radius = ir * grid_r_max / grid_n_r
    theta = atan(radius/(kick_ele%value(l$)/2 + lens_ele%s))
    omega_t = omega_u * 2 * gamma**2 / (1 + k_wig**2/2 + (theta * gamma)**2)
    if (use_helical_und == 1 ) omega_t = omega_u * 2 * gamma**2 / (1 + k_wig**2 + (theta * gamma)**2)

    do it = 1, grid_n_theta
      if (ir == 0 .and. it /= 1) exit
      phi = it * twopi / grid_n_theta

      r(1) = radius * cos(phi) - (track%pt(i)%orb%vec(1) - x_ave)
      r(2) = radius * sin(phi) - track%pt(i)%orb%vec(3)
      r(3) = lens_ele%s + (track%pt(i)%orb%s - kick_ele%s_start)

      dist = norm2(r)

      do ifreq = 0, num_freq-1
        omega_f = omega_t * freq_ratio(ifreq)
!        use Kirchoff formula in form given by Lebedev: equation 41  (OSC_for_ICFA_NewsLet.pdf)
        coef_tmp = 1/(twopi * c_light) * radius*grid_r_max/grid_n_r * twopi/grid_n_theta * omega_f/dist
        phase_tmp = omega_f*(t1-t0)-1.*omega_f/c_light*dist-pi/2.

!        get the correct phase
        phase_tmp = omega_f*(t1-t0)-1.*omega_f/c_light*dist-pi/2. + lens_n_index*lens_thickness/c_light + phase_delay
        if (side_from_peak == -1 ) phase_tmp = phase_tmp - pi/2
        if (side_from_peak == 1 )  phase_tmp = phase_tmp + pi/2

        exp_integral = cmplx(cos(phase_tmp), sin(phase_tmp))
        field_x_tmp = coef_tmp * Ex_field_3D(ir, it, ifreq) * exp_integral
        field_y_tmp = coef_tmp * Ey_field_3D(ir, it, ifreq) * exp_integral

        field_x_obs = field_x_obs + 2 * light_amplification_ratio * field_x_tmp * omega_t * (freq_ratio(1)-freq_ratio(0))
        field_y_obs = field_y_obs + 2 * light_amplification_ratio * field_y_tmp * omega_t * (freq_ratio(1)-freq_ratio(0))
      enddo
    enddo
  enddo
  if (i>0) e_transfer = e_transfer + (track%pt(i)%orb%vec(2)*real(field_x_obs) + track%pt(i)%orb%vec(4)*real(field_y_obs))*(track%pt(i)%orb%t-track%pt(i-1)%orb%t)

  print '(1i6, 6es17.6)', i, track%pt(i)%orb%s-kick_ele%s_start, real(field_x_obs), real(field_y_obs), e_transfer
enddo



!---------------------------------------------------------
contains

subroutine fft_field (omega, Ex_fft, Ey_fft)

real(rp) omega, factor
complex(rp) Ex_fft, Ey_fft

!

Ex_fft = 0
Ey_fft = 0

do i = 0, track%n_pt
  if (i == 0) then
    factor = (time(i+1) - time(i)) / 2
  elseif (i == track%n_pt) then
    factor = (time(i) - time(i-1)) / 2
  else
    factor = (time(i+1) - time(i-1)) / 2
  endif
  Ex_fft = Ex_fft + factor * cmplx(cos(time(i)*omega) * Ex(i), sin(time(i)*omega) * Ex(i))
  Ey_fft = Ey_fft + factor * cmplx(cos(time(i)*omega) * Ey(i), sin(time(i)*omega) * Ey(i))
enddo

Ex_fft = 2 * Ex_fft / (time(track%n_pt) - time(0))
Ey_fft = 2 * Ey_fft / (time(track%n_pt) - time(0))

end subroutine fft_field



subroutine field_at_lens (omega, time_0, Ex_fft_lens, Ey_fft_lens)

  real(rp) omega, factor, time_0
  complex(rp) Ex_fft_lens, Ey_fft_lens

  Ex_fft_lens = 0
  Ey_fft_lens = 0

  do i = 0, track%n_pt
    if (i == 0) then
      factor = (time(i+1) - time(i)) / 2
    elseif (i == track%n_pt) then
      factor = (time(i) - time(i-1)) / 2
    else
      factor = (time(i+1) - time(i-1)) / 2
    endif
    Ex_fft_lens = Ex_fft_lens + factor * cmplx(cos((time(i)-time_0)*omega) * Ex(i), -1.*sin((time(i)-time_0)*omega) * Ex(i))
    Ey_fft_lens = Ey_fft_lens + factor * cmplx(cos((time(i)-time_0)*omega) * Ey(i), -1.*sin((time(i)-time_0)*omega) * Ey(i))
  enddo

  Ex_fft_lens = Ex_fft_lens / (2 * pi)
  Ey_fft_lens = Ey_fft_lens / (2 * pi)

end subroutine field_at_lens


!---------------------------------------------------------
! contains

function brent_fft (omega) result (y_val)

real(rp), intent(in) :: omega
real(rp) :: y_val
complex(rp) Ex_fft, Ey_fft

!

call fft_field (omega, Ex_fft, Ey_fft)
y_val = -sqrt(abs(Ex_fft)**2 + abs(Ey_fft)**2)

end function brent_fft

end program