# Parallel network simulation 

using DifferentialEquations
using DynamicalNetworks
using Graphs
using SimpleWeightedGraphs
using LinearAlgebra
using Distributed
using Logging
using LoggingExtras
using JLD2
using Dates

# Includes 
include("getclargs.jl")
include("setlogger.jl")

# Settings
clargs = getclargs()
mode = clargs["mode"]
numexp = clargs["ntrials"]
minsnr = clargs["minsnr"]
stepsnr = clargs["stepsnr"]
maxsnr = clargs["maxsnr"]
nbits = clargs["nbits"]
tbit = clargs["tbit"]
ncores = clargs["ncores"]
high = clargs["high"]
low = clargs["low"]
ϵ = clargs["coupling-strength"]
srcnode = clargs["srcnode"]
dstnode = clargs["dstnode"]
simdir = clargs["simdir"]
simprefix = clargs["simprefix"]
loglevel = clargs["loglevel"]
savenoise = clargs["savenoise"]
P = diagm([1.0, 0.0, 0.0])
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
]

# Construct simulation path 
simpath = joinpath(simdir, simprefix * string(Dates.format(now(), "yy-mm-dd-HH-MM-SS")))
isdir(simpath) || mkpath(simpath)

# Set the logger
logger = setlogger(simpath, loglevel)

# Start simulation
tinit = time()
@info "Started simulation with settings" clargs

# Initializations 
@info "Initializing simulation..."
ti = 0.0
tf = nbits * tbit
params = (
    mode = mode,
    nbits = nbits,
    tbit = tbit,
    bits = rand(Bool, nbits),
    high = high,
    low = low,
    ϵ = ϵ,
    tswitch = ti:tbit:tf-tbit,
    srcnode = srcnode,
    dstnode = dstnode,
    ctrlidx = findall(!iszero, H),
    snrrange = minsnr:stepsnr:maxsnr,
    noiserange = vcat(fill((minsnr:stepsnr:maxsnr)', numexp)...),
    ncores = ncores,
    simdir = simdir,
    simprefix = simprefix
)
@info "Done."

# Network 
@info "Constructing network..."
n = 8
graph = cycle_graph(n) |> SimpleWeightedGraph
add_edge!(graph, 2, 7)
add_edge!(graph, 3, 6)
net = SDENetwork(Lorenz, ϵ, graph, H, P)
@info "Done."

# Loads 
if params.mode == "distributed"
    @info "Loading processess..."
    using Distributed
    na = length(Sys.cpu_info()) - 1 - nprocs()
    params.ncores ≤ na ? addprocs(params.ncores) : addprocs(na)
    @everywhere begin
        using Pkg
        dev_env_path = joinpath(Pkg.envdir(), "dev-env")
        dirname(Pkg.project().path) == dev_env_path || Pkg.activate(dev_env_path)
        include("load.jl")
    end
    @info "Done."
else
    include("load.jl")
end

# Get base problem 
@info "Constructing base problem..."
prob = netproblem(net, (ti, tf), params, noise_rate_prototype = get_noise_rate_prototype(net),
    callback = DiscreteCallback(condition, action!), tstops = params.tswitch)
@info "Done."

# Get ensemble problem 
@info "Getting ensemble problem..."
ensembleprob = EnsembleProblem(prob, prob_func = prob_func)
@info "Done."

# Simulate ensemle problem 
@info "Solving ensemble problems..."
if params.mode == "distributed"
    ensemblealg = EnsembleDistributed()
elseif params.mode == "threads"
    ensemblealg = EnsembleThreads()
else
    ensemblealg = EnsembleSerial()
end
sim = solve(ensembleprob, SRA(), ensemblealg, trajectories = length(params.noiserange))
@info "Done."

# Save simulation 
@info "Writing simulation..."
cartidx = CartesianIndices(params.noiserange)
for (i, sol) in enumerate(sim)
    trial, snr = cartidx[i].I
    snrpath = joinpath(simpath, "$snr-dB")
    isdir(snrpath) || mkpath(snrpath)
    trialpath = joinpath(snrpath, "Traial-$trial")
    isdir(trialpath) || mkpath(trialpath)
    @save joinpath(trialpath, "sol.jld2") sol
end
@save joinpath(simpath, "bits.jld2") params.bits
@info "Done."
