program monte_carlo

    use, intrinsic :: iso_fortran_env
    implicit none

    !> Initialize some variables for neutrons
    integer(int64), parameter :: n = 1e6
    integer(int64), parameter :: bins = 800

    !> 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 x-position and cosine of angle for each neutron
    real(real64), dimension(n) :: px, mu
    !> Track whether or not the neutron is alive
    logical, dimension(n) :: is_alive

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

    !> Initialize RNG to be actually random
    call random_seed()
    call initialize_neutrons(n, px, mu)

    !> The main monte carlo loop 
    main: do i = 1,n

        single: do while (is_alive(i))

            !> Find what region the neutron is in and get XS
            region = count(px(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 = 1e10_real64
            else
                !> Else, calculate with d = -ln(xi)/sigma_t
                call random_number(xi)
                d = -log(xi) / sigma_t
            end if

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

            !> If it goes through a wall, restart the particle at the new boundary
            if(d_boundary.le.d) then
                px(i) = px(i) + d_boundary*mu(i)
                if(px(i).le.0.0_real64) then
                    px(i) = 0.0_real64
                    mu(i) = -mu(i)
                else if (px(i).ge.L) then
                    is_alive = .false.
                    exit single
                end if
                cycle single
            end if

            !> If not, determine what type of interaction happens
            call random_number(xi)
            if(xi.le.sigma_s/sigma_t) then
                !> Scattering interaction
                !> Sample new isotropic angle
                call random_number(xi)
                mu(i) = 2.0_real64*xi - 1.0_real64
            else
                !> Absorption interaction
                is_alive(i) = .false.
            end if

        end do single


    end do main


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, px, mu)
        integer(int64), intent(in) :: n
        real(real64), intent(out), dimension(n) :: px, mu
        real(real64), dimension(n) :: r_region, r_position, r_mu

        !> 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 * 106.0_real64

        !> Assign a random position in each source region
        where(r_region.lt.100)
            px = r_position * 2.0_real64
        else where(r_region.lt.105)
            px = 2.0_real64 + r_position
        else where
            px = 5.0_real64 + r_position
        end where

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

    end subroutine initialize_neutrons

end program monte_carlo