using Pkg

for p in ["GLMakie", "FFMPEG"]
    try
        Base.require(Symbol(p))
    catch
        Pkg.add(p)
    end
end

using GLMakie
using FFMPEG

############################################################
# SETTINGS
############################################################

const FPS = 60
const FRAMES = 700
const OUTFILE = "4granular_heap_final.gif"

const N = 800
const R = 0.16

const DT = 0.016
const G = 0.22

const RESTITUTION = 0.35
const FRICTION = 0.03

############################################################
# CYLINDER (STRAIGHT WALLS)
############################################################

const LEFT_WALL  = -3.0
const RIGHT_WALL =  3.0

############################################################
# GRID HASH (fast neighbor search)
############################################################

cellsize = 2R

hash(x,y) = (floor(Int, x/cellsize), floor(Int, y/cellsize))

############################################################
# MAIN
############################################################

function main()

    fig = Figure(resolution=(900,900))

    ax = Axis(fig[1,1],
        title="Granular Heap Formation (angle of repose emerges)",
        aspect=DataAspect()
    )

    limits!(ax, -5,5,-7,7)

    ########################################################
    # WALLS
    ########################################################

    lines!(ax, [LEFT_WALL, LEFT_WALL], [-7,7], linewidth=3)
    lines!(ax, [RIGHT_WALL, RIGHT_WALL], [-7,7], linewidth=3)

    ########################################################
    # PARTICLES
    ########################################################

    x = rand(N).*2 .- 1
    y = rand(N).*3 .+ 4

    vx = zeros(N)
    vy = zeros(N)

    pos = Observable(Point2f.(x,y))
    scatter!(ax, pos, markersize=7)

    display(fig)

    ########################################################
    # SIMULATION LOOP
    ########################################################

    record(fig, OUTFILE, 1:FRAMES; framerate=FPS) do i

        ####################################################
        # BUILD SPATIAL GRID
        ####################################################

        grid = Dict{Tuple{Int,Int}, Vector{Int}}()

        for j in 1:N
            c = hash(x[j], y[j])
            push!(get!(grid, c, Int[]), j)
        end

        ####################################################
        # PARTICLE COLLISIONS (LOCAL ONLY)
        ####################################################

        for j in 1:N

            cx, cy = hash(x[j], y[j])

            for dx in -1:1, dy in -1:1

                cell = get(grid, (cx+dx, cy+dy), nothing)
                cell === nothing && continue

                for k in cell
                    k == j && continue

                    dxv = x[k] - x[j]
                    dyv = y[k] - y[j]

                    dist2 = dxv^2 + dyv^2
                    min_d = 2R

                    if dist2 < min_d^2 && dist2 > 1e-10

                        dist = sqrt(dist2)
                        overlap = min_d - dist

                        nx = dxv / dist
                        ny = dyv / dist

                        vx[j] -= nx * overlap * 0.5
                        vy[j] -= ny * overlap * 0.5

                        vx[k] += nx * overlap * 0.5
                        vy[k] += ny * overlap * 0.5
                    end
                end
            end
        end

        ####################################################
        # UPDATE PARTICLES
        ####################################################

        for j in 1:N

            # gravity
            vy[j] -= G

            # integrate
            x[j] += vx[j] * DT
            y[j] += vy[j] * DT

            ####################################################
            # WALL COLLISIONS (VERTICAL CYLINDER)
            ####################################################

            if x[j] < LEFT_WALL + R
                x[j] = LEFT_WALL + R
                vx[j] = abs(vx[j]) * RESTITUTION
            end

            if x[j] > RIGHT_WALL - R
                x[j] = RIGHT_WALL - R
                vx[j] = -abs(vx[j]) * RESTITUTION
            end

            ####################################################
            # DEPTH-DEPENDENT FRICTION
            ####################################################

            f = clamp((y[j] + 7) / 14, 0, 1)

            vx[j] *= (1 - FRICTION*f)
            vy[j] *= (1 - 0.5*FRICTION*f)

            ####################################################
            # RESET TOP INJECTION
            ####################################################

            if y[j] < -7
                y[j] = 7
                x[j] = rand()*2 - 1
                vx[j] = 0
                vy[j] = 0
            end
        end

        ####################################################
        # UPDATE VISUALS
        ####################################################

        pos[] = Point2f.(x,y)

    end

    println("Saved => $OUTFILE")
    println("Press ENTER")
    readline()
end

main()