using Plots
using DifferentialEquations
import LinearAlgebra

function neighbors(n::Int, i::Int)::Tuple
    # Returns neighbors in bottom/top/right/left
    a = (i + n) % n^2
    a = (a==0) ? n^2 : a

    b = (i - n) < 1 ? (n*(n-1)+i) : (i-n)


    c = i + 1 % n^2
    c = (div(c-1, n) > div(i-1, n)) ? (i-n+1) : c

    d = (i - 1) < 1 ? n : (i-1)
    d = (div(d-1, n) < div(i-1, n)) ? (i+n-1) : d

    (a,b,c,d)
end

function position(n::Int, i::Int, a::Float64)::Tuple
    ((i-1)%n, (div(i-1, n)))
end

# Each oscillator has two equations: one for x and one for the y direction.
# Assuming the neighbors L/R (x), T/B (y), we get the following equations of motion for oscillator i:
#
# ̈x_i = k (x_R + x_L - 2 x_i)
# ̈y_i = k (y_T + y_B - 2 y_i)
#
# This assumes constant spring constant k for all springs, obviously.
# In the vector notation, oscillator i has its x component at index 2i-1 and its y component at 2i.

function coupling_matrix(n::Int, k::Float64)::Matrix
    N = n^2
    A = zeros(4N, 4N)

    # Link velocity to displacement.
    A[1:2N, 2N+1:4N] = LinearAlgebra.I(2N)

    for i in 1:N
        (b,t,r,l) = neighbors(n, i)
        # Diagonal neighbors
        (_, _, r1, l1) = neighbors(n, t)
        (_, _, r2, l2) = neighbors(n, b)

        x = 2i-1
        y = 2i

        # TO DO: Introduce non-linear coupling

        A[2N+x, 2r-1] += k
        A[2N+x, 2l-1] += k
        A[2N+x, 2i-1] -= 2k

        A[2N+y, 2t] += k
        A[2N+y, 2b] += k
        A[2N+y, 2i] -= 2k

        # Add diagonal neighbors
        for e in (r1, r2, l1, l2)
            # Diagonal dependence on √2 neighbors.
            A[2N+x, 2*e-1] += k/4
            A[2N+x, 2*e]   += k/8 # Interaction between diag. neighbor y and this element's x is weaker.
            A[2N+y, 2*e-1] += k/8
            A[2N+y, 2*e]   += k/4
        end
        A[2N+x, 2i-1] -= 3/2 * k
        A[2N+y, 2i] -= 3/2 * k
    end
    A
end

function initial_state(n::Int, displacements::Vector{Tuple{Int, Int, Float64, Float64}})::Vector
    u0 = zeros(4n^2)
    N = n^2
    for (x, y, dx, dy) in displacements
        i = n*(x-1) + y
        ix, iy = 2i-1, 2i
        #@show (i, ix, iy)
        u0[ix] += dx
        u0[iy] += dy
    end
    u0
end

function plot_state(n::Int, u::Vector, i::Int=0; a=1.)
    N = n^2
    dis = u[1:2N]
    dxy = reshape(dis, 2, N)

    is = collect(1:N)
    xs = ((is .- 1) .% n)
    ys = div.(is .- 1, n)

    xs += dxy[1,:]
    ys += dxy[2,:]
    scatter(xs .* a, ys .* a, xlim=(-a, (n)*a), ylim=(-2a, (n)*a), title="Positions $(i)")
end

function f2d(du, u, A, t)
    LinearAlgebra.mul!(du, A, u)
end

n = 5

# These are two different initial configurations:

# All elements on two orthogonal sides are displaced initially.
initial_xy_displacements = vcat(
    [(1, i, 0., .2) for i in 1:n],
    [(i, 2, .2, .0) for i in 1:n])
# One element is displaced diagonally.
one_corner_displacement = [(1,1,0.2,0.2)]

u0 = initial_state(n, initial_xy_displacements)

coupling = coupling_matrix(n, 1.)

# Look at a few eigenvibrations
ED = LinearAlgebra.eigen(coupling[2n^2+1:end, 1:2n^2])

# Enumerate all 25 eigenmodes of the 5x5 lattice.
for i in 1:n^2
    @show i
    ev = abs(ED.values[i])
    u0 = vcat(zeros(2n^2), ED.vectors[:,i])

    problem = ODEProblem(f2d, u0, (0, 4*2pi/sqrt(ev)), coupling)
    @time sol = solve(problem);

    a = @animate for (j,u) in enumerate(sol.u)
        plot_state(n, u, j)
    end
    gif(a, "anim_$(n)_ev_$(i).gif", fps=12)
end