# Paralel network simulation 

# Add process 
using Distributed
nprocs() == Sys.CPU_THREADS || addprocs(Sys.CPU_THREADS - nprocs() + 1)

# Load packages
@everywhere begin
    using Pkg
    Pkg.activate(joinpath(Pkg.envdir(), "dev-env"))
end

@everywhere begin
    import Cairo, Fontconfig
    using GraphPlot
    using Graphs
    using SimpleWeightedGraphs
    using LinearAlgebra
    using DynamicalNetworks
    using DifferentialEquations
end

@everywhere begin
    σr = collect(0.0:0.1:1.0)
    function prob_func(prob, i, repeat)
        @show σr[i]
        prob.p[1].noisemat .*= σr[i]
        prob
    end

    # Settings 
    nbits = 10
    tbit = 50.0
    bits = rand(Bool, nbits)
    ti = 0.0
    tf = nbits * tbit
    tswitch = ti:tbit:tf-tbit
    srcnode1, dstnode1 = 5, 8     # Control link nodes 
    high = 100.0            # Control coupling strength 
    low = eps()             # Control coupling strength
    condition(u, t, integrator) = integrator.t ∈ tswitch
    function action!(integrator)
        net, bits = integrator.sol.prob.p
        if popfirst!(bits)
            add_edge!(net.graph, srcnode1, dstnode1, high)
        else
            add_edge!(net.graph, srcnode1, dstnode1, low)
        end
    end
end

# Topology 
n = 8
graph = cycle_graph(n) |> SimpleWeightedGraph
add_edge!(graph, 2, 7)
add_edge!(graph, 3, 6)
ϵ = 20.0

# Network 
σ = σr[1]
H = [
    1 0 0 0
    -1 0 0 0
    0 1 0 0
    0 -1 0 0
    0 0 1 0
    0 0 -1 0
    0 0 0 1
    0 0 0 -1
] * σ
P = diagm([1.0, 0.0, 0.0])
net = SDENetwork(Lorenz, ϵ, graph, H, P)
callback = DiscreteCallback(condition, action!)

# Simulation 
prob = netproblem(net, (ti, tf), copy(bits), noise_rate_prototype = get_noise_rate_prototype(net), callback = callback, tstops = tswitch)
ensembleprob = EnsembleProblem(prob, prob_func = prob_func)
sim = solve(ensembleprob, SRA(), EnsembleDistributed(), trajectories = length(σr))

# Plots 
using PyPlot
pygui(true)
close("all")
getwave(x, i, d) = getindex.(x, (i - 1) * d + 1)
figure()
d = 3
subplot(221)
for (i, sol) in enumerate(sim[[10, 1]])
    plot(sol.t, abs.(getwave(sol.u, 1, d) - getwave(sol.u, 8, d)), label = "$i")
    axvline.(collect(ti:tbit:tf), color = "red", linestyle = "dashed")
    legend()
end
