# A coplanar double pendulum

# https://discourse.julialang.org/t/symbolics-jl-modelingtoolkit-jl-differentialequations-jl-and-the-lagrangian-approach/89629/3

using Symbolics, ModelingToolkit
using Latexify # only needed to 'inspect' intermediate equations
using OrdinaryDiffEq
using PyPlot
using Colors
using Printf

@variables t θ(t) ϕ(t) 
@parameters m1, m2, l1, l2, g

# Differential operators
d = Differential(t)
D1 = Differential(θ)
D2 = Differential(ϕ)
D11 = Differential(d(θ))
D22 = Differential(d(ϕ))

# Cartesian coordinates
x1 =  l1*sin(θ)
y1 = -l1*cos(θ)
x2 =  l2*sin(ϕ) + x1
y2 = -l2*cos(ϕ) + y1

# Kinetic energy
T = 0.5*m1*(expand_derivatives(d(x1))^2 + expand_derivatives(d(y1))^2) +
    0.5*m2*(expand_derivatives(d(x2))^2 + expand_derivatives(d(y2))^2)
T = simplify(T)

# Potential energy
V = m1*g*y1 + m2*g*y2

# Lagrangian
L = T - V

# Lagrange equations
E1 = expand_derivatives(d(expand_derivatives(D11(L)))) - expand_derivatives(D1(L))
E2 = expand_derivatives(d(expand_derivatives(D22(L)))) - expand_derivatives(D2(L))

# Solve Lagrange equations for the second order time derivatives
E11, E22 = Symbolics.solve_for([E1 ~ 0, E2 ~ 0], [d(d(θ)), d(d(ϕ))])

# ODE problem
system = [d(d(θ)) ~ E11, d(d(ϕ)) ~ E22]

@named DP = ODESystem(system, t, [θ, ϕ], [m1, m2, l1, l2, g])

DP = structural_simplify(DP)

# Parameters
M = [m1 => 1.0, m2 => 1.0, l1 => 1.0, l2 => 1.0, g => 10.0]

tspan = (0.0, 10.0)

# Initial conditions
X₀ = [d(θ) => 0.0, d(ϕ) => 0.0, θ => pi/2, ϕ => -pi/2]

prob = ODEProblem(DP, X₀, tspan, M, abstol = 1e-7, reltol = 1e-7)
sol = solve(prob, Tsit5());

# Recalculate solution for equidistant time steps for smooth visualization
np = 300
t = range(tspan[1], tspan[2], np)
sl = sol(t);

# Cartesian coordinates
xx1 =  substitute(l1, M).val * sin.(sl[1,:]);
yy1 = -substitute(l1, M).val * cos.(sl[1,:]);
xx2 =  substitute(l2, M).val * sin.(sl[3,:]) .+ xx1;
yy2 = -substitute(l2, M).val * cos.(sl[3,:]) .+ yy1;

# Animation

"""
    snapshot_ft(nt, pts_disp, xl, xr, yb, yt, t, x1, y1, x2, y2)

Display the current positions (x1[nt], y1[nt]) and (x2[nt], y2(nt]), and 
the trajectories' fading tails - the most recent pts_disp points
"""
function snapshot_ft(nt, pts_disp, box, t, x1, y1, x2, y2)
    nt_start = max(1, nt-pts_disp)
    xl, xr, yb, yt = box
    axis("square")
    xlim(xl, xr)
    ylim(yb, yt)
    titl = @sprintf("Double pendulum t = %5.2f .. %5.2f",
                    t[nt_start], t[nt])
    title(titl, family="monospace")
    xlabel(L"$x$")
    ylabel(L"$y$")
    grid(true)
    
    # pivot
    plot(0.0, 0.0, color="black", marker="o", markersize=5.0)

    # two rods
    plot([0.0, x1[nt],x2[nt]], [0.0, y1[nt], y2[nt]], color="black")

    # mass 1
    c1 = "green"
    plot(x1[nt], y1[nt], color=c1, marker="o", markersize=5.0)

    # mass 2
    c2 = "blue"
    plot(x2[nt], y2[nt], color=c2, marker="o", markersize=5.0)

    # (possibly) fading tail trajectory of mass 2
    c2p = parse(RGBA, c2)
    for i = nt_start:nt-1
        transp = 1.0 - ((nt-i-1)/(pts_disp-1))^2
        cc = (c2p.r, c2p.g, c2p.b, transp)
        plot([x2[i], x2[i+1]], [y2[i], y2[i+1]], color=cc, linewidth=0.5)
    end
    
end

# Define the plotting limits
axis_lim = substitute(l1 + l2, M).val*1.2
xl = -axis_lim
xr = -xl
yb = -axis_lim
yt = -yb
box = (xl, xr, yb, yt)

f = 0.1                      # Fraction of displayed points
pts_disp = round(Int, f*np)  # Number of displayed points

fig = figure()
for nt = 1:np
    snapshot_ft(nt, pts_disp, box, t, xx1, yy1, xx2, yy2)
    sleep(0.001)
    if length(get_fignums()) == 0  # figure window has been closed
        break
    end
    clf()
end
close(fig)