module und_energy_trans_mod

use bmad

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

contains

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

!subroutine und_plane_kick (lat, closed_orb, source_ele, kick_ele, lens_ele, delta_s, delta_z, light_amplification_ratio, energy_kick)
!subroutine und_plane_kick (lat, closed_orb, source_ele, kick_ele, delta_s, delta_z, light_amplification_ratio, energy_kick)
subroutine und_plane_kick (lat, source_start, kick_start, source_ele, kick_ele, delta_s, delta_z, light_amplification_ratio, energy_kick)

use bmad

type (lat_struct), target :: lat
!type (coord_struct), allocatable, target :: closed_orb(:)
type (coord_struct) source_start, kick_start
type (coord_struct) start_orb, end_orb
type (ele_struct), pointer, intent(in) :: source_ele, kick_ele
type (ele_struct), pointer :: ele1, ele2
type (track_struct) track

type (lens_grid_struct), target :: grid
type (lens_grid1_struct), pointer :: g1
type (lens_grid1_struct) g1_max

real(rp), intent(in) :: delta_s
real(rp), intent(in) :: delta_z
real(rp), intent(in) :: light_amplification_ratio
real(rp), intent(out) :: energy_kick

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, und_period, 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(:)
real(rp) lens_s

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, time_delay
real(rp) t1_photon, delta_t, delta_z0

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


use_helical_und = 0

n_track_pts = 14000
grid_n_r = 100
grid_n_theta = 100
grid_r_max = 0.01737

!
lens_s = (abs(kick_ele%floor%r(3) - source_ele%floor%r(3)) - source_ele%value(l$)/2 - kick_ele%value(l$)/2 ) / 2
!print *, lens_s

time_delay = delta_s/c_light !
delta_z0 = 0 !2.15E-7 !-3.5e-08 ! small correction for the maximum energy transfer
time_delay = time_delay + ( delta_z + delta_z0) / c_light

!call find_element_ends(source_ele, ele1, ele2)
start_orb = source_start !closed_orb(ele1%ix_ele)

gamma = (1 + start_orb%vec(6)) * start_orb%p0c / (mass_of(start_orb%species) * start_orb%beta)
und_period = source_ele%value(l$) / source_ele%value(n_period$)
k_wig = charge_of(lat%param%particle) * source_ele%value(b_max$) * und_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 * und_period / & 
                          (twopi * mass_of(lat%param%particle) / c_light)
endif

k_u = twopi / und_period

! average longitudinal velocity - see eqtn. (2.11) of https://link.springer.com/chapter/10.1007/978-3-319-04081-3_2
v_bar = c_light * (1 - (1 + k_wig**2/2) / (2 * gamma**2))
! For helical undulator use (1 + k_wig**2) factor
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)

!write (6,'(3(a,es12.4))') 'k_wig: ', k_wig, '  gamma: ', gamma, '  lambda: ', twopi*c_light/omega0 

l_tmp=source_ele%value(l$)/2 + lens_s
theta = atan(grid_r_max /(source_ele%value(l$)/2 + lens_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)

! number of divisions 0.5*omega to 1.5*omega, will be used for the spectal integral
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
enddo

! x_ave is needed to have the center of the particle oscillations and center of the lens on the same level
!x_ave = sum(track%pt(0:track%n_pt)%orb%vec(1)) / (track%n_pt + 1)
! focal length of our lens
foc_dist = (lens_s + source_ele%value(l$) / 2.) /2.
e_transfer = 0

x_ave = -2.164226497269471E-004

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_s

    ! Field in time domain

    do i = 0, track%n_pt
      ! r() is the distance from track i-th step to the lens (it, ir) grid point
      r(1) = g1%x - (track%pt(i)%orb%vec(1) - x_ave)
      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 when the light arrives to the lens (it, ir) grid point
      time(i) = track%pt(i)%orb%t + dr / c_light

      ! time when the light from start of the undulator hits the center of the lens
      if ( i == 0 .and. ir == 0 ) then
         t0 = r(3) / c_light + track%pt(i)%orb%t
!         print *, it, track%pt(i)%orb%t, t0
      end if

      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)

!     Lienard-Wiechert formula See Eq. 37 in News letter
      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_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(track, omega_f, time, t0, Ex, Ey, 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
  enddo
enddo

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

! get field at observation point

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_photon = t0 + lens_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

!  delta_t = t1 - t1_photon - time_delay
!  if (i == 0)   write (6,*) i, t1, t1 - t1_photon

  ! get field at observation point

  do ir = 0, grid_n_r
    ! distance from center of lens
    radius = ir * grid_r_max / grid_n_r
    theta = atan(radius/(kick_ele%value(l$)/2 + lens_s))

    ! observed frequency of radiation by lab-frame observor somewhere on the lens - see derivation of
    ! eqtn 2.19 in https://link.springer.com/chapter/10.1007/978-3-319-04081-3_2, but use the average
    ! longitudinal velocity of an electron in a helical undulator
    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
      ! azimuthal angle where we strike lens
      phi = it * twopi / grid_n_theta

      ! get distance from point on lens to final observation point
      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_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
        ! https://www.classe.cornell.edu/~dlr/osc/documentation/references/OSC_for_ICFA_NewsLet.pdf -
        ! note that division by i is applied by using extra pi/2 in phase advance
        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+time_delay))-1.*omega_f/c_light*dist-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 '(1i16, 6es17.6)', i, track%pt(i)%orb%s-kick_ele%s_start, real(field_x_obs), real(field_y_obs), e_transfer
enddo

energy_kick = e_transfer*c_light

end subroutine und_plane_kick

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

subroutine field_at_lens (track, omega, time, time_0, Ex, Ey, Ex_fft_lens, Ey_fft_lens)

  type (track_struct), target :: track

  real(rp), allocatable, target :: time(:)
  real(rp), allocatable, target :: Ex(:), Ey(:)
  real(rp), target :: omega, time_0
  real(rp) factor

  complex(rp) Ex_fft_lens, Ey_fft_lens

  integer i

  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


!---------------------------------------------
end module und_energy_trans_mod