# This example simulates cluster synchronization using pinning control 

import Cairo, Fontconfig
using LinearAlgebra
using DifferentialEquations
using PyPlot
using BlockArrays
using Graphs, GraphPlot
pygui(true)
close("all")

# Helper functions 
getwave(x, i, d) = getindex.(x, (i - 1) * d + 1)

# Define node dynamics 
h(x, a = -1 / 7, b = 2 / 7) = b * x + 1 / 2 * (a - b) * (abs(x + 1) - abs(x - 1))
function chua(dx, x, α = 9, β = 100 / 7, h = h)
    dx[1] = α * (x[2] - h(x[1]))
    dx[2] = x[1] - x[2] + x[3]
    dx[3] = -β * x[2]
end
d = 3

function lorenz(dx, x, σ = 10, β = 8 / 3, ρ = 35)
    dx[1] = σ * (x[2] - x[1])
    dx[2] = x[1] * (ρ - x[3]) - x[2]
    dx[3] = x[1] * x[2] - β * x[3]
end

# Define clusters
cls = [1:3, 4:6]
ks = length.(cls)
n = sum(ks)
l = length(ks)

# Define passivity index.
f = lorenz
Δ = diagm(ones(d)) * 10
P = diagm(ones(d))

# Define topology of the network 
μ(A) = 1 / 2 * maximum(size(A)) * maximum(abs.(A))
Φ = BlockArray([
        -1.0 1 0 0 0 0
        1 -2 1 0 -1 1
        0 1 -1 0 1 -1
        0 0 0 -1 1 0
        0 -1 1 1 -2 1
        0 1 -1 0 1 -1
    ], [3, 3], [3, 3])

μs = zeros(l * l - l)
k = 1
for i = 1:l, j = 1:l
    global k
    if j != i
        μs[k] = μ(getblock(Φ, i, j))
        k += 1
    end
end
μm = maximum(μs)
δm = maximum(diag(Δ))

# Define initial constants 
γs = [1, 1] * 200
Γiis = [diagm([zeros(k - 1); γ]) for (k, γ) in zip(ks, γs)]
Φiis = [getblock(Φ, i, i) for i = 1:l]
λs = [-maximum(eigvals(Φii - Γii)) for (Φii, Γii) in zip(Φiis, Γiis)]
βis = (δm + 2 * (l - 1) * μm) ./ λs
β = maximum(βis) * 1  # The condition is not statisfied.
αs = β * γs
E = copy(Φ)
for i = 1:l
    setblock!(E, getblock(E, i, i) * β, i, i)
end
E = [zeros(l, l) zeros(l, n); zeros(n, l) E]
H1 = zeros(n, l)
for (k, (j, α)) in zip(cumsum(ks), enumerate(αs))
    H1[k, j] = -α
end
H2 = zeros(n, n)
for (k, α) in zip(cumsum(ks), αs)
    H2[k, k] = α
end
H = [zeros(l, l + n); [H1 H2]]

# Define coupled network dynamics
function statefunc(dx, x, u, t, f = f, E = E, H = H, P = P)
    n = size(E, 1)
    d = size(P, 1)
    # Individual evolution
    for i = 1:n
        f(view(dx, (i-1)*d+1:i*d), view(x, (i-1)*d+1:i*d))
    end
    # Add coupling and control inputs
    dx .+= kron(E - H, P) * x
end

# Simulation network
tspan = (0.0, 100.0)
x0 = rand((n + l) * d) * 1e-3
sol = solve(ODEProblem(statefunc, x0, tspan), saveat = 0.001)
t = sol.t
x = sol.u

# Plot the results
ni = 1
nf = 50
graph = SimpleDiGraph(E - H)
gplot(graph, nodelabel = 1:nv(graph), layout = circular_layout) |> display
subplot(5, 2, 1)
plot(t, getindex.(x, 1))
subplot(5, 2, 2)
plot(getindex.(x, 1), getindex.(x, 2))
subplot(5, 2, 3)
plot(t, abs.(getwave(x, 1, d) - getwave(x, 3, d)))
subplot(5, 2, 4)
plot(t, abs.(getwave(x, 1, d) - getwave(x, 4, d)))
subplot(5, 2, 5)
plot(t, abs.(getwave(x, 1, d) - getwave(x, 5, d)))
subplot(5, 2, 6)
plot(t, abs.(getwave(x, 2, d) - getwave(x, 7, d)))
subplot(5, 2, 7)
plot(t, abs.(getwave(x, 2, d) - getwave(x, 7, d)))
subplot(5, 2, 8)
plot(t, abs.(getwave(x, 2, d) - getwave(x, 8, d)))
subplot(5, 2, 9)
plot(t, abs.(getwave(x, 1, d) - getwave(x, 2, d)))
subplot(5, 2, 10)
plot(t, abs.(getwave(x, 4, d) - getwave(x, 7, d)))
tight_layout()
