program langevin_motion

    use, intrinsic :: iso_fortran_env
    use :: omp_lib 

    implicit none

    !> Define some basic parameters
    integer(int64), parameter :: n = 1000
    real(real64), parameter :: pi = 2.0_real64*asin(1.0_real64)
    real(real64), parameter :: L = 1.0_real64

    !> Numerical constants and such
    real(real64), parameter :: dt = 0.00001_real64
    real(real64), parameter :: t_max = 1.0_real64
    real(real64), parameter :: kT = 10.0_real64
    real(real64), parameter :: g = 1.0_real64
    real(real64), parameter :: m = 1.0_real64
    real(real64), parameter :: eps = 1.0_real64
    real(real64), parameter :: sigma = 1E-03_real64
    real(real64), parameter :: rc = sigma*2**(1.0_real64/6.0_real64)
    real(real64), parameter :: pref1 = g, pref3 = sqrt(24.0_real64*kT*g/dt)

    !> And non-constants
    real(real64) :: t = 0.0_real64
    real(real64), dimension(n) :: dx, dy, dn
    real(real64), dimension(n) :: F
    real(real64), dimension(n) :: ran1, ran2
    integer(int64) :: i

    !> Vectors for positon, velocity, acceleration, and tracking 
    real(real64), dimension(n) :: px, py, vx, vy, ax, ay, vhx, vhy
    logical, dimension(n) :: is_tracked = .TRUE., are_interacting

    !> Prepare random numbers
    call random_seed()

    !> Start by initializing the particles with a random pos. and vel.
    call initialize_particles(px, py, vx, vy, ax, ay, kT, L, pi)

    !> Now simulate!
    do while (t .lt. t_max)

        !> Update half-step velocities
        vhx = vx + 0.5_real64*ax*dt
        vhy = vy + 0.5_real64*ay*dt

        !> Update the positions
        px = px + vhx*dt
        py = py + vhy*dt

        !> Impose the boundary conditions
        call impose_BC(px, py, vhx, vhy, is_tracked, L)

        !> Find the new acceleration
        ax = 0.0_real64
        ay = 0.0_real64

        !> Random numbers
        call random_number(ran1)
        ran1 = ran1 - 0.5_real64
        call random_number(ran2)
        ran2 = ran2 - 0.5_real64

        !> Random force
        !ax = ax - pref1*vhx + pref2*ran1
        !ay = ay - pref1*vhy + pref2*ran2

        !> Particle-particle interaction force
        interactions: do i = 1,n
            
            !> Calculate the distance between the particles:
            dx = px(i) - px
            dy = py(i) - py
            dn = sqrt(dx**2 + dy**2)

            !> Check which particles are interacting
            are_interacting = (dn .gt. 0.0) .and. (dn .lt. rc) .and. is_tracked

            !> Start with no forces
            F = 0.0_real64

            where (are_interacting)
                !> Lennard-Jones force
                F = 4.0_real64*eps*(12.0_real64*sigma**12/dn**13 - 6.0_real64*sigma**6/dn**7)
            end where

            !> Update the accelerations with the new force when a force was calculated
            ax(i) = ax(i) + sum(F*dx / (dn*m), mask=are_interacting)
            ay(i) = ay(i) + sum(F*dy / (dn*m), mask=are_interacting)

        end do interactions

        !> Update full velocities
        vx = vhx + 0.5_real64*ax*dt
        vy = vhy + 0.5_real64*ay*dt

        !> Update timestep
        t = t + dt

        !> Print posiitons to stdout
        !print '(e24.12,e24.12)', (px(i), py(i), i=1,n)

        print *, sum(0.5_real64*m*vx**2 + 0.5_real64*m*vy**2, mask=is_tracked)

        !> This splits timesteps into data blocks for Gnuplot
        !print *
        !print *

    end do

contains

    impure elemental subroutine initialize_particles(px, py, vx, vy, ax, ay, kT, L, pi)
        real(real64), intent(in out) :: px, py, vx, vy, ax, ay
        real(real64), intent(in) :: kT, L, pi
        real(real64) :: ran1, ran2

        !> Random number seed
        !> And 2 random numbers
        call random_number(ran1)
        call random_number(ran2)

        !> Initial position
        px = L*(ran1-0.5_real64)
        py = L*(ran2-0.5_real64)

        !> Initial acceleration is zero
        ax = 0.0_real64
        ay = 0.0_real64

        !> New random numbers
        call random_number(ran1)
        call random_number(ran2)

        !> Initial velocity using Box-Muller transform
        vx = sqrt(kT/m)*sqrt(-2.0_real64*log(ran1))*cos(2.0_real64*pi*ran2)
        vy = sqrt(kT/m)*sqrt(-2.0_real64*log(ran1))*sin(2.0_real64*pi*ran2)

    end subroutine initialize_particles

    pure elemental subroutine impose_BC(px, py, vhx, vhy, is_tracked, L)
        real(real64), intent(in out) :: px, py, vhx, vhy
        real(real64), intent(in) :: L
        logical, intent(in out) :: is_tracked
        
        if (is_tracked) then
            !> Top boundary
            if (py.gt.L/2) then
                py = L - py
                vhy = -vhy
            end if

            !> Bottom boundary
            if (py.lt.-L/2) then
                py = -L - py
                vhy = -vhy
            end if

            !> Left boundary
            if (px.gt.L/2) then
                px = L - px
                vhx = -vhx
            end if

            !> Right boundary
            if (px.lt.-L/2) then
                px = -L - px
                vhx = -vhx
            end if

            !> If the particle is still outside, it goes into the garbage
            if (abs(px).gt.L/2 .or. abs(py) .gt. L/2) is_tracked = .false.

        end if

    end subroutine impose_BC

end program langevin_motion