import LinearAlgebra
using DifferentialEquations
using Plots
theme(:ggplot2)

# 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; fixed_ends=false)
    # 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
        prevx = (x-2) < 1 ? (2n-1) : (x-2)
        prevy = (y-2) < 1 ? (2n) : (y-2)
        if !fixed_ends || ((x-2) >= 1)
            A[2n+x, prevx] += k1/masses[i]
        end
        if !fixed_ends || ((y-2) >= 1)
            A[2n+y, prevy] += k1/masses[i]
        end


        # Link next oscillator
        nextx = ((x+2)%(2n))
        nexty = max(((y+2-1)%(2n)+1), 2)
        if !fixed_ends || (nextx == x+2)
            A[2n+x, nextx] += k2/masses[i]
        end
        if !fixed_ends || (nexty == y+2)
            A[2n+y, nexty] += k2/masses[i]
        end

        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)::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

function plot_website_chain_animation(n::Int, solution)
    xs = .4 .* collect(1:n)
    ys = zeros(n)
    T = length(solution.u)

    trajectories_x, trajectories_y = zeros(n, T), zeros(n, T)

    anim = @animate for (i,u) in enumerate(solution.u)
        xpos = u[1:2:2n]
        ypos = u[2:2:2n]
        trajectories_x[:, i] = xs .+ xpos
        trajectories_y[:, i] = ypos .+ ones(n)
        plot(vcat([0], xs .+ xpos, [(n+1)*.4]), ones(n+2) .+ vcat([0], ypos, [0]), xlim=(0, .4*(n+1)), ylim=(.8, 1.2), size=(800, 200), legend=false, markersize=8, markershape=:circle)
        plot!(trajectories_x[:, 1:i]', trajectories_y[:, 1:i]')
    end
    gif(anim, "anim_$(n).gif", fps=12)
end

function initial_wave_displacement(N::Int, q::Float64, ampl::Float64=.25, a::Float64=1.; trans=false)::Vector{Tuple}
    N = div(N, 2)
    ampls = [ampl * sin(q * a * i) for i in 1:N]
    [(i, trans ? 0. : ampls[i], trans ? ampls[i] : 0.) for i in 1:N]
end

N = 20
k1, k2 = 1., 2.
mass = 1.
initial_displacement = [(1, 0.15, 0.15), (N, 0., 0.)]
initial_displacement_2 = [(1, .5, 0.)]
initial_wave_4 = initial_wave_displacement(N, 2pi/(N), .1; trans=true)
timespan = (0, 90)
coupling_H = n_coupled_masses(N, k1 * ones(N+1), ones(N) * mass; fixed_ends=true)

E = eigenmodes(N, coupling_H)

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

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)