import LinearAlgebra
using DifferentialEquations
using Plots
theme(:juno)

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(2n, 2n)
    A[1:n, n+1:2n] = LinearAlgebra.I(n)

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

        if i > 1
            A[b+i, i-1] += k1/masses[i]
        end
        if i < n
            A[b+i, i+1] += k2/masses[i]
        end
        A[b+i, i] -= (k1+k2)/masses[i]
    end
    A
end

# Create an x^3 coupling matrix. l is the l in U = k/2 x^2 - l x^3 + ...
# so that ̈x = k (x_i+1 + x_i-1 - 2 x_i) - 3l ((x_i - x_i-1)² + (x_i - x_i+1)²)
function anharmonic_contribution(n::Int, l::Real)::Matrix
    B = zeros(2n, 2n)

    b = n

    for i in 1:n
        B[b+i, i] = -3l
        if i < n
            B[b+i, i+1] = -3l
        end
    end
    B
end

function initial_state(n::Int, displacements::Vector{Tuple{Int,Float64}})::Vector{Float64}
    u0 = zeros(2n)
    for (di, du) in displacements
        u0[di] = du
    end
    u0
end

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

    u_2 = vcat([0], (u[2:end] - u[1:end-1]).^2)
    du .= A[1] * u + A[2] * u_2
end

N = 7
k1, k2 = 1, 2
mass = 1
initial_displacement = [(1, 1.), (N, 0.)]
timespan = (0, 25)
coupling_H = n_coupled_masses(N, [k1, k1, k1, k1, k1, k1, k1, k1], mass * ones(N))
coupling_AH = anharmonic_contribution(N, -0.1)

problem = ODEProblem(f, initial_state(N, initial_displacement), timespan, (coupling_H, coupling_AH))
solution = solve(problem)

plot(solution, vars=[(0, x) for x in 1:N],
     size=(1800,900), title="Trajectories")