!+ 
! Program particle_generator
!
! Parses lattice file and outputs particle .dat file with
! particles on lattice walls.
!
! Places particles around regions between RF cells (ie at irises).
! Populates 1/lattice_part of lattice elements that are RF cells,
! each with iris_pts emitters and numPhases of phases
!
! Modules Needed:
!   use bmad
!   use beam_def_struct
!
! Input (namelist)
!   particle_generator.in
!
! Output
!   i#_e#_p#.dat or i#_e#_p#_all.dat
!-

program particle_generator

use bmad
use beam_def_struct

implicit none

type (lat_struct) :: lat
type (ele_struct), pointer :: ele
type (coord_struct) :: particle
character(100) :: in_file, lat_name, particle_file_name, &
     numIris_str, iris_pts_str, numPhases_str, &
     numParticles_str
integer :: ix_ele, numIris, numParticles, &
     phase_ct, numPhases, &
     section_id, lattice_part, lattice_shift, &
     iris_pts, iris_pts_max, iris_ct, &
     top_pts, top_ct
real (rp) :: tiny, d_radius, r, phi, beta_fn, A_fn
integer, parameter :: particle_file = 1, namelist_file = 2

namelist / particle_generator_params / lat_name, numPhases, &
     lattice_part, lattice_shift, iris_pts, iris_pts_max, tiny, beta_fn, A_fn

!------------------------------------------
!Defaults for namelist
lat_name = 'la.bmad'
particle_file_name = "EmitterParticles.dat"
numPhases = 1
lattice_part = 30
lattice_shift = 0
iris_pts = 1
iris_pts_max = 10
tiny = 0.001
beta_fn = 300
A_fn = 3e-17

!Read namelist
in_file = 'particle_generator.in'
print *, 'Opening: ', trim(in_file)
open (namelist_file, file = in_file, status = "old")
read (namelist_file, nml = particle_generator_params)
close (namelist_file)

!Parse lattice
call bmad_parser (lat_name, lat)

!Initialize pseudorandom number generator
call init_random_seed()

!-----------------------
!Count number of particles
numParticles = 0
numIris = 0
do ix_ele = 1 + lattice_shift, lat%n_ele_track / lattice_part + lattice_shift
   if (lat%ele(ix_ele)%key .eq. lcavity$) then
      numParticles = numParticles + numPhases * (iris_pts + 1)
      numIris = numIris + 1
   end if
end do


!Prepare output file name: i#_e#_p#_all.dat, with _all present if we place
! particles all along inside wall (ie. iris_pts = iris_pts_max)
write(numIris_str, *) numIris
numIris_str = adjustl(numIris_str)
write(iris_pts_str, *) iris_pts + 1
iris_pts_str = adjustl(iris_pts_str)
write(numPhases_str, *) numPhases
numPhases_str = adjustl(numPhases_str)
if (iris_pts .eq. iris_pts_max) then
   particle_file_name = "i" // trim(numIris_str) // "_e" // &
        trim(iris_pts_str) // "_p" // trim(numPhases_str) // "_all.dat"
else
   particle_file_name = "i" // trim(numIris_str) // "_e" // &
        trim(iris_pts_str) // "_p" // trim(numPhases_str) // ".dat"
end if


!Open output file
open(particle_file, file = particle_file_name)


!Write number of particles to file
write(numParticles_str, *) numParticles
write (particle_file, '(a)') trim(adjustl(numParticles_str))
print *, "Number of particles generated: ", numParticles


!Make particles and write to file
!First set invariant coordinates (momenta)
particle%vec(2) = 0.0
particle%vec(4) = 0.0
particle%vec(6) = 0.0


!Set rest of coordinates; assume pillbox cavities
do ix_ele = 1 + lattice_shift, lat%n_ele_track / lattice_part + lattice_shift

   ele => lat%ele(ix_ele)

   if (ele%key .eq. lcavity$) then

      d_radius = ele%wall3d%section(2)%v(1)%radius_x - &
           ele%wall3d%section(1)%v(1)%radius_x

      particle%vec(5) = ele%s - ele%value(l$) + tiny**3 
      r = ele%wall3d%section(1)%v(1)%radius_x - tiny**3

      !Make particles along first edge of cell
      do iris_ct = 1, iris_pts

         !Place particle at random angle phi; random_number has range {0,1}
         call random_number(phi)
         phi = phi * pi/2
!         phi = 0.0
         particle%vec(1) = cos(phi) * r 
         particle%vec(3) = sin(phi) * r 

         do phase_ct = 1, numPhases
            particle%t = (real(phase_ct) / numPhases) / (ele%value(rf_frequency$) * ele%em_field%mode(1)%harmonic)

            !Determine weight (charge)
            call Fowler_Nordheim_current_calc(lat, ele, particle, numPhases, &
                 A_fn, beta_fn)

            !Write to file: starting coordinates
!            if (.not. (particle%charge .eq. 0.0)) &
                 write (particle_file, '(8es18.10)') particle%vec, &
                 particle%t, particle%charge

         end do

         r = r + d_radius / iris_pts_max
         particle%vec(5) = particle%vec(5) + tiny / iris_pts_max

      end do

      !Now r should be maximum; make particles along top of pillbox cavity
      if (iris_pts .eq. iris_pts_max) then

         !Calculate number of points on top for same point density as sides
         top_pts = (ele%value(l$) - 2*tiny) * (iris_pts / d_radius)

         do top_ct = 1, top_pts

            !Place particle at random angle phi; random_number has range {0,1}
            call random_number(phi)
            phi = phi * pi/2
!            phi = 0.0
            particle%vec(1) = cos(phi) * r 
            particle%vec(3) = sin(phi) * r 

            do phase_ct = 1, numPhases
               particle%t = (real(phase_ct) / numPhases) / (ele%value(rf_frequency$) * ele%em_field%mode(1)%harmonic)

               !Determine weight (charge)
               call Fowler_Nordheim_current_calc(lat, ele, particle, &
                    numPhases, A_fn, beta_fn)

!               if (.not. (particle%charge .eq. 0.0)) &
                    write(particle_file, '(8es18.10)') particle%vec, &
                    particle%t, particle%charge
            end do

            particle%vec(5) = particle%vec(5) + d_radius / iris_pts_max
         end do

      end if


      !Now make particles along opposite edge of cavity
      particle%vec(5) = ele%s - tiny - tiny**3

      do iris_ct = 1, iris_pts

         !Place particle at random angle phi; random_number has range {0,1}
         call random_number(phi)
         phi = phi * pi/2
!         phi = 0.0
         particle%vec(1) = cos(phi) * r 
         particle%vec(3) = sin(phi) * r 

         do phase_ct = 1, numPhases
            particle%t = (real(phase_ct) / numPhases) / (ele%value(rf_frequency$) * ele%em_field%mode(1)%harmonic)

            !Determine weight (charge)
            call Fowler_Nordheim_current_calc(lat, ele, particle, numPhases, A_fn, beta_fn)

            !Write to file: starting coordinates
!            if (.not. (particle%charge .eq. 0.0)) &
                 write (particle_file, '(8es18.10)') particle%vec, &
                 particle%t, particle%charge
         end do

         r = r  - d_radius / iris_pts_max
         particle%vec(5) = particle%vec(5) + tiny / iris_pts_max

      end do

   end if

end do


!Close data file
close(particle_file)
print *, "Written: ", particle_file_name



contains

!+
! Subroutine Fowler_Nordheim_current_calc
!
! Routine to calculate the field emitted current (in amperes)
! according to the Fowler Nordheim model
!
! Modules needed
!   use bmad
!   use capillary_mod
!   use beam_def_struct
!
! Input
!   lat        -- lat_struct: Accelerator lattice
!   ele        -- ele_struct: current lattice element
!   particle   -- particle_struct: dark current particle
!     %r       -- coord_struct: particle coordinates
!   numPhases  -- integer: number of phases over 2 pi
!   A_fn       -- real(rp): FN effective emitter area
!   beta_fn    -- real(rp): FN field enhancement factor
!
! Output
!   particle   -- coord_struct: dark current particle
!     %charge  -- real(rp): particle weight!
!
!-

subroutine Fowler_Nordheim_current_calc(lat, ele, particle, numPhases, &
     A_fn, beta_fn)

use bmad
use capillary_mod
use beam_def_struct

implicit none

type (lat_struct) :: lat
type (ele_struct) :: ele
type (coord_struct) :: particle
type (photon_coord_struct) :: particle_p
type (em_field_struct) :: field

integer :: numPhases
real(rp) :: A_fn, beta_fn, e_tot, d_radius, norm(3), y, I_fn
!real(rp), parameter :: e = 1.60217646e-19


!------------------
!Calculate E field at point
call em_field_calc (ele, lat%param, &
     particle%vec(5) - (ele%s - ele%value(l$)), particle%t, particle, &
     .false., field, .false.)


!Change coordinates to "photon" coords for capillary subroutines
!Get e_tot from momentum, calculate beta_i = c*p_i / e_tot
particle_p%orb = particle
e_tot = sqrt(particle_p%orb%vec(2)**2 + particle_p%orb%vec(4)**2 + &
     particle_p%orb%vec(6)**2 + mass_of(lat%param%particle)**2)
particle_p%orb%vec(2) = particle_p%orb%vec(2) / e_tot
particle_p%orb%vec(4) = particle_p%orb%vec(4) / e_tot
particle_p%orb%vec(6) = particle_p%orb%vec(6) / e_tot

!Get normal vector to wall
d_radius = capillary_photon_d_radius(particle_p, ele, norm)


!If E field pushes particle into wall, charge is 0
if (field%E(1) * norm(1) + &
     field%E(2) * norm(2) + &
     field%E(3) * norm(3) < 0) then
!     sqrt(field%E(1)**2 + field%E(2)**2 + field%E(3)**2) * cos(pi/3)) then
   particle%charge = 0

   return
end if

!Otherwise, use FN emission formula to determine current
!Formula from p.94 of RF Superconductivity by Hassan Padamsee
!Assume a triangular potential barrier
I_fn = 3.85e-7 * A_fn * &
     (beta_fn**2 * (field%E(1)**2 + field%E(2)**2 + field%E(3)**2)) * &
     exp( -5.464e10 / &
     (beta_fn * sqrt((field%E(1)**2 + field%E(2)**2 + field%E(3)**2))))

!Q_fn = I_fn * delta_t
particle%charge = I_fn * (1.0_rp / numPhases) / (ele%value(rf_frequency$) * ele%em_field%mode(1)%harmonic)

!Weight charge by radius
particle%charge = particle%charge * &
     sqrt(particle%vec(1)**2 + particle%vec(3)**2)

end subroutine



!Small subroutine to get us random numbers
! taken from fortranwiki.org/fortran/show/random_seed
subroutine init_random_seed()
  integer :: i, n, clock
  integer, dimension(:), allocatable :: seed

  call random_seed(size = n)
  allocate(seed(n))

  call system_clock(count=clock)

  seed = clock + 37 * (/ (i - 1, i = 1, n) /)
  call random_seed(put = seed)

  deallocate(seed)
end subroutine


end program