using Plots
using LinearAlgebra
using Random

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

const g = 9.81

# Malý valec
const Rsmall = 0.10
const msmall = 1.0
const Ismall = 0.5 * msmall * Rsmall^2

# Veľký valec
const Rbig = 1.0

# Trenie
const μ = 0.42

# Odrazivosť
const restitution = 0.08

# Fyzikálny krok
const dt = 0.0015

# Fyzikálne kroky na jeden obrázok
const STEPS_PER_FRAME = 12

# Počet veľkých valcov
const Nbig = 11

# Malé valce
const NROWS = 7
const NCOLS = 55

# Začiatočný a koncový sklon
const INITIAL_ANGLE = deg2rad(18.0)
const FINAL_ANGLE   = deg2rad(42.0)

# Otočenie celej scény
const SCENE_ROTATION = deg2rad(270.0)
const ROTATION_START = deg2rad(90.0)
const ROTATION_END = deg2rad(0.0)

# Pevný bod otáčania
const PIVOT = [0.0, 0.0]

Random.seed!(1234)

# ============================================================
# VEKTOROVÉ FUNKCIE
# ============================================================

function rotate(v, θ)

    c = cos(θ)
    s = sin(θ)

    return [
        c * v[1] - s * v[2],
        s * v[1] + c * v[2]
    ]
end

function cross2(a, b)

    return a[1] * b[2] -
           a[2] * b[1]
end

function rotate_around(point, pivot, θ)

    return pivot +
           rotate(
               point - pivot,
               θ
           )
end

# ============================================================
# UHOL ROVINY
# ============================================================

function plane_angle(frame)
    t = frame / 80.0
    rotation = ROTATION_START + t * (ROTATION_END - ROTATION_START)
    return SCENE_ROTATION + rotation
end

# ============================================================
# LOKÁLNE -> SVETOVÉ SÚRADNICE
# ============================================================

function world_position(local_pos, θ)

    return rotate_around(
        local_pos,
        PIVOT,
        θ
    )
end

# ============================================================
# MALÝ VALec
# ============================================================

mutable struct SmallCylinder

    pos::Vector{Float64}
    vel::Vector{Float64}
    omega::Float64

end

# ============================================================
# VEĽKÝ VALec
# ============================================================

struct BigCylinder

    pos::Vector{Float64}

end

# ============================================================
# GRAVITÁCIA V LOKÁLNYCH SÚRADNICIACH
# ============================================================

function gravity_local(θ)

    gworld = [0.0, -g]

    return rotate(
        gworld,
        -θ
    )
end

# ============================================================
# VEĽKÉ VALCE
# ============================================================

function create_big_cylinders()

    result = BigCylinder[]

    spacing =
        2.0 * Rbig * 0.98

    total_width =
        (Nbig - 1) * spacing

    # Pravý koniec radu je vľavo od pivotu.
    right_x = -1.5

    first_x =
        right_x - total_width

    for i in 1:Nbig

        x =
            first_x +
            (i - 1) * spacing

        # Veľké valce ležia na rovine.
        y = Rbig

        push!(
            result,
            BigCylinder([x, y])
        )
    end

    return result
end

# ============================================================
# MALÉ VALCE
#
# NOVÉ:
#
# Malé valce nezačínajú vo vzduchu.
#
# Prvý rad je položený PRIAMO NA VEĽKÝCH VALCOCH.
#
# Stred malého valca:
#
# y = Rbig + Rsmall
#
# teda sa dotýka vrcholu veľkého valca.
# ============================================================

function create_small_cylinders(big)

    result = SmallCylinder[]

    dx =
        2.0 * Rsmall * 1.04

    dy =
        2.0 * Rsmall * 1.08

    # --------------------------------------------------------
    # Malé valce budú rozložené nad radom veľkých valcov.
    # --------------------------------------------------------

    first_big_x =
        big[1].pos[1]

    last_big_x =
        big[end].pos[1]

    # Začneme trochu vľavo od prvého veľkého valca.
    x_start =
        first_big_x - 0.25

    # Končíme trochu pred pravým koncom.
    x_end =
        last_big_x + 0.25

    # --------------------------------------------------------
    # Prvý rad sedí priamo na veľkých valcoch.
    # --------------------------------------------------------

    y_start =
        Rbig + Rsmall

    for row in 1:NROWS

        # Každý druhý rad je posunutý.
        offset =
            isodd(row) ?
            0.0 :
            Rsmall

        # Výška ďalších vrstiev.
        y =
            y_start +
            (row - 1) * dy

        x = x_start

        while x <= x_end

            xx =
                x + offset

            # Ak je mimo oblasti, preskočí.
            if xx <= x_end

                # Malá náhodná odchýlka.
                xx +=
                    (rand() - 0.5) * 0.025

                yy =
                    y +
                    (rand() - 0.5) * 0.015

                # Na začiatku prakticky bez
                # vlastnej rýchlosti.
                vx =
                    (rand() - 0.5) * 0.02

                vy =
                    (rand() - 0.5) * 0.02

                # Žiadne umelé roztočenie.
                omega = 0.0

                push!(
                    result,
                    SmallCylinder(
                        [xx, yy],
                        [vx, vy],
                        omega
                    )
                )
            end

            x += dx
        end
    end

    return result
end

# ============================================================
# KONTAKT S ROVINOU
# ============================================================

function collide_with_plane!(p)

    n = [0.0, 1.0]
    t = [1.0, 0.0]

    # Spodný kontaktný bod.
    r = [0.0, -Rsmall]

    contact =
        p.pos + r

    penetration =
        -contact[2]

    if penetration > 0

        p.pos[2] += penetration

        vc = [
            p.vel[1] -
            p.omega * r[2],

            p.vel[2] +
            p.omega * r[1]
        ]

        vn =
            dot(vc, n)

        if vn < 0

            Jn =
                -(1.0 + restitution) *
                vn /
                (1.0 / msmall)

            p.vel +=
                Jn * n / msmall

            # ----------------------------
            # TRENIE
            # ----------------------------

            vc = [
                p.vel[1] -
                p.omega * r[2],

                p.vel[2] +
                p.omega * r[1]
            ]

            vt =
                dot(vc, t)

            inv_mass_t =
                1.0 / msmall +
                Rsmall^2 / Ismall

            Jt =
                -vt /
                inv_mass_t

            Jt =
                clamp(
                    Jt,
                    -μ * Jn,
                     μ * Jn
                )

            impulse =
                Jt * t

            p.vel +=
                impulse / msmall

            p.omega +=
                cross2(
                    r,
                    impulse
                ) / Ismall
        end
    end
end

# ============================================================
# MALÝ - MALÝ
# ============================================================

function collide_small_small!(a, b)

    d =
        b.pos - a.pos

    dist2 =
        dot(d, d)

    min_dist =
        2.0 * Rsmall

    if dist2 >= min_dist^2 ||
       dist2 < 1e-14

        return
    end

    dist =
        sqrt(dist2)

    n =
        d / dist

    t =
        [-n[2], n[1]]

    # Oprava preniknutia
    penetration =
        min_dist - dist

    correction =
        0.5 *
        penetration *
        n

    a.pos -= correction
    b.pos += correction

    # Kontaktné body
    ra =
        Rsmall * n

    rb =
        -Rsmall * n

    va = [
        a.vel[1] -
        a.omega * ra[2],

        a.vel[2] +
        a.omega * ra[1]
    ]

    vb = [
        b.vel[1] -
        b.omega * rb[2],

        b.vel[2] +
        b.omega * rb[1]
    ]

    relative =
        vb - va

    vn =
        dot(relative, n)

    if vn > 0
        return
    end

    # Normálový impulz
    Jn =
        -(1.0 + restitution) *
        vn /
        (2.0 / msmall)

    impulse_n =
        Jn * n

    a.vel -=
        impulse_n / msmall

    b.vel +=
        impulse_n / msmall

    # Trecí impulz
    va = [
        a.vel[1] -
        a.omega * ra[2],

        a.vel[2] +
        a.omega * ra[1]
    ]

    vb = [
        b.vel[1] -
        b.omega * rb[2],

        b.vel[2] +
        b.omega * rb[1]
    ]

    relative =
        vb - va

    vt =
        dot(relative, t)

    inv_mass_t =
        2.0 / msmall +
        2.0 * Rsmall^2 / Ismall

    Jt =
        -vt /
        inv_mass_t

    Jt =
        clamp(
            Jt,
            -μ * Jn,
             μ * Jn
        )

    impulse_t =
        Jt * t

    a.vel -=
        impulse_t / msmall

    b.vel +=
        impulse_t / msmall

    a.omega +=
        cross2(
            ra,
            -impulse_t
        ) / Ismall

    b.omega +=
        cross2(
            rb,
            impulse_t
        ) / Ismall
end

# ============================================================
# MALÝ - VEĽKÝ
# ============================================================

function collide_small_big!(
    p,
    big
)

    d =
        p.pos - big.pos

    dist2 =
        dot(d, d)

    min_dist =
        Rsmall + Rbig

    if dist2 >= min_dist^2 ||
       dist2 < 1e-14

        return
    end

    dist =
        sqrt(dist2)

    n =
        d / dist

    t =
        [-n[2], n[1]]

    # Oprava preniknutia
    penetration =
        min_dist - dist

    p.pos +=
        penetration * n

    # Kontaktný bod malého valca
    r =
        -Rsmall * n

    vc = [
        p.vel[1] -
        p.omega * r[2],

        p.vel[2] +
        p.omega * r[1]
    ]

    relative =
        -vc

    vn =
        dot(relative, n)

    if vn >= 0
        return
    end

    # Normálový impulz
    Jn =
        -(1.0 + restitution) *
        vn /
        (1.0 / msmall)

    p.vel +=
        Jn * n / msmall

    # Trecí impulz
    vc = [
        p.vel[1] -
        p.omega * r[2],

        p.vel[2] +
        p.omega * r[1]
    ]

    vt =
        dot(vc, t)

    inv_mass_t =
        1.0 / msmall +
        Rsmall^2 / Ismall

    Jt =
        -vt /
        inv_mass_t

    Jt =
        clamp(
            Jt,
            -μ * Jn,
             μ * Jn
        )

    impulse =
        Jt * t

    p.vel +=
        impulse / msmall

    p.omega +=
        cross2(
            r,
            impulse
        ) / Ismall
end

# ============================================================
# POHYB
# ============================================================

function integrate!(
    p,
    θ
)

    glocal =
        gravity_local(θ)

    p.vel +=
        glocal * dt

    p.pos +=
        p.vel * dt

    p.vel *=
        0.99995
end

# ============================================================
# FYZIKÁLNY KROK
# ============================================================

function physics_step!(
    particles,
    big,
    θ
)

    # Gravitácia
    for p in particles
        integrate!(
            p,
            θ
        )
    end

    # Rovina
    for p in particles
        collide_with_plane!(p)
    end

    # Malé - veľké
    for p in particles

        for b in big

            collide_small_big!(
                p,
                b
            )
        end
    end

    # Malé - malé
    N =
        length(particles)

    for i in 1:N-1

        for j in i+1:N

            collide_small_small!(
                particles[i],
                particles[j]
            )
        end
    end
end

# ============================================================
# KRUHY
# ============================================================

function circle_points(
    center,
    radius,
    θplane
)

    angles =
        range(
            0,
            2π,
            length=32
        )

    x =
        Float64[]

    y =
        Float64[]

    for a in angles

        point_local = [
            radius * cos(a),
            radius * sin(a)
        ]

        world =
            world_position(
                center + point_local,
                θplane
            )

        push!(
            x,
            world[1]
        )

        push!(
            y,
            world[2]
        )
    end

    return x, y
end

# ============================================================
# ZNAČKA ROTÁCIE MALÉHO VALCA
# ============================================================

function roller_marker(
    center,
    radius,
    omega,
    θplane
)

    marker_angle =
        omega * 0.04

    point =
        rotate(
            [0.0, radius * 0.75],
            marker_angle
        )

    p1 =
        world_position(
            center,
            θplane
        )

    p2 =
        world_position(
            center + point,
            θplane
        )

    return (
        [p1[1], p2[1]],
        [p1[2], p2[2]]
    )
end

# ============================================================
# VYTVORENIE OBJEKTOV
# ============================================================

big =
    create_big_cylinders()

particles =
    create_small_cylinders(big)

println(
    "Počet malých valcov: ",
    length(particles)
)

# ============================================================
# ANIMÁCIA
# ============================================================

frames =
    0:80

anim =
    @animate for frame in frames

        θplane =
            plane_angle(frame)

        # ----------------------------
        # FYZIKA
        # ----------------------------

        for step in 1:STEPS_PER_FRAME

            physics_step!(
                particles,
                big,
                θplane
            )
        end

        # ----------------------------
        # PLOT
        # ----------------------------

        plt =
            plot(
                aspect_ratio = :equal,
                xlim = (-12, 12),
                ylim = (-12, 10),
                legend = false,
                axis = false,
                grid = false,
                size = (1100, 750),
                background_color = :white
            )

        # ----------------------------
        # ROVINA
        # ----------------------------

        plane_half_width =
            13.0

        plane_a =
            world_position(
                [-plane_half_width, 0.0],
                θplane
            )

        plane_b =
            world_position(
                [plane_half_width, 0.0],
                θplane
            )

        plot!(
            [
                plane_a[1],
                plane_b[1]
            ],
            [
                plane_a[2],
                plane_b[2]
            ],
            linewidth = 7,
            color = :black
        )

        # ----------------------------
        # PIVOT
        # ----------------------------

        scatter!(
            [PIVOT[1]],
            [PIVOT[2]],
            markersize = 7,
            color = :red
        )

        # ----------------------------
        # VEĽKÉ VALCE
        # ----------------------------

        for b in big

            x, y =
                circle_points(
                    b.pos,
                    Rbig,
                    θplane
                )

            plot!(
                x,
                y,
                linewidth = 2,
                color = :black,
                fill = :gray70
            )
        end

        # ----------------------------
        # MALÉ VALCE
        # ----------------------------

        for p in particles

            x, y =
                circle_points(
                    p.pos,
                    Rsmall,
                    θplane
                )

            plot!(
                x,
                y,
                linewidth = 1,
                color = :black,
                fill = :lightblue
            )

            mx, my =
                roller_marker(
                    p.pos,
                    Rsmall,
                    p.omega,
                    θplane
                )

            plot!(
                mx,
                my,
                linewidth = 1,
                color = :red
            )
        end

        title!(
            "Valce — malé valce sa kĺžu po veľkých"
        )
    end

# ============================================================
# VIDEO
# ============================================================

mp4(
    anim,
    "valce_na_velkych.mp4",
    fps = 12
)

println()
println(
    "Hotovo: valce_na_velkych.mp4"
)