#julia --project=hadice_env hadice_A_vs_B.jl
using Pkg

# Prvýkrát odkomentuj:
 Pkg.add(["GLMakie", "DifferentialEquations"])

#GLMakie.activate!(force=true)

using GLMakie
using DifferentialEquations
using LinearAlgebra
using Printf
using Statistics


# ============================================================
# SIMULÁCIA: dve identické hadice
#
# A = tri prípoje z rôznych smerov
# B = tri prípoje z jednej strany
#
# Hadica:
#   - dĺžka približne 70 m
#   - obdĺžnikový tvar
#   - zaoblené rohy
#
# ============================================================


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

const N = 90

# Celková dĺžka hadice
const L = 70.0

# Obdĺžnik
const RECT_WIDTH  = 25.0
const RECT_HEIGHT = 12.0

# Polomer zaoblených rohov
#
# Pri týchto rozmeroch vychádza obvod približne 70 m.
#
const CORNER_RADIUS = 2.329


# Dĺžka jedného segmentu
const ds = L / (N - 1)


# ------------------------------------------------------------
# Mechanické parametre
# ------------------------------------------------------------

const segment_mass = 0.055

const k_stretch = 5000.0
const k_bend = 35.0

const damping = 1.8
const friction = 0.55

const g = 9.81


# ------------------------------------------------------------
# Voda
# ------------------------------------------------------------

const ρ = 1000.0
const Cd = 0.82


# ------------------------------------------------------------
# Čas zapnutia vody
# ------------------------------------------------------------

const water_on = 0.7


# ============================================================
# TLAKY PRÍPOJOV
# ============================================================

# A:
# každý prípoj má iný tlak

const pressures_A = [
    2.0e5,     # 2 bar
    3.0e5,     # 3 bar
    4.0e5      # 4 bar
]


# B:
# poradie tlakov je opačné

const pressures_B = [
    4.0e5,     # 4 bar
    3.0e5,     # 3 bar
    2.0e5      # 2 bar
]


# ============================================================
# SILA VODNÉHO PRÍPOJA
# ============================================================

function water_reaction_force(p, d)

    A = π * d^2 / 4

    v = Cd * sqrt(2p / ρ)

    Q = A * v

    mdot = ρ * Q

    # Reakcia na zmenu hybnosti prúdu
    return mdot * v

end


# ============================================================
# PRIEMERY PRÍPOJOV
# ============================================================

const diameters = [
    0.012,
    0.012,
    0.012
]


# ============================================================
# ZAOBLENÝ OBDĹŽNIK
# ============================================================
#
# Vytvorí sa presná geometria približne 70 m dlhej hadice.
#
# Následne sa geometria znovu vzorkuje tak, aby boli
# jednotlivé body približne rovnomerne rozmiestnené.
#
# ============================================================

function rounded_rectangle_shape(offset_x, offset_y)

    W = RECT_WIDTH
    H = RECT_HEIGHT
    r = CORNER_RADIUS


    # --------------------------------------------------------
    # Pomocná funkcia na pridávanie bodov
    # --------------------------------------------------------

    points_x = Float64[]
    points_y = Float64[]


    function add_point(x, y)

        push!(points_x, x)
        push!(points_y, y)

    end


    # --------------------------------------------------------
    # Koľko bodov použijeme na jednotlivé úseky
    # --------------------------------------------------------

    nline = 80
    narc  = 50


    # ========================================================
    # HORNÁ STRANA
    # ========================================================

    for q in range(
        0.0,
        1.0,
        length=nline
    )

        x = -W/2 + r + q * (W - 2r)
        y = H/2

        add_point(x, y)

    end


    # ========================================================
    # PRAVÝ HORNÝ ROH
    # ========================================================

    for θ in range(
        π/2,
        0.0,
        length=narc
    )

        x = W/2 - r + r*cos(θ)
        y = H/2 - r + r*sin(θ)

        add_point(x, y)

    end


    # ========================================================
    # PRAVÁ STRANA
    # ========================================================

    for q in range(
        0.0,
        1.0,
        length=nline
    )

        x = W/2
        y = H/2 - r - q * (H - 2r)

        add_point(x, y)

    end


    # ========================================================
    # PRAVÝ DOLNÝ ROH
    # ========================================================

    for θ in range(
        0.0,
        -π/2,
        length=narc
    )

        x = W/2 - r + r*cos(θ)
        y = -H/2 + r + r*sin(θ)

        add_point(x, y)

    end


    # ========================================================
    # DOLNÁ STRANA
    # ========================================================

    for q in range(
        0.0,
        1.0,
        length=nline
    )

        x = W/2 - r - q * (W - 2r)
        y = -H/2

        add_point(x, y)

    end


    # ========================================================
    # ĽAVÝ DOLNÝ ROH
    # ========================================================

    for θ in range(
        -π/2,
        -π,
        length=narc
    )

        x = -W/2 + r + r*cos(θ)
        y = -H/2 + r + r*sin(θ)

        add_point(x, y)

    end


    # ========================================================
    # ĽAVÁ STRANA
    # ========================================================

    for q in range(
        0.0,
        1.0,
        length=nline
    )

        x = -W/2
        y = -H/2 + r + q * (H - 2r)

        add_point(x, y)

    end


    # ========================================================
    # ĽAVÝ HORNÝ ROH
    # ========================================================

    for θ in range(
        π,
        π/2,
        length=narc
    )

        x = -W/2 + r + r*cos(θ)
        y = H/2 - r + r*sin(θ)

        add_point(x, y)

    end


    # ========================================================
    # POSLEDNÝ BOD NECHCEME DUPLIKOVAŤ
    # ========================================================

    pop!(points_x)
    pop!(points_y)


    # ========================================================
    # PREPOČET NA ROVNOMERNÉ VZDIALENOSTI
    # ========================================================

    cumulative = zeros(length(points_x))

    for i in 2:length(points_x)

        dx = points_x[i] - points_x[i-1]
        dy = points_y[i] - points_y[i-1]

        cumulative[i] =
            cumulative[i-1] +
            hypot(dx, dy)

    end


    total_length = cumulative[end]


    # Požadované rovnomerné vzdialenosti
    target = range(
        0.0,
        total_length,
        length=N
    )


    x = zeros(N)
    y = zeros(N)


    # --------------------------------------------------------
    # Lineárna interpolácia po obvode
    # --------------------------------------------------------

    for k in 1:N

        s = target[k]

        j = searchsortedlast(
            cumulative,
            s
        )

        j = clamp(
            j,
            1,
            length(cumulative) - 1
        )


        local_length =
            cumulative[j+1] -
            cumulative[j]


        if local_length < 1e-12

            α = 0.0

        else

            α =
                (s - cumulative[j]) /
                local_length

        end


        x[k] =
            (1-α) * points_x[j] +
            α * points_x[j+1]


        y[k] =
            (1-α) * points_y[j] +
            α * points_y[j+1]

    end


    # --------------------------------------------------------
    # Korekcia mierky tak, aby dĺžka bola presne L
    # --------------------------------------------------------

    current_length = 0.0

    for i in 1:N-1

        current_length += hypot(
            x[i+1] - x[i],
            y[i+1] - y[i]
        )

    end


    scale = L / current_length


    x .*= scale
    y .*= scale


    x .+= offset_x
    y .+= offset_y


    return x, y

end


# ============================================================
# SMERY PRÍPOJOV A
# ============================================================

# Tri rôzne smery

const dirs_A = [

    [ 0.0, -1.0],

    [ 1.0,  0.25],

    [-0.8,  0.6]

]


# ============================================================
# SMERY PRÍPOJOV B
# ============================================================

# Všetky tri prípoje smerujú rovnakým smerom

const dirs_B = [

    [0.0, -1.0],

    [0.0, -1.0],

    [0.0, -1.0]

]


# Normalizácia

for j in eachindex(dirs_A)

    dirs_A[j] ./= norm(dirs_A[j])
    dirs_B[j] ./= norm(dirs_B[j])

end


# ============================================================
# MIESTA PRÍPOJOV
# ============================================================

# Tri body na obvode hadice

const connection_indices = [
    18,
    45,
    72
]


# ============================================================
# DYNAMIKA HADICE
# ============================================================

function hose_dynamics!(du, u, p, t)

    # p obsahuje:
    #
    # p[1] = smery prípojov
    # p[2] = tlaky prípojov
    # p[3] = zapnutie vody

    directions = p[1]

    pressures = p[2]

    water_enabled = p[3]


    # --------------------------------------------------------
    # Stav
    # --------------------------------------------------------

    x = @view u[1:N]
    y = @view u[N+1:2N]

    vx = @view u[2N+1:3N]
    vy = @view u[3N+1:4N]


    # --------------------------------------------------------
    # Derivácie
    # --------------------------------------------------------

    dx = @view du[1:N]
    dy = @view du[N+1:2N]

    dvx = @view du[2N+1:3N]
    dvy = @view du[3N+1:4N]


    dx .= vx
    dy .= vy


    # --------------------------------------------------------
    # Sily
    # --------------------------------------------------------

    Fx = zeros(N)
    Fy = zeros(N)


    # ========================================================
    # GRAVITÁCIA
    # ========================================================

    Fy .-= segment_mass * g


    # ========================================================
    # POZDLŽNE PRUŽINY
    # ========================================================

    for i in 1:N-1

        rx = x[i+1] - x[i]
        ry = y[i+1] - y[i]

        dist = hypot(rx, ry)

        dist = max(dist, 1e-8)


        nx = rx / dist
        ny = ry / dist


        F = k_stretch * (dist - ds)


        Fx[i] += F * nx
        Fy[i] += F * ny


        Fx[i+1] -= F * nx
        Fy[i+1] -= F * ny

    end


    # ========================================================
    # OHYBOVÁ TUHOSŤ
    # ========================================================

    for i in 2:N-1

        bx =
            x[i-1] -
            2x[i] +
            x[i+1]


        by =
            y[i-1] -
            2y[i] +
            y[i+1]


        Fx[i] +=
            k_bend *
            bx /
            ds


        Fy[i] +=
            k_bend *
            by /
            ds

    end


    # ========================================================
    # TLMIENIE
    # ========================================================

    for i in 1:N

        Fx[i] -= damping * vx[i]
        Fy[i] -= damping * vy[i]

    end


    # ========================================================
    # TRENIE
    # ========================================================

    for i in 1:N

        speed = hypot(
            vx[i],
            vy[i]
        )


        if speed > 1e-4

            Fx[i] -=
                friction *
                segment_mass *
                g *
                vx[i] /
                speed


            Fy[i] -=
                friction *
                segment_mass *
                g *
                vy[i] /
                speed

        end

    end


    # ========================================================
    # VODNÉ PRÍPOJE
    # ========================================================

    if water_enabled && t >= water_on

        for j in 1:3

            i =
                connection_indices[j]


            F =
                water_reaction_force(
                    pressures[j],
                    diameters[j]
                )


            Fx[i] +=
                F *
                directions[j][1]


            Fy[i] +=
                F *
                directions[j][2]

        end

    end


    # ========================================================
    # URÝCHLENIE
    # ========================================================

    dvx .= Fx ./ segment_mass
    dvy .= Fy ./ segment_mass

end


# ============================================================
# VYTVORENIE DVOCH HADÍC
# ============================================================

# Väčší odstup, aby sa obdĺžniky neprekrývali.

xA, yA =
    rounded_rectangle_shape(
        -40.0,
        0.0
    )


xB, yB =
    rounded_rectangle_shape(
         40.0,
        0.0
    )


# ============================================================
# POČIATOČNÉ STAVY
# ============================================================

u0A = [
    xA
    yA
    zeros(N)
    zeros(N)
]


u0B = [
    xB
    yB
    zeros(N)
    zeros(N)
]


# ============================================================
# DIFERENCIÁLNE PROBLÉMY
# ============================================================

probA = ODEProblem(
    hose_dynamics!,
    u0A,
    (0.0, 10.0),
    (
        dirs_A,
        pressures_A,
        true
    )
)


probB = ODEProblem(
    hose_dynamics!,
    u0B,
    (0.0, 10.0),
    (
        dirs_B,
        pressures_B,
        true
    )
)


# ============================================================
# SIMULÁCIA
# ============================================================

println()
println("Spúšťam simuláciu A...")


solA = solve(
    probA,
    Tsit5(),
    saveat=0.025,
    maxiters=10^7
)


println("Spúšťam simuláciu B...")


solB = solve(
    probB,
    Tsit5(),
    saveat=0.025,
    maxiters=10^7
)


println("Simulácia dokončená.")


# ============================================================
# VÝPOČET ŤAŽISKA
# ============================================================

function center_of_mass(sol, frame)

    x = sol[1:N, frame]
    y = sol[N+1:2N, frame]

    return mean(x), mean(y)

end


# ============================================================
# VIZUALIZÁCIA
# ============================================================

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


ax = Axis(
    fig[1, 1],
    aspect=DataAspect(),

    title=
        "A: rôzne smery a tlaky   vs.   " *
        "B: rovnaký smer a rôzne tlaky",

    xlabel="x [m]",
    ylabel="y [m]"
)


# Stred medzi hadicami

vlines!(
    ax,
    0.0,
    linestyle=:dash
)


# ============================================================
# HADICA A
# ============================================================

xObsA =
    Observable(
        solA[1:N, 1]
    )


yObsA =
    Observable(
        solA[N+1:2N, 1]
    )


lines!(
    ax,
    xObsA,
    yObsA,
    linewidth=8
)


# ============================================================
# PRÍPOJE A
# ============================================================

for j in eachindex(
    connection_indices
)

    i =
        connection_indices[j]


    dir =
        dirs_A[j]


    # --------------------------------------------------------
    # Bod prípoja
    # --------------------------------------------------------

    point_x =
        lift(
            xObsA,
            yObsA
        ) do x, y

            [x[i]]

        end


    point_y =
        lift(
            xObsA,
            yObsA
        ) do x, y

            [y[i]]

        end


    scatter!(
        ax,
        point_x,
        point_y,
        markersize=15
    )


    # --------------------------------------------------------
    # Šípka
    # --------------------------------------------------------

    arrow_x =
        lift(xObsA) do x

            [
                x[i],
                x[i] +
                dir[1] * 2.0
            ]

        end


    arrow_y =
        lift(yObsA) do y

            [
                y[i],
                y[i] +
                dir[2] * 2.0
            ]

        end


    lines!(
        ax,
        arrow_x,
        arrow_y,
        linewidth=3
    )

end


# ============================================================
# HADICA B
# ============================================================

xObsB =
    Observable(
        solB[1:N, 1]
    )


yObsB =
    Observable(
        solB[N+1:2N, 1]
    )


lines!(
    ax,
    xObsB,
    yObsB,
    linewidth=8
)


# ============================================================
# PRÍPOJE B
# ============================================================

for j in eachindex(
    connection_indices
)

    i =
        connection_indices[j]


    dir =
        dirs_B[j]


    point_x =
        lift(
            xObsB,
            yObsB
        ) do x, y

            [x[i]]

        end


    point_y =
        lift(
            xObsB,
            yObsB
        ) do x, y

            [y[i]]

        end


    scatter!(
        ax,
        point_x,
        point_y,
        markersize=15
    )


    arrow_x =
        lift(xObsB) do x

            [
                x[i],
                x[i] +
                dir[1] * 2.0
            ]

        end


    arrow_y =
        lift(yObsB) do y

            [
                y[i],
                y[i] +
                dir[2] * 2.0
            ]

        end


    lines!(
        ax,
        arrow_x,
        arrow_y,
        linewidth=3
    )

end


# ============================================================
# POPIS POD GRAFOM
# ============================================================

Label(
    fig[2, 1],

    "A: tlaky = 2 / 3 / 4 bar, rôzne smery    |    " *
    "B: tlaky = 4 / 3 / 2 bar, rovnaký smer    |    " *
    "dĺžka hadice = $(L) m    |    " *
    "voda ON = $(water_on) s",

    tellwidth=false
)


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

# 15 fps = pomalšie prehrávanie.
#
# Fyzikálny čas simulácie:
# 10 sekúnd
#
# Počet uložených snímok:
# približne 401
#
# Výsledné video:
# približne 27 sekúnd

record(
    fig,
    "hadice_A_vs_B.mp4",

    1:length(solA.t);

    framerate=15

) do frame


    # --------------------------------------------------------
    # Hadica A
    # --------------------------------------------------------

    xObsA[] =
        solA[
            1:N,
            frame
        ]


    yObsA[] =
        solA[
            N+1:2N,
            frame
        ]


    # --------------------------------------------------------
    # Hadica B
    # --------------------------------------------------------

    xObsB[] =
        solB[
            1:N,
            frame
        ]


    yObsB[] =
        solB[
            N+1:2N,
            frame
        ]


    
# Nadpis s aktuálnym časom
# --------------------------------------------------------

t_val = solA.t[frame]

ax.title = string(
    "A: 2/3/4 bar – rôzne smery   |   ",
    "B: 4/3/2 bar – jeden smer   |   ",
    "t = ",
    round(t_val, digits=2),
    " s"
)

end


# ============================================================
# KONIEC
# ============================================================

println()
println("==============================================")
println("HOTOVO")
println("==============================================")
println()
println("Video:")
println("hadice_A_vs_B.mp4")
println()
println("Dĺžka hadice:")
println("$(L) m")
println()
println("Tvar:")
println("obdĺžnik so zaoblenými rohmi")
println()
println("A:")
println("  prípoj 1 = 2 bar")
println("  prípoj 2 = 3 bar")
println("  prípoj 3 = 4 bar")
println("  rôzne smery")
println()
println("B:")
println("  prípoj 1 = 4 bar")
println("  prípoj 2 = 3 bar")
println("  prípoj 3 = 2 bar")
println("  všetky rovnaký smer")
println()
println("Voda sa zapína:")
println("t = $(water_on) s")
println()
println("Video:")
println("15 fps")
println()
println("Poznámka:")
println("Toto je zjednodušený 2D mechanický model,")
println("nie CFD simulácia.")