using Plots
using DifferentialEquations
import LinearAlgebra

using SparseArrays

theme(:ggplot2)

# Masses schema:
#
#  7 -- 8 -- 9
#  |  X |  X |
#  4 -- 5 -- 6
#  |  X |  X |
#  1 -- 2 -- 3
# 
# where 'X' is a cross connection.
# We use Born-van Karman boundary conditions, i.e.
# a "Pacman world" where the area behaves like it is on a torus.

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.
# Additional terms for diagonal springs are introduced. The spring constants for this are lower.

function coupling_matrix(n::Int, k::Float64; damping::Float64=0., full=false)
    N = n^2
    if full
        A = zeros(4N, 4N)
    else
        A = spzeros(4N, 4N)
    end

    # Link velocity to displacement.
    A[1:2N, 2N+1:end] = 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

        # Next-neighbor 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

        # Next-neighbor out-dimensional coupling: i.e., top neighbor causes x displacement.
        # The coupling constant is just half here. In reality, there would be a nonlinear dependency.
        # A[2N+x, 2t-1] += k/4
        # A[2N+x, 2b-1] += k/4
        # A[2N+x, 2i-1] -= k/2

        # A[2N+y, 2r] += k/4
        # A[2N+y, 2l] += k/4
        # A[2N+y, 2i] -= k/2

        # 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, x] -= 3/2 * k
        A[2N+y, y] -= 3/2 * k

        # Add damping terms.
        A[2N+1:end, 2N+1:end] -= LinearAlgebra.I(2N) * damping
    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*(y-1) + x
        ix, iy = 2i-1, 2i
        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=(-a, n*a), title="Positions $(i)")
end

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

# There will be n x n oscillators.
n = 10

# These are two different initial configurations:
#  All elements on two orthogonal sides are displaced initially.
initial_xy_longitudinal_displacements = vcat(
    [(1, i, 0., .5) for i in 1:n],
    [(i, 2, .5, .0) for i in 1:n])
initial_xy_transversal_displacements = vcat(
    [(1, i, 0.5, 0.0) for i in 1:n],
    [(i, 1, 0.0, 0.5) for i in 1:n])
initial_x_longitudinal_displacements = [(i, 1, .2, 0.) for i in 1:n]
initial_x_transversal_displacements = [(i, 1, 0., 0.2) for i in 1:n]
#  One element is displaced diagonally.
initial_corner_displacements = [(1,1,0.3,0.3), (n, 1, 0.3, -0.3)]
initial_center_displacements = [(div(n,2), div(n,2), .7, .7)]


u0 = initial_state(n, initial_xy_transversal_displacements)

# Essential part: the coupling matrix!
coupling = coupling_matrix(n, 1.; damping=0.0000, full=true)

# Alternatively: Use an eigenvector as initial state, exciting an eigenmode.
ED = LinearAlgebra.eigen(coupling[2n^2+1:end, 1:2n^2]);
u1 = vcat(ED.vectors[:,150] ./ 2, zeros(2n^2));

# ODE formalities
problem = ODEProblem(f2d, u0, (0, 50), coupling);
@time solution = solve(problem);

@gif for (j, u) in enumerate(solution.u)
    plot_state(n, u, j)
end

#plot(solution, vars=[(0, 1), (0, 2)])

# Plot animations for all eigenmodes given an n x n array of oscillators.
# Careful: there are n^2 eigenmodes! (can be restricted to at most `max` eigenmodes being calculated)
function plot_all_eigenmodes(n=5; max=-1)
    # Enumerate all eigenmodes
    coupling = coupling_matrix(n, 1.; full=true)
    ED = LinearAlgebra.eigen(coupling[2n^2+1:end, 1:2n^2])
    @show size(coupling)

    for i in 1:(max < 0 ? n^2 : max)
        @show i
        ev = abs(ED.values[i])
        u0 = vcat(ED.vectors[:,i] ./ 2, zeros(2n^2))

        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
end

# Given a solution, extracts a trajectory for every nth oscillator. Returns two x, y matrices
# suitable for plotting (each column is one timeseries).
function extract_2d_trajectories(solution, nth::Int)::Tuple{Matrix,Matrix}
    N = div(length(solution.u[1]), 4)
    n = round(Int, sqrt(N))
    m = div(N, nth)
    T = length(solution.t)
    x = zeros(T, m)
    y = zeros(T, m)

    j = 1
    for i in 1:nth:N
        x0, y0 = position(n, i, 1.)
        x[:,j] = [solution.u[k][2i-1] for k in 1:T] .+ x0
        y[:,j] = [solution.u[k][2i]   for k in 1:T] .+ y0
        j += 1
    end

    (x, y)
end