module particle_simulation
  implicit none
  private
  public :: init_particles, update_system, output_particles, output_container

  ! Simulation parameters
  double precision, parameter :: transition_height = 0.4d0  ! Height of transition region
  integer, parameter, public :: num_particles = 1000
  double precision, parameter, public :: dt = 0.00001d0
  double precision, parameter, public :: end_time = 2.0d0
  integer, parameter, public :: output_interval = 1000
    
  ! Container shape parameters
  double precision, parameter :: upper_width = 0.2d0
  double precision, parameter :: narrow_width = 0.02d0
  double precision, parameter :: container_height = 2.0d0
  double precision, parameter :: narrow_section_height = 0.02d0
  
  ! Particle parameters
  double precision, parameter :: particle_diameter = 0.02d0
  double precision, parameter :: particle_radius = particle_diameter / 2.0d0
  double precision, parameter :: particle_restitution = 0.5d0
  double precision, parameter :: wall_restitution = 0.7d0
  
  ! Physical parameters
  double precision, parameter :: gravity = 9.81d0
  double precision, parameter :: pi = 3.14159265358979d0
  
  type particle
    double precision :: pos(3)    ! Position
    double precision :: vel(3)    ! Velocity
    logical :: active            ! Activity flag
  end type particle
  
  type(particle), allocatable :: particles(:)

contains
subroutine init_particles()
  implicit none
  integer :: i
  double precision :: theta, r, z
  double precision :: rand_val1, rand_val2, rand_val3, rand_val4
  double precision :: initial_height = 0.2d0
  double precision :: x, y
  
  allocate(particles(num_particles))
  
  do i = 1, num_particles
    call random_number(rand_val1)
    call random_number(rand_val2)
    call random_number(rand_val3)
    call random_number(rand_val4)
    
    theta = 2.0d0 * pi * rand_val1
    r = upper_width * sqrt(rand_val2)
    z = container_height/2.0d0 - initial_height * rand_val3
    
    x = r * cos(theta)
    y = r * sin(theta)
    
    particles(i)%pos(1) = x
    particles(i)%pos(2) = y
    particles(i)%pos(3) = z
    
    particles(i)%vel(1) = 0.0d0
    particles(i)%vel(2) = 0.0d0
    particles(i)%vel(3) = 0.0d0
    particles(i)%active = .true.
  end do
end subroutine init_particles

subroutine update_system()
  implicit none
  call update_particles()
  call handle_collisions()
end subroutine update_system

subroutine update_particles()
  implicit none
  integer :: i
  
  do i = 1, num_particles
    if (.not. particles(i)%active) cycle
    
    particles(i)%vel(3) = particles(i)%vel(3) - gravity * dt
    particles(i)%pos = particles(i)%pos + particles(i)%vel * dt
    
    call check_boundary_collision(i)
  end do
end subroutine update_particles

subroutine check_boundary_collision(i)
  implicit none
  integer, intent(in) :: i
  double precision :: r, z, normal(3), boundary_radius, t
  
  r = sqrt(particles(i)%pos(1)**2 + particles(i)%pos(2)**2)
  z = particles(i)%pos(3)
  
  ! Calculate radius at current height (smooth transition)
  if (abs(z) < narrow_section_height/2) then
    boundary_radius = narrow_width
  else if (abs(z) < (narrow_section_height/2 + transition_height)) then
    t = (abs(z) - narrow_section_height/2) / transition_height
    boundary_radius = narrow_width + (upper_width - narrow_width) * &
                   (3.0d0 * t**2 - 2.0d0 * t**3)
  else
    boundary_radius = upper_width
  endif
   
  ! Wall collision
  if (r > boundary_radius - particle_radius) then
    normal = [-particles(i)%pos(1)/r, -particles(i)%pos(2)/r, 0.0d0]
    call apply_boundary_collision(i, normal)
    particles(i)%pos(1:2) = particles(i)%pos(1:2) * &
                           ((boundary_radius - particle_radius)/r)
  endif
  
  ! Bottom collision
  if (z < -container_height/2 + particle_radius) then
    normal = [0.0d0, 0.0d0, 1.0d0]
    call apply_boundary_collision(i, normal)
    particles(i)%pos(3) = -container_height/2 + particle_radius
  endif
  
  ! Top collision
  if (z > container_height/2 - particle_radius) then
    normal = [0.0d0, 0.0d0, -1.0d0]
    call apply_boundary_collision(i, normal)
    particles(i)%pos(3) = container_height/2 - particle_radius
  endif
end subroutine check_boundary_collision

subroutine apply_boundary_collision(i, normal)
  implicit none
  integer, intent(in) :: i
  double precision, intent(in) :: normal(3)
  double precision :: v_normal
  
  v_normal = dot_product(particles(i)%vel, normal)
  
  if (v_normal < 0.0d0) then
    particles(i)%vel = particles(i)%vel - (1.0d0 + wall_restitution) * v_normal * normal
  endif
end subroutine apply_boundary_collision

subroutine handle_collisions()
  implicit none
  integer :: i, j
  double precision :: dx(3), dist, overlap
  double precision :: relative_vel(3), normal(3)
  double precision :: reduced_mass, j_n
  
  do i = 1, num_particles-1
    do j = i+1, num_particles
      if (.not. (particles(i)%active .and. particles(j)%active)) cycle
      
      dx = particles(i)%pos - particles(j)%pos
      dist = sqrt(sum(dx**2))
      
      if (dist < particle_diameter) then
        overlap = particle_diameter - dist
        
        if (dist > 1.0d-10) then
          normal = dx/dist
          particles(i)%pos = particles(i)%pos + 0.5d0 * overlap * normal
          particles(j)%pos = particles(j)%pos - 0.5d0 * overlap * normal
          
          relative_vel = particles(i)%vel - particles(j)%vel
          reduced_mass = 0.5d0
          j_n = -(1.0d0 + particle_restitution) * dot_product(relative_vel, normal) * reduced_mass
          
          particles(i)%vel = particles(i)%vel + j_n * normal
          particles(j)%vel = particles(j)%vel - j_n * normal
        end if
      end if
    end do
  end do
end subroutine handle_collisions

subroutine output_particles(step)
  implicit none
  integer, intent(in) :: step
  character(len=100) :: filename
  integer :: i, unit
  
  write(filename, '(A,I6.6,A)') 'particles_', step, '.vtk'
  open(newunit=unit, file=filename, status='replace')
  
  write(unit, '(A)') '# vtk DataFile Version 2.0'
  write(unit, '(A)') 'Particle Data'
  write(unit, '(A)') 'ASCII'
  write(unit, '(A)') 'DATASET POLYDATA'
  write(unit, '(A,I0,A)') 'POINTS ', num_particles, ' float'
  
  do i = 1, num_particles
    write(unit, '(3F12.6)') particles(i)%pos
  end do
  
  close(unit)
end subroutine output_particles

subroutine output_container()
  implicit none
  integer :: unit, i, j
  integer, parameter :: n_circle = 32
  integer, parameter :: n_vertical = 50
  double precision :: theta, z, r, t
  double precision :: points(3, (n_circle+1)*(n_vertical+1))
  integer :: point_count, cell_count
  
  open(newunit=unit, file='container.vtk', status='replace')
  
  write(unit, '(A)') '# vtk DataFile Version 2.0'
  write(unit, '(A)') 'Container Shape'
  write(unit, '(A)') 'ASCII'
  write(unit, '(A)') 'DATASET UNSTRUCTURED_GRID'
  
  point_count = 0
  do i = 0, n_vertical
    z = container_height/2.0d0 - container_height * dble(i)/dble(n_vertical)
    
    if (abs(z) < narrow_section_height/2) then
      r = narrow_width
    else if (abs(z) < (narrow_section_height/2 + transition_height)) then
      t = (abs(z) - narrow_section_height/2) / transition_height
      r = narrow_width + (upper_width - narrow_width) * &
          (3.0d0 * t**2 - 2.0d0 * t**3)
    else
      r = upper_width
    end if
    
    do j = 0, n_circle
      theta = 2.0d0 * pi * dble(j)/dble(n_circle)
      point_count = point_count + 1
      points(:, point_count) = [r*cos(theta), r*sin(theta), z]
    end do
  end do
  
  write(unit, '(A,I0,A)') 'POINTS ', point_count, ' float'
  do i = 1, point_count
    write(unit, '(3F12.6)') points(:,i)
  end do
  
  cell_count = n_vertical * n_circle
  write(unit, '(A,I0,A,I0)') 'CELLS ', cell_count, ' ', cell_count*5
  
  do i = 0, n_vertical-1
    do j = 0, n_circle-1
      write(unit, '(5I5)') 4, &
        i*(n_circle+1) + j, &
        i*(n_circle+1) + j + 1, &
        (i+1)*(n_circle+1) + j + 1, &
        (i+1)*(n_circle+1) + j
    end do
  end do
  
  write(unit, '(A,I0)') 'CELL_TYPES ', cell_count
  do i = 1, cell_count
    write(unit, '(A)') '9'
  end do
  
  close(unit)
end subroutine output_container

end module particle_simulation

program main
  use particle_simulation
  implicit none
  integer :: step
  double precision :: time
  
  call random_seed()
  call init_particles()
  call output_container()
  
  time = 0.0d0
  step = 0
  do while (time < end_time)
    call update_system()
    
    if (mod(step, output_interval) == 0) then
      call output_particles(step/output_interval)
    endif
    
    time = time + dt
    step = step + 1
  end do

end program main
