program monte_carlo

    use, intrinsic :: iso_fortran_env
    implicit none

    !> Initialize some variables for neutrons
    integer(int64), parameter :: n = 1e7
    integer(int64), parameter :: bins = 800
    real(real64), parameter :: eps = 1.0e-6_real64
    real(real64), parameter :: pi = 4.0_real64 * atan(1.0_real64)

    !> And for geometry
    real(real64), parameter :: L = 8.0_real64
    real(real64), parameter :: bin_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 cosine of the neutron direction vector
    real(real64), allocatable, dimension(:) :: x, u
    !> Track whether or not the neutron is alive
    logical, allocatable, dimension(:) :: is_alive
    !> And for tallying the flux
    real(real64), allocatable, dimension(:) :: neutron_bin_length, flux_tally, flux_sq_tally, flux_std_dev

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

    !> Allocate arrays
    allocate(x(n), u(n), is_alive(n), neutron_bin_length(bins), flux_tally(bins), flux_sq_tally(bins), flux_std_dev(bins))
    is_alive = .true.
    flux_tally = 0.0_real64
    flux_sq_tally = 0.0_real64
    flux_std_dev = 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, u)

    !> The main monte carlo loop 
    main: do i = 1,n
        !> Initialize empty neutron length
        neutron_bin_length = 0.0_real64

        !> Simulate the neutron 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)

            if(region .ge. size(boundaries) .or. region .le. 0) then
                !> Put neutrons back inside if they get lost
                u(i) = eps
            end if

            !> 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
                !> Moving in the +x direction
                d_boundary = (boundaries(region + 1) - x(i)) / u(i)
            else if(u(i) .lt. 0.0_real64) then
                !> Moving in the -x direction
                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)

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

            !> Update the neutron position
            x(i) = x(i) + d_step * u(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(u(i) .gt. 0.0_real64 .and. abs(x(i)-L) .lt. eps) then
                    !> Bin neutron at vacuum boundary at x=8.0
                    is_alive(i) = .false.
                else if(u(i) .lt. 0.0_real64 .and. x(i) .lt. eps) then
                    !> Reflect on specular boundary at x=0.0
                    u(i) = -u(i)
                    x(i) = -x(i)
                else
                    !> Internal boundary
                    if(u(i) .gt. 0.0_real64) then
                        x(i) = boundaries(region+1) + 2.0_real64*eps
                    else
                        x(i) = boundaries(region) - 2.0_real64*eps
                    end if
                end if
            else
                !> A collision happens
                call random_number(xi)
                if(xi .le. sigma_s/sigma_t) then
                    !> This is a scattering interaction
                    call random_number(xi)
                    u(i) = 2.0_real64 * xi - 1.0_real64
                else
                    !> The particle is absorbed and lost
                    is_alive(i) = .false.
                end if
            end if

        end do single

        !> Update tallies
        flux_tally = flux_tally + neutron_bin_length
        flux_sq_tally = flux_sq_tally + neutron_bin_length**2

    end do main

    !> sqrt(x²/n - (x/n)²)/(dx²)
    flux_std_dev = sqrt((flux_sq_tally/real(n, real64) - (flux_tally/real(n, real64))**2)) / bin_dx**2

    !> Average region 1 flux
    avg_r1_flux = sum(flux_tally(1:floor(boundaries(2)/bin_dx)))/real(floor(boundaries(2)/bin_dx), real64)

    !> Normalize to region 1
    flux_tally = flux_tally / avg_r1_flux
    flux_std_dev = flux_std_dev / avg_r1_flux

    open(unit=10, file='flux.out', status='replace')
    write: do i = 1,bins
        write(10,'(3E15.7)') (i-0.5_real64)*bin_dx, flux_tally(i), flux_std_dev(i)
    end do write

    !> Deallocate everyting
    deallocate(x, u, is_alive, flux_tally)


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

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

    impure subroutine initialize_neutrons(n, x, u)
        integer(int64), intent(in) :: n
        real(real64), intent(out), dimension(n) :: x, u
        real(real64), dimension(n) :: r_region, r_position, r_mu
        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)

        !> Consider the total integrated source-distance for regions
        r_region = r_region * 101.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. 101.0_real64)
            x = 5.0_real64 + eps + r_position * (1.0_real64 - 2.0_real64*eps)
        end where

        !> Assign random angles to each particle
        u = 2.0_real64 * r_mu - 1.0_real64

    end subroutine initialize_neutrons

    pure subroutine tally_segment(x, u, d_step, bins, bin_dx, flux_tally)
        integer(int64), intent(in) :: bins
        real(real64), intent(in) :: x, u, d_step, bin_dx
        real(real64), intent(in out), dimension(bins) :: flux_tally
        real(real64) :: x_start, x_end, dx
        integer(int64) :: start_bin, end_bin, bin

        x_start = x
        x_end = x + d_step*u
        if (x_start .gt. x_end) then
            call swap(x_start, x_end)
        end if

        start_bin = ceiling(x_start/bin_dx)
        end_bin   = ceiling(x_end/bin_dx)

        !> Clamp to valid bin range
        start_bin = max(1, min(bins, start_bin))
        end_bin   = max(1, min(bins, end_bin))

        if (start_bin .eq. end_bin) then
            flux_tally(start_bin) = flux_tally(start_bin) + d_step
        else
            !> First partial bin
            dx = (start_bin*bin_dx - x_start)/abs(u)
            flux_tally(start_bin) = flux_tally(start_bin) + dx

            !> Full bins in between
            do bin = start_bin+1, end_bin-1
                flux_tally(bin) = flux_tally(bin) + bin_dx/abs(u)
            end do

            !> Last partial bin
            dx = (x_end - (end_bin-1)*bin_dx)/abs(u)
            flux_tally(end_bin) = flux_tally(end_bin) + dx
        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

end program monte_carlo