program monte_carlo

    use omp_lib
    use, intrinsic :: iso_fortran_env
    implicit none

    !> Initialize some variables for neutrons
    integer(int64), parameter :: n = 1e5
    integer(int64), parameter :: bins = 800
    real(real64), parameter :: eps = 1.0e-12_real64

    !> And for geometry
    real(real64), parameter :: L = 8.0_real64
    real(real64), parameter :: dx = L / bins
    real(real64), dimension(6), parameter :: boundaries = [0.0_real64, 2.0_real64, 3.0_real64, 5.0_real64, 6.0_real64, 8.0_real64]

    !> These will contain the position and cosines of direction vector
    real(real64), allocatable, dimension(:) :: x, y, z, u, v, w
    !> Track whether or not the neutron is alive
    logical, allocatable, dimension(:) :: is_alive
    !> And for tallying the flux
    real(real64), allocatable, dimension(:) :: flux_tally, collision_tally
    real(real64), allocatable, dimension(:) :: thread_tally, col_thread_tally

    !> Some local trackers to be updated via subroutine call
    real(real64) :: sigma_t, sigma_s, q, d_flight, d_boundary, d_step, xi, next_boundary
    integer(int64) :: i, region

    !> Allocate arrays
    allocate(x(n), y(n), z(n), u(n), v(n), w(n), is_alive(n), flux_tally(n), collision_tally(n), thread_tally(n), col_thread_tally(n))
    is_alive = .true.
    flux_tally = 0.0_real64
    collision_tally = 0.0_real64
    thread_tally = 0.0_real64
    col_thread_tally = 0.0_real64

    !> Initialize RNG to be actually random
    call random_seed()
    call random_init(.false., .true.)

    !> Provide all neutrons an initial location/angle
    call initialize_neutrons(n, x, y, z, u, v, w)

    !$OMP PARALLEL PRIVATE(thread_tally, col_thread_tally, sigma_t, sigma_s, q, d_flight, d_boundary, d_step, xi, region, i)
    thread_tally=0.0_real64

    !> The main monte carlo loop 
    !$OMP DO
    main: do i = 1,n
    !write(*,'("n = ",i0)') i
    !write(*,'("Before: x=",E9.3," y=",E9.3," z=",E9.3)') x(i), y(i), z(i)
    !write(*,'("Before: u=",E9.3," v=",E9.3," w=",E9.3)') u(i), v(i), w(i)

        !> Skip if the neutron is already lost
        if(.not. is_alive(i)) cycle main

        !> Otherwise, simulate it until it is lost
        single: do while(is_alive(i))

            !> Find what region the neutron is in and get XS
            region = count(x(i) .ge. boundaries)
            call get_region_info(region, sigma_t, sigma_s, q)

            !> Determine the distance before collision
            if(sigma_t .eq. 0.0_real64) then
                !> If we are in a void, make it huge
                d_flight = huge(real64)
            else
                !> Else, calculate with d = -ln(xi)/sigma_t
                call random_number(xi)
                d_flight = -log(xi) / sigma_t
            end if

            !> Determine where the next boundary is and the distance
            if(u(i) .gt. 0.0_real64) then
                d_boundary = (boundaries(region + 1) - x(i)) / u(i)
            else if(u(i) .lt. 0.0_real64) then
                d_boundary = (boundaries(region) - x(i)) / u(i)
            else
                !> In the rare chance it is perfectly parallel with the walls
                d_boundary = huge(real64)
            end if

            !> How far we will actually go
            d_step = min(d_flight,d_boundary)
            !write(*,'("Step distance = ",E9.3)') d_step

            !> Tally the track-length contributions for the bin
            call tally_segment(x(i), u(i), d_step, bins, dx, thread_tally)

            !> Update the neutron position
            x(i) = x(i) + d_step * u(i)
            y(i) = y(i) + d_step * v(i)
            z(i) = z(i) + d_step * w(i)

            !> If the neutron crosses a boundary, restart the simulation at the new location
            if(d_boundary .le. d_flight) then
                !> We cross a boundary
                if(abs(x(i) - L) .lt. eps) then
                    !write(*,'(A)') "Escaped!"
                    !> Bin neutron at vacuum boundary at x=8.0
                    is_alive(i) = .false.
                    cycle main
                else if(x(i) .lt. eps) then
                    !write(*,'(A)') "Reflected!"
                    !> Reflect on specular boundary at x=0.0
                    u(i) = -u(i)
                    v(i) = -v(i)
                    w(i) = -w(i)
                    x(i) = eps*0.1_real64
                    cycle single
                else
                    !write(*,'(A)') "Internal boundary!"
                    !> Internal boundary
                    if(u(i) .gt. 0.0_real64) then
                        !write(*,'(A)') "Forwards reflection!"
                        x(i) = boundaries(region+1) + 2*eps
                    else
                        !write(*,'(A)') "Backwards reflection!"
                        x(i) = boundaries(region) - 2*eps
                    end if
                    cycle single
                end if
            end if

            !> If not, a collision happens. Tally it:
            !> The current 'bin' is ceiling(x/dx)
            col_thread_tally(max(ceiling(x/dx),0)) = col_thread_tally(ceiling(x/dx)) + 1

            !> And determine the type:
            call random_number(xi)
            if(xi .le. sigma_s / sigma_t) then
                !> Scattering interaction
                !> Sample new isotropic angle
                call scatter(u(i),v(i),w(i))
            else
                !> Absorption interaction
                is_alive(i) = .false.
            end if

            !write(*,'("x=",E9.3," y=",E9.3," z=",E9.3)') x(i), y(i), z(i)
            !write(*,'("u=",E9.3," v=",E9.3," w=",E9.3)') u(i), v(i), w(i)
        end do single

    end do main
    !$OMP END DO

    !$OMP CRITICAL
    flux_tally = flux_tally + thread_tally
    collision_tally = collision_tally + thread_tally
    !$OMP END CRITICAL
    !$OMP END PARALLEL
    !flux_tally = flux_tally / flux_tally(bins/8)
    flux_tally = flux_tally / (real(n, real64) * dx)
    collision_tally = collision_tally / (real(n, real64)*4000)

    open(unit=10, file='flux.out', status='replace')
    open(unit=11, file='collision.out', status='replace')
    write: do i = 1,bins
        write(10,'(2E15.7)') (i-0.5_real64)*dx, flux_tally(i)
        write(11,'(2E15.7)') (i-0.5_real64)*dx, collision_tally(i)
    end do write

    !> Deallocate everyting
    deallocate(x, y, z, u, v, w, is_alive, flux_tally, collision_tally, thread_tally, col_thread_tally)


contains
        
    pure elemental subroutine get_region_info(region, sigma_t, sigma_s, q)
        integer(int64), intent(in) :: region
        real(real64), intent(out) :: sigma_t, sigma_s, q

        if(region .eq. 1) then
            sigma_t = 50.0_real64; sigma_s = 0.0_real64; q = 50.0_real64
        else if(region .eq. 2) then
            sigma_t = 5.0_real64; sigma_s = 0.0_real64; q = 5.0_real64
        else if(region .eq. 3) then
            sigma_t = 0.0_real64; sigma_s = 0.0_real64; q = 0.0_real64
        else if(region .eq. 4) then
            sigma_t = 1.0_real64; sigma_s = 0.9_real64; q = 1.0_real64
        else if(region .eq. 5) then
            sigma_t = 1.0_real64; sigma_s = 0.9_real64; q = 0.0_real64
        end if
    end subroutine get_region_info

    impure subroutine initialize_neutrons(n, x, y, z, u, v, w)
        integer(int64), intent(in) :: n
        real(real64), intent(out), dimension(n) :: x, y, z, u, v, w
        real(real64), dimension(n) :: r_region, r_position, r_mu, r_phi
        real(real64), parameter :: pi = 4.0_real64 * atan(1.0_real64)

        !> Make some uniform random numbers
        call random_number(r_region)
        call random_number(r_position)
        call random_number(r_mu)
        call random_number(r_phi)

        !> Consider the total integrated source-distance for regions
        r_region = r_region * 106.0_real64

        !> Assign a random position in each source region
        where(r_region .lt. 100.0_real64)
            x = eps + r_position * (2.0_real64 - 2.0_real64*eps)
        else where(r_region .lt. 105.0_real64)
            x = 2.0_real64 + eps + r_position * (1.0_real64 - 2.0_real64*eps)
        else where(r_region .lt. 106.0_real64)
            x = 5.0_real64 + eps + r_position * (1.0_real64 - 2.0_real64*eps)
        end where

        !> Start on the x-axis
        y = 0.0_real64
        z = 0.0_real64

        !> Assign random angles to each particle
        u = sqrt(max(1.0_real64 - r_mu**2.0_real64, 0.0_real64))*cos(2.0_real64*pi*r_phi)
        v = sqrt(max(1.0_real64 - r_mu**2.0_real64, 0.0_real64))*sin(2.0_real64*pi*r_phi)
        w = r_mu

    end subroutine initialize_neutrons

    pure subroutine tally_segment(px, mu, d_step, bins, dx, thread_tally)
        integer(int64), intent(in) :: bins
        real(real64), intent(in) :: px, mu, d_step, dx
        real(real64), intent(in out), dimension(bins) :: thread_tally
        real(real64) :: x_start, x_end, x_left, x_right
        integer(int64) :: i_start, i_end, i

        !> If no distance was travelled
        if (d_step <= 0.0_real64) return

        !> Start and end position
        x_start = px
        x_end = px + mu * d_step

        !> Swap them if we're going backwards
        if (x_start > x_end) call swap(x_start, x_end)

        !> Start and end indices for bins
        i_start = ceiling(x_start / dx)
        i_end   = ceiling(x_end / dx)

        !> Clamp them too
        i_start = max(1, min(bins, i_start))
        i_end   = max(1, min(bins, i_end))

        if (i_start == i_end) then
            !> No bins were crossed. Add the length
            thread_tally(i_start) = thread_tally(i_start) + d_step
        else
            !> First partial bin exited
            x_right = i_start * dx
            thread_tally(i_start) = thread_tally(i_start) + (x_right - x_start)/abs(mu)

            !> Going across the interior bins
            do i = i_start+1, i_end-1
                thread_tally(i) = thread_tally(i) + dx/abs(mu)
            end do

            !> Last partial bin entered
            x_left = (i_end - 1) * dx
            thread_tally(i_end) = thread_tally(i_end) + (x_end - x_left)/abs(mu)
        end if

    end subroutine tally_segment

    pure elemental subroutine swap(x, y)
        real(real64), intent(in out) :: x, y
        real(real64) :: temp
        temp = x
        x = y
        y = temp
    end subroutine swap

    impure elemental subroutine scatter(u, v, w)
        real(real64), intent(in out) :: u, v, w
        real(real64) :: cos_theta, phi, sin_theta, cos_phi, sin_phi
        real(real64) :: up, vp, wp
        real(real64) :: sqrt_1_minus_w2
        real(real64) :: uprime, vprime, wprime
        real(real64) :: xi
        real(real64), parameter :: pi = 4.0_real64 * atan(1.0_real64)

        !> Generate a random mu
        call random_number(xi)
        cos_theta = 2.0_real64 * xi - 1.0_real64

        !> And a random phi
        call random_number(xi)
        phi = 2.0_real64 * pi * xi

        !> Introduce some quantities
        sin_theta = sqrt(1.0_real64 - cos_theta**2)
        cos_phi = cos(phi)
        sin_phi = sin(phi)
        sqrt_1_minus_w2 = sqrt(1.0_real64 - w**2)

        !> Lewis and Miller, Eq. 7-163
        up = sin_theta / sqrt_1_minus_w2 * (v*sin_phi - v*u*cos_phi) + u*cos_theta
        vp = sin_theta / sqrt_1_minus_w2 * (-u*sin_phi - w*v*cos_phi) + v*cos_theta
        wp = sin_theta * sqrt_1_minus_w2 * cos_phi + w*cos_theta

        !> Update new direction cosines
        u = up
        v = vp
        w = wp

    end subroutine scatter

end program monte_carlo