!+ 
! Program particle_generator_cavity7
!
! Parses lattice file and outputs particle .dat file with
! numSources particles evenly spaced on cavity walls.
!
! Places particles along all cavity walls, at random angles.
! Populates 1/lattice_part of lattice elements that are RF
! cavities, each with numPhases of phases
!
! Modules Needed:
!   use bmad
!
! Input (namelist)
!   particle_generator_cavity7.in
!
! Output
!   l#_n#_p#.dat
!-

program particle_generator_cavity7

use bmad

implicit none

type (lat_struct) :: lat
type (ele_struct), pointer :: ele
type (coord_struct), allocatable :: particle_list(:)
type (ele_pointer_struct), allocatable :: eles(:)

character(100) :: in_file, lat_name, particle_file_name, dct_file_name, &
     numPhases_str, numSources_str, file_body
character(400) :: element_monitor_list
character(400) :: elements

integer :: numPhases, numSources
integer :: i, ix_ele, lattice_part, lattice_shift, ix_particle
integer :: n_particles, total_particles
integer :: ix_pass

real(rp) :: tmp_vec(8)
!Namelist items
real(rp) :: buffer, beta_fn, A_fn, min_charge_ratio, Num_emitters_per_m2
real(rp) 	:: total_charge
real(rp) ::  min_allowed_charge, angle0

logical :: err
logical :: angles_on
logical :: make_dark_current_tracker_input_files

integer :: particle_file, dct_file, scratch_file, namelist_file
integer :: n_ele_use, file_number, n_files, n_ele_in_file, n_ele_chunk, total_ele, n_loc
integer :: n_links



namelist / particle_generator_params / lat_name, numSources, &
     numPhases, n_files, buffer, beta_fn, A_fn, angle0, elements, &
     angles_on, min_charge_ratio, Num_emitters_per_m2, make_dark_current_tracker_input_files,  element_monitor_list

!------------------------------------------
!Defaults for namelist
lat_name = 'lat.bmad'
numSources = 100
numPhases = 1
n_files = 1
buffer = 1e-8
beta_fn = 100
A_fn = 1e-15
Num_emitters_per_m2 = 1000
angles_on = .True.
angle0 = 0  ! initial angle / 2 pi, used if angles_on is false. 
min_charge_ratio = 1e-6
make_dark_current_tracker_input_files = .false.
element_monitor_list = ''
elements = 'lcavity::*'

!Read namelist from command line. Otherwise use default
in_file = 'particle_generator_cavity7.in'
if (command_argument_count() > 0) call get_command_argument(1, in_file)
print '(a100)', ' ___  __  ___ _____ _  ___ _   ___       __ ___ __  _ ___ ___  __ _____ __  ___  '
print '(a100)', '| _,\/  \| _ \_   _| |/ _/| | | __|     / _] __|  \| | __| _ \/  \_   _/__\| _ \ '
print '(a100)', '| v_/ /\ | v / | | | | \__| |_| _| ____| [/\ _|| |   | _|| v / /\ || || \/ | v / '
print '(a100)', '|_| |_||_|_|_\ |_| |_|\__/|___|___|____|\__/___|_|\__|___|_|_\_||_||_| \__/|_|_\ '
print '(a100)', ''
print '(2a)', 'Using input file: ', in_file

namelist_file = lunget()
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)

!Force absolute time tracking
bmad_com%absolute_time_tracking = .true. 

!Initialize pseudorandom number generator
call init_random_seed()

!-----------------------
!Prepare output file tail: _n#_p#.dat
write(numSources_str, *) numSources
numSources_str = adjustl(numSources_str)
write(numPhases_str, *) numPhases
numPhases_str = adjustl(numPhases_str)
file_body =  "_n" //  trim(numSources_str) // "_p" // trim(numPhases_str)

!Initialize counters
total_particles = 0
total_charge = 0

! Use element locator string to pre-select elements
call lat_ele_locator (elements, lat, eles, n_loc, err)

!Count elements to use (ele%key == lcavity$)
total_ele = 0
do i = 1, n_loc
  ele => eles(i)%ele
  ! Check for first pass only
  call multipass_chain (ele, ix_pass, n_links)
  if (ix_pass > 1) cycle
  if (ele%key == lcavity$) total_ele = total_ele + 1
end do
n_ele_in_file = ceiling(real(total_ele)/real(n_files))
print *, 'Approximate number of LCAVITY elements in a file: ', n_ele_in_file

!Make array for single element
allocate( particle_list(1:numSources*numPhases))

!Use a scratch file for particle data
scratch_file = lunget()
open(scratch_file, status = 'scratch')

!Get file units
particle_file = lunget()
dct_file = lunget()

!Loop over elements

!
n_ele_chunk = 0
n_ele_use   = 0
file_number = 0
do i = 1, n_loc

  ele => eles(i)%ele
  
  !Only use lcavities for now
  if (ele%key .ne. lcavity$) cycle
  
  ! Check for first pass only
  call multipass_chain (ele, ix_pass, n_links)
  if (ix_pass > 1) cycle
  
  n_ele_chunk = n_ele_chunk + 1
  n_ele_use = n_ele_use + 1

  call generate_field_emission_particles(ele, particle_list, n_particles)
  
  min_allowed_charge =  min_charge_ratio*maxval(particle_list(1:n_particles)%charge)

  !Write coordinates to file
  do ix_particle = 1, n_particles
    if (particle_list(ix_particle)%charge >= min_allowed_charge) then
     	 write (scratch_file, '(8es19.10E3)') &
           particle_list(ix_particle)%vec, &
           particle_list(ix_particle)%t, &
           particle_list(ix_particle)%charge
           total_particles = total_particles + 1
           total_charge = total_charge + particle_list(ix_particle)%charge
     endif  
  end do
  
  !Check for output to file if chunk is full or on the last ele
  if ((n_ele_chunk >= n_ele_in_file) .or. (n_ele_use == total_ele )) then
    !Reset chunk
    n_ele_chunk = 0
    
    !Prepare file name
    file_number = file_number + 1    
    write( particle_file_name, '(a, i0, 2a)') 'l', file_number, trim(file_body), '.dat'

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

    !Write number of particles to file
    write (particle_file, '(i0)') total_particles
  
    !Copy from scratch file to particle file
    rewind(scratch_file)
    do ix_particle = 1, total_particles
      read(scratch_file,  '(8es19.10E3)')  tmp_vec
      write(particle_file,  '(8es19.10E3)')  tmp_vec
    end do
    rewind(scratch_file)
  
    write (*, '(a, i8)')      'Number of particles   : ', total_particles
    write  (*,'(a, es13.3)')  'Total charge (C) 	  : ', total_charge
    print *, "Written: ", particle_file_name
  
    close(particle_file)
   
   
    !Make input file for dark current if asked
    if (make_dark_current_tracker_input_files) then
      write( dct_file_name, '(a, i0, 2a)') 'dct_l', file_number, trim(file_body), '.in'
          !Open output file
      open(dct_file, file = dct_file_name)
      write (dct_file, '(a)') '&dark_current_tracker_params'
      write (dct_file, '(3a)') ' lat_name           = "', trim(lat_name), '"'
      write (dct_file, '(3a)') ' particle_file_name = "', trim(particle_file_name), '"'
      write (dct_file, '(3a)') ' element_monitor_list = "', trim(element_monitor_list), '"'
      write (dct_file, '(a, es19.10E3)')  ' min_charge_ratio   = ', min_charge_ratio 
      write (dct_file, '(a)') '/'
      
      print *, "Written: ", dct_file_name
    end if
 
 
    !Reset counts for next file
    total_particles = 0
    total_charge = 0
  
  end if
  

    
end do


!Close scratch file
close(scratch_file)

print *, 'Total particle files: ', file_number

contains

!+
! Subroutine 
!
! Input
!   lat        -- lat_struct: Accelerator lattice
!   ele        -- ele_struct: current lattice element
!   particle   -- coord_struct: dark current particle
!   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!
!   field      -- em_field_struct : Local field used
!-

subroutine generate_field_emission_particles(ele, particle_list, n_particle_list)

type(ele_struct)       :: ele
type (coord_struct) :: particle_list(:)
integer                :: n_particle_list

type (coord_struct) :: particle
real(rp) :: total_len, d_len, d_section
real(rp) :: r, phi, freq

integer ix_wall_section
integer phase_ct, particle_ct_prelim

! This needs global variables: 
! numSources
! numPhases
! buffer
! angles_on
! A_fn, beta_fn
! Num_emitters_per_m2



total_len = 0

!Find length of cavity boundary (assume all cavities identical)
do ix_wall_section = 2, ubound(ele%wall3d(1)%section, 1)
total_len = total_len + sqrt( &
    (ele%wall3d(1)%section(ix_wall_section)%v(1)%radius_x - &
     ele%wall3d(1)%section(ix_wall_section - 1)%v(1)%radius_x)**2 + &
    (ele%wall3d(1)%section(ix_wall_section)%s - &
     ele%wall3d(1)%section(ix_wall_section - 1)%s)**2)
end do

print *, 'Wall length (m) = ', total_len

!Find distance between start and first particle, equal to half of
! the distance between consecutive particles
d_len = total_len / numSources / 2


!Start with 3D wall section 1
ix_wall_section = 1

!Particle momemta are all zero
particle%vec(2) = 0.0
particle%vec(4) = 0.0
particle%vec(6) = 0.0
particle%species = electron$

!Initialize particle counter
n_particle_list = 0

do 
  !Find the length of the current section (old: Find length to next section)
  d_section = sqrt( &
	   (ele%wall3d(1)%section(ix_wall_section + 1)%v(1)%radius_x - &
	   ele%wall3d(1)%section(ix_wall_section)%v(1)%radius_x)**2 + &
	   (ele%wall3d(1)%section(ix_wall_section + 1)%s - &
	   ele%wall3d(1)%section(ix_wall_section)%s)**2)

  !CASE 1: wall section is shorter than distance to next particle
  !------------------------------------------------------------
  if (d_section < d_len) then
	 d_len = d_len - d_section
	 ix_wall_section = ix_wall_section + 1

  !CASE 2 and 3: wall section is equal to distance to next particle
  !------------------------------------------------------------
  else    
	!====
	if (d_section .eq. d_len) then
		!Make a particle, all phases, and write to file
		 r = ele%wall3d(1)%section(ix_wall_section)%v(1)%radius_x - buffer
	     particle%vec(5) = ele%s - ele%value(l$) + ele%wall3d(1)%section(ix_wall_section)%s 
		!Update distance to next particle and wall section
		d_len = total_len / numSources
		ix_wall_section = ix_wall_section + 1
	 else !if (d_section > d_len) then
		!Make a particle, all phases, and write to file
		r = ele%wall3d(1)%section(ix_wall_section)%v(1)%radius_x + &
			  (ele%wall3d(1)%section(ix_wall_section + 1)%v(1)%radius_x - &
				ele%wall3d(1)%section(ix_wall_section)%v(1)%radius_x) / &
				d_section * d_len - buffer
		 particle%vec(5) = ele%s - ele%value(l$) + &
		 ele%wall3d(1)%section(ix_wall_section)%s + &
			(ele%wall3d(1)%section(ix_wall_section + 1)%s - &
			 ele%wall3d(1)%section(ix_wall_section)%s) / &
			d_section * d_len
		!Update distance to next particle; wall section stays same
		d_len = d_len + total_len / numSources
	 end if
	 !====
	! print *, 'particle at ', particle%vec(5)
	 
	 !Place particle at random angle phi; random_number has range {0,1}
	 if (angles_on) then
		call random_number(phi)
		phi = phi * twopi
	 else
		phi = angle0 * twopi
	 endif
		   
	 !Find x and y coordinates corresponding to angle phi
	 particle%vec(1) = cos(phi) * r 
	 particle%vec(3) = sin(phi) * r 

	 !Find charge and write to file for all phases
	 freq = ele%value(rf_frequency$) * ele%grid_field(1)%harmonic
	 
	 do phase_ct = 1, numPhases
		particle%t = (real(phase_ct) / numPhases) / freq + ele%value(ref_time_start$)

		!Determine weight (charge)
		call Fowler_Nordheim_current_calc(ele, particle, &
			 numPhases, A_fn, beta_fn)
		!Multiply by effective area. The total charge then needs
		!  to be multiplied by the number of emitters per unit area. 
		! \approx 2\pi \delta_l r
		particle%charge = particle%charge * twopi & 
						*( total_len / numSources) &
						* sqrt(particle%vec(1)**2 + particle%vec(3)**2) &
						* Num_emitters_per_m2            
				
		!Save particle to array
		if (particle%charge > 0 .and. particle%vec(5) < ele%s ) then
		  n_particle_list = n_particle_list + 1
		  particle_list(n_particle_list) = particle
		end if

	 end do
	end if

  !Stop once we hit the last section
  if (ix_wall_section .eq. ubound(ele%wall3d(1)%section, 1)) exit
end do

end subroutine generate_field_emission_particles


!+
! Subroutine Fowler_Nordheim_current_calc(ele, particle, numPhases,  A_fn, beta_fn)
!
! Routine to calculate the field emitted current (in amperes)
! according to the Fowler Nordheim model. 
!
! Note: The particle is assumed to be an electron
!
! Modules needed
!   use bmad
!
! Input
!   ele        -- ele_struct: current lattice element
!   particle   -- coord_struct: dark current particle
!   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!
!   field      -- em_field_struct : Local field used
!-

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

use bmad
use wall3d_mod

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, E_dot_n, freq
real(rp) :: t_rel, s_rel
integer, parameter :: particle_type = electron$

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



!Calculate E field at point
s_rel = particle%vec(5) - (ele%s - ele%value(l$)) 
t_rel = particle%t !Using absolute time tracking, no: - ele%value(ref_time_start$)

call em_field_calc (ele, lat%param, s_rel, particle, .false., field, .false., rf_time = t_rel)

!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(particle_type)**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(5) = s_rel
!particle_p%orb%vec(6) = particle_p%orb%vec(6) / e_tot

!Get normal vector to wall
d_radius = wall3d_d_radius([particle%vec(1), 0.0_rp, particle%vec(3), & 
           0.0_rp, s_rel, particle%vec(6) ], &
           ele, 1, norm)

!FOR CHECKING: Get norm . rhat to determine outward or inward facing normal
!print *, particle_p%orb%vec(1)*norm(1) + particle_p%orb%vec(3)*norm(2) 

!Particles are electrons
E_dot_n = dot_product(field%E, norm)

!If E field pushes particle into wall, charge is 0
!normal is outward facing? CHECK
if (E_dot_n <= 0) 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)))

I_fn = 3.85e-7 * A_fn * &
     (beta_fn * E_dot_n)**2  &
     * exp( -5.464e10 / (beta_fn * E_dot_n) )


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

!This was moved away: 
!Weight charge by radius
!particle%charge = particle%charge * &
!     sqrt(particle%vec(1)**2 + particle%vec(3)**2)

end subroutine



!+
! Subroutine init_random_seed
!
! Small subroutine to get us random numbers
! taken from fortranwiki.org/fortran/show/random_seed
!
! Input:
!   (none)
!
! Output:
!   (none)
!
!-

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