# Spring pendulum.
#
# A pendulum made of a spring with a mass m on the end. The spring is arranged
# to lie in a straight line (which we can arrange by, say, wrapping the spring
# around a rigid massless rod). The equilibrium length of the spring is l, the
# elastic cobstant is k. Assume that the motion takes place in a vertical plane.
# The acceleration of gravity is g.
# Let the spring length be r(t), and the angle between the spring and the vertical
# be θ(t). Write down the equations of motion for r(t), θ(t) and solve them (numerically).
#

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

using Symbolics, ModelingToolkit
using Latexify
using OrdinaryDiffEq
using PyPlot
using Colors
using Printf

@variables t θ(t) r(t) 
@parameters m, k, l, g

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

# Kinetic and potential energies
T = 0.5*m*((d(r))^2 + r^2*(d(θ))^2)
V = -m*g*r*cos(θ) + 0.5*k*(r - l)^2

# 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(r))])

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

@named SP = ODESystem(system, t, [θ, r], [m, k, l, g])

SP = structural_simplify(SP)

# Parameters
M = [m => 1.0, k => 1.5, l => 1.0, g => 10.0]

tspan = (0.0, 40.0)

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

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

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

# Cartesian coordinates
x =  sl[1,:] .* sin.(sl[3,:]);
y = -sl[1,:] .* cos.(sl[3,:]);

# Animation

"""
    snapshot_ft(nt, pts_disp, box, t, x1, y1)

Display the current positions (x1[nt], y1[nt]) and 
the trajectory's fading tail - the most recent pts_disp points
"""
function snapshot_ft(nt, pts_disp, box, t, x1, y1)
    nt_start = max(1, nt-pts_disp)
    xl, xr, yb, yt = box
    axis("square")
    xlim(xl, xr)
    ylim(yb, yt)
    titl = @sprintf("Spring 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="x", markersize=5.0)

    # the spring
    plot([0.0, x1[nt]], [0.0, y1[nt]], color="black", linewidth=0.5)

    # mass 
    c1 = "red"
    plot(x1[nt], y1[nt], color=c1, marker="o", markersize=3.0)

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

# Define the plotting limits
axis_lim = substitute(m*g/k, M).val*3
xl = -axis_lim*0.75
xr = -xl
yb = -axis_lim
yt = -yb/2
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, x, y)
    sleep(0.001)
    if length(get_fignums()) == 0  # figure window has been externally closed
        break
    end
    clf()
end
close(fig)