using Plots
using LinearAlgebra
using Random

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

# Gravitácie pre jednotlivé simulácie
const GRAVITIES = [9.81, 20.0, 4.0, 15.0]

# 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)

# Aktuálne:
# 270° + 90° = 360° na začiatku
# 270° + 0°  = 270° na konci
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(θ, g_value)

    gworld = [0.0, -g_value]

    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
# ============================================================

function create_small_cylinders(big)

    result = SmallCylinder[]

    dx =
        2.0 * Rsmall * 1.04

    dy =
        2.0 * Rsmall * 1.08

    first_big_x =
        big[1].pos[1]

    last_big_x =
        big[end].pos[1]

    x_start =
        first_big_x - 0.25

    x_end =
        last_big_x + 0.25

    y_start =
        Rbig + Rsmall

    for row in 1:NROWS

        offset =
            isodd(row) ?
            0.0 :
            Rsmall

        y =
            y_start +
            (row - 1) * dy

        x = x_start

        while x <= x_end

            xx =
                x + offset

            if xx <= x_end

                xx +=
                    (rand() - 0.5) * 0.025

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

                vx =
                    (rand() - 0.5) * 0.02

                vy =
                    (rand() - 0.5) * 0.02

                omega = 0.0

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

            x += dx
        end
    end

    return result
end

# ============================================================
# KOPÍROVANIE MALÝCH VALCOV
# ============================================================

function copy_particles(particles)

    result = SmallCylinder[]

    for p in particles

        push!(
            result,
            SmallCylinder(
                copy(p.pos),
                copy(p.vel),
                p.omega
            )
        )
    end

    return result
end

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

function collide_with_plane!(p)

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

    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

            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]]

    penetration =
        min_dist - dist

    correction =
        0.5 *
        penetration *
        n

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

    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

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

    impulse_n =
        Jn * n

    a.vel -=
        impulse_n / msmall

    b.vel +=
        impulse_n / msmall

    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]]

    penetration =
        min_dist - dist

    p.pos +=
        penetration * n

    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

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

    p.vel +=
        Jn * n / msmall

    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,
    θ,
    g_value
)

    glocal =
        gravity_local(
            θ,
            g_value
        )

    p.vel +=
        glocal * dt

    p.pos +=
        p.vel * dt

    p.vel *=
        0.99995
end

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

function physics_step!(
    particles,
    big,
    θ,
    g_value
)

    # Gravitácia
    for p in particles

        integrate!(
            p,
            θ,
            g_value
        )
    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 ZÁKLADNÉHO STAVU
# ============================================================

# Veľké valce sú rovnaké pre všetky simulácie.
big =
    create_big_cylinders()

# Vytvoríme JEDEN počiatočný stav.
# Ten potom skopírujeme pre všetky štyri g.
base_particles =
    create_small_cylinders(big)

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

# ============================================================
# ŠTYRI SAMOSTATNÉ SIMULÁCIE
# ============================================================

particles_981 =
    copy_particles(base_particles)

particles_20 =
    copy_particles(base_particles)

particles_4 =
    copy_particles(base_particles)

particles_15 =
    copy_particles(base_particles)

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

frames =
    0:80

anim =
    @animate for frame in frames

        θplane =
            plane_angle(frame)

        # ----------------------------------------------------
        # FYZIKA PRE g = 9.81
        # ----------------------------------------------------

        for step in 1:STEPS_PER_FRAME

            physics_step!(
                particles_981,
                big,
                θplane,
                9.81
            )
        end

        # ----------------------------------------------------
        # FYZIKA PRE g = 20
        # ----------------------------------------------------

        for step in 1:STEPS_PER_FRAME

            physics_step!(
                particles_20,
                big,
                θplane,
                20.0
            )
        end

        # ----------------------------------------------------
        # FYZIKA PRE g = 4
        # ----------------------------------------------------

        for step in 1:STEPS_PER_FRAME

            physics_step!(
                particles_4,
                big,
                θplane,
                4.0
            )
        end

        # ----------------------------------------------------
        # FYZIKA PRE g = 15
        # ----------------------------------------------------

        for step in 1:STEPS_PER_FRAME

            physics_step!(
                particles_15,
                big,
                θplane,
                15.0
            )
        end

        # ====================================================
        # FUNKCIA NA VYKRESLENIE JEDNÉHO GRAFU
        # ====================================================

        function draw_scene!(
            plt,
            particles,
            g_value
        )

            # -----------------------------------------------
            # 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!(
                plt,
                [
                    plane_a[1],
                    plane_b[1]
                ],
                [
                    plane_a[2],
                    plane_b[2]
                ],
                linewidth = 5,
                color = :black
            )

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

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

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

            for b in big

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

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

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

            for p in particles

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

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

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

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

            # -----------------------------------------------
            # NÁZOV
            # -----------------------------------------------

            title!(
                plt,
                "g = $(g_value) m/s²",
                fontsize = 11
            )

            return plt
        end

        # ====================================================
        # 4 GRAFY
        # ====================================================

        plt1 =
            plot(
                aspect_ratio = :equal,
                xlim = (-12, 12),
                ylim = (-12, 10),
                legend = false,
                axis = false,
                grid = false,
                background_color = :white
            )

        plt2 =
            plot(
                aspect_ratio = :equal,
                xlim = (-12, 12),
                ylim = (-12, 10),
                legend = false,
                axis = false,
                grid = false,
                background_color = :white
            )

        plt3 =
            plot(
                aspect_ratio = :equal,
                xlim = (-12, 12),
                ylim = (-12, 10),
                legend = false,
                axis = false,
                grid = false,
                background_color = :white
            )

        plt4 =
            plot(
                aspect_ratio = :equal,
                xlim = (-12, 12),
                ylim = (-12, 10),
                legend = false,
                axis = false,
                grid = false,
                background_color = :white
            )

        # Vykreslenie jednotlivých simulácií

        draw_scene!(
            plt1,
            particles_981,
            9.81
        )

        draw_scene!(
            plt2,
            particles_20,
            20
        )

        draw_scene!(
            plt3,
            particles_4,
            4
        )

        draw_scene!(
            plt4,
            particles_15,
            15
        )

        # ====================================================
        # SPOJENIE DO 2 × 2
        # ====================================================

        plot(
            plt1,
            plt2,
            plt3,
            plt4,
            layout = (2, 2),
            size = (1400, 900),
            margin = 2Plots.mm
        )
    end

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

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

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