using GLMakie
using LinearAlgebra
using Random

# =========================
# PARAMETRE
# =========================

Random.seed!(11)

N = 40
R = 0.04

m_ball = 1.0
I_ball = 0.5 * m_ball * R^2

dt = 0.004
T = 8.0

# pružina drôtu
Nx = 80
kx = 900.0        # tuhosť medzi bodmi
dx = 1.0 / (Nx-1)

wire_mass = 0.3

# kolízie
K = 12000.0
μ = 0.35

# =========================
# GULE (p, v, ω)
# =========================

balls = [
    (
        Point2f(0.2 + 0.6rand(), 0.1 + 0.8rand()),
        Point2f(0, 0),
        0.0
    )
    for _ in 1:N
]

# =========================
# PRUŽNÝ DRÔT (2D VLNA)
# =========================

wire_x = range(0, 1, Nx)
wire_y = [0.5 + 0.1*sin(6*x) for x in wire_x]
wire_v = zeros(Float64, Nx)

function wire_height(xi)
    wire_y[xi]
end

# =========================
# OBSERVABLES
# =========================

ball_obs = Observable(Point2f[])
wire_obs = Observable(Point2f[])

# =========================
# FIGÚRA
# =========================

fig = Figure(size=(1500,700))

ax = Axis(fig[1,1],
    title = "Elastic wire bouncing through granular field",
    aspect = DataAspect()
)

ax2 = Axis(fig[1,2],
    title = "System energy proxy",
    xlabel = "t",
    ylabel = "E"
)

limits!(ax, 0, 1, 0, 1)

E_hist = Float64[]
t_hist = Float64[]

E_obs = Observable(Point2f[])

scatter!(ax, ball_obs, markersize=18)
lines!(ax, wire_obs, linewidth=3)
lines!(ax2, E_obs)

# =========================
# FUNKCIE
# =========================

function update_wire!()
    # spring force (discrete laplacian)
    new_acc = zeros(Float64, Nx)

    for i in 2:Nx-1
        left = wire_y[i-1]
        mid  = wire_y[i]
        right= wire_y[i+1]

        new_acc[i] = kx * ((left + right - 2*mid))
    end

    for i in 2:Nx-1
        wire_v[i] += new_acc[i] * dt
        wire_v[i] *= 0.99
        wire_y[i] += wire_v[i] * dt
    end
end

function update_balls!()
    for i in 1:N
        p, v, w = balls[i]

        v += Point2f(0, -0.0012)

        if p[1] < 0.05 || p[1] > 0.95
            v = Point2f(-0.9v[1], v[2])
        end
        if p[2] < 0.05 || p[2] > 0.95
            v = Point2f(v[1], -0.9v[2])
        end

        balls[i] = (p + v, v, w)
    end
end

function collisions!()
    for i in 1:N, j in i+1:N
        p1, v1, w1 = balls[i]
        p2, v2, w2 = balls[j]

        d = p2 - p1
        dist = norm(d)

        if dist < 2R && dist > 1e-6
            n = d / dist
            t = Point2f(-n[2], n[1])

            rel = v2 - v1

            vn = dot(rel, n)
            vt = dot(rel, t)

            Jn = -1.2 * vn
            Jt = -μ * Jn * sign(vt)

            imp = Jn*n + Jt*t

            v1 -= imp / m_ball
            v2 += imp / m_ball

            w1 -= Jt / I_ball
            w2 += Jt / I_ball

            balls[i] = (p1, v1, w1)
            balls[j] = (p2, v2, w2)
        end
    end
end

function wire_ball_interaction!()
    E = 0.0

    for i in 1:N
        p, v, w = balls[i]

        idx = Int(clamp(round(Int, p[1]*(Nx-1))+1, 2, Nx-1))

        wy = wire_y[idx]
        diff = p[2] - wy

        if abs(diff) < R
            force = K * (R - abs(diff))
            dir = sign(diff + 1e-6)

            v += Point2f(0, force * dir * 0.0004)

            balls[i] = (p, v, w)

            # prenos energie do pružiny
            wire_v[idx] += force * dir * 0.0002

            E += force^2
        end
    end

    push!(E_hist, E)
end

# =========================
# MP4
# =========================

frames = Int(T/dt)

record(fig, "elastic_wire_bounce.mp4", 1:frames; framerate=60) do frame

    t = frame * dt

    update_wire!()
    update_balls!()
    collisions!()
    wire_ball_interaction!()

    ball_obs[] = [p for (p,_,_) in balls]

    wire_obs[] = [
        Point2f(wire_x[i], wire_y[i]) for i in 1:Nx
    ]

    push!(t_hist, t)

    E_obs[] = [
        Point2f(t_hist[i], E_hist[i]) for i in eachindex(E_hist)
    ]

    notify(ball_obs)
    notify(wire_obs)
    notify(E_obs)

    sleep(0.001)
end