using DifferentialEquations
using Plots
theme(:ggplot2)

function cho(du, u, p, t)
    du[1] = u[3]
    du[2] = u[4]
    du[3] = -p[:k]/p[:m] * (2*u[1] - u[2])
    du[4] = -p[:k]/p[:m] * (-u[1] + 2*u[2])
end

function plot_chain_animation(n::Int, solution)
    xs = 2 .* collect(1:n)
    ys = zeros(n)

    anim = @animate for (i,u) in enumerate(solution.u)
        pos = u[1:n]
        plot(vcat([0], xs .+ pos, [6]), ones(n+2), xlim=(0, 2*(n+1)), ylim=(.9, 1.1), size=(800, 200), title="t = $(solution.t[i])", legend=false, markersize=8, markershape=:circle)
    end
    gif(anim, "anim_$(n).gif", fps=12)
end

p = Dict(:k => 1., :m => 1.)
u0 = [1/sqrt(2), 1/sqrt(2), 0, 0]

problem = ODEProblem(cho, u0, (0, 3 * 2pi/sqrt(1)), p)
@time solution1 = solve(problem)

plot(
    plot(solution1, vars=[(0, 1), (0, 2)], xlabel="t", ylabel="x", label=["\$1: x_1(t)\$" "\$1: x_2(t)\$"]),
    plot(solution2, vars=[(0, 1), (0, 2)], label=["\$2: x_1(t)\$" "\$2: x_2(t)\$"]),
    layout=(2,1))
#plot(solution2, vars=[(0, 1), (0, 2)], label=["\$2: x_1(t)\$" "\$2: x_2(t)\$"])