import LinearAlgebra
using DifferentialEquations
using Plots
theme(:juno)

# The coupling matrix has 4nx4n elements for a linear chain of two degrees of freedom per oscillator.
# We use Born-von Karman boundary conditions, i.e. the chain is a ring.
function n_coupled_masses(n::Int, ks::Vector, masses::Vector)
    # Prepares coupling matrix.
    if length(ks) != (n+1)
        @show (length(ks), n+1)
        error("Length of ks must be n+1.")
    end
    if length(masses) != n
        @show n
        error("Length of masses must be n.")
    end

    A = zeros(4n, 4n)
    A[1:2n, 2n+1:4n] = LinearAlgebra.I(2n)

    for i = 1:n  # oscillators
        k1, k2 = ks[i], ks[i+1]

        x = 2i-1
        y = 2i

        # Link previous oscillator
        A[2n+x, (x-2) < 1 ? (2n-1) : (x-2)] += k1/masses[i]
        A[2n+y, (y-2) < 1 ? (2n) : (y-2)] += k1/masses[i]

        # Link next oscillator
        A[2n+x, ((x+2)%(2n))] += k2/masses[i]
        A[2n+y, max(((y+2-1)%(2n)+1), 2)] += k2/masses[i]

        A[2n+x, x] -= (k1+k2)/masses[i]
        A[2n+y, y] -= (k1+k2)/masses[i]
    end
    A
end

# The state vector has 4n elements: 2n elements containing x/y positions (alternating, 2i-1 is x),
# and another 2n elements containing x/y velocities (alternating).
# The initial displacements are given as (i, x, y) tuples representing a
# (x,y) displacement of the i'th oscillator.
function initial_state(n::Int, displacements::Vector{Tuple{Int,Float64,Float64}})::Vector{Float64}
    u0 = zeros(4n)
    for (di, dx, dy) in displacements
        u0[2di-1] += dx
        u0[2di] += dy
    end
    u0
end

# The differential equation.
function f(du::Vector{Float64}, u::Vector{Float64}, A, t)
    # u has N positions, N velocities.
    # A gives: du/dt = A u

    du .= A * u
end

# Find eigenmodes of system.
function eigenmodes(n::Int, coupling::Matrix)
    sub = @view coupling[2n+1:end, 1:2n]
    LinearAlgebra.eigen(sub)
end

# Convert an eigenvector E.vectors[:,i] into a system state.
function initial_state_from_eigenmode(em::Vector)::Vector
    vcat(em, zeros(size(em)))
end

# Animate a solution of the system.
function plot_chain_animation(n::Int, solution)
    xs = collect(1:n)
    ys = zeros(n)

    anim = @animate for (i,u) in enumerate(solution.u)
        pos = reshape(u[1:2n], 2, n)
        scatter(xs .+ pos[1,:], pos[2,:], xlim=(-1, N+1), ylim=(-1, 1))
    end
    gif(anim, "anim_$(n).gif")
end

N = 10
k1, k2 = 1, 2
mass = 1
initial_displacement = [(1, .4, .0), (N, -.0, .0)]
timespan = (0, 20)
coupling_H = n_coupled_masses(N, [k1, k2, k1, k2, k1, k2, k1, k2, k1, k2, k1], mass * ones(N))

E = eigenmodes(N, coupling_H)

# Two different initial configurations as choice.
u0 = initial_state(N, initial_displacement)
u1 = initial_state_from_eigenmode(E.vectors[:,14] ./ 3)

problem = ODEProblem(f, u0, timespan, coupling_H)
@time solution = solve(problem)

# 3D trajectory plot
plot(solution, vars=[(0, 2x-1, 2x) for x in 1:N],
     size=(900,450), title="x trajectories")
savefig("trajectory_3d_$(N).svg")

# phase image
plot(solution, vars=[(1,2), (3,4), (5,6), (7,8), (9,10)])
savefig("phasefig_$(N).svg")

plot(solution, vars=[(0, 2i-1) for i in 1:N])

# Animation
plot_chain_animation(N, solution)