# This file includes the network 

# --------------------------------------- ODENetwork construction and modificatio ------------------------------- # 

"""
    $TYPEDEF 

# Fields 

    $TYPEDFIELDS
"""
mutable struct ODENetwork{T1<:AbstractNode,T2<:AbstractSimpleWeightedGraph,T3<:AbstractMatrix} <: AbstractNetwork
    "Dynamical nodes"
    nodes::Vector{T1}
    "Coupling strength"
    ϵ::Float64
    "Graph explaining topology"
    graph::T2
    "Inner coupling matrix"
    coupling::T3
end
ODENetwork(nodes = Lorenz[], ϵ = 1.0, graph = zeros(length(nodes), length(nodes)), coupling = diagm(ones(dimension(eltype(nodes))))) =
    ODENetwork(nodes, ϵ, SimpleWeightedGraph(graph), coupling)
ODENetwork(::Type{T}, ϵ = 1.0, graph = fill(0.0, 0, 0), coupling = diagm(ones(dimension(T)))) where {T<:AbstractNode} =
    ODENetwork([T() for i = 1:size(graph)[1]], ϵ, graph, coupling)

"""
    $TYPEDEF 

# Fields 

    $TYPEDFIELDS
"""
mutable struct SDENetwork{T1<:AbstractNode,T2<:AbstractSimpleWeightedGraph,T3<:AbstractMatrix,T4<:AbstractMatrix} <: AbstractNetwork
    "Dynamical nodes"
    nodes::Vector{T1}
    "Coupling strength"
    ϵ::Float64
    "Graph explaining topology"
    graph::T2
    "Noise coupling matrix"
    noisemat::T3
    "Inner coupling matrix"
    coupling::T4
end
function SDENetwork(nodes = Lorenz[], ϵ = 1.0, graph = zeros(length(nodes), length(nodes)), noisemat = nothing, coupling = fullstatecoupling(nodes))
    weightedgraph = SimpleWeightedGraph(graph)
    if noisemat === nothing
        noisemat = getnoisemat(weightedgraph)
    end
    SDENetwork(nodes, ϵ, weightedgraph, noisemat, coupling)
end
SDENetwork(::Type{T}, ϵ = 1.0, graph = fill(0.0, 0, 0), noisemat = nothing, coupling = diagm(ones(dimension(T)))) where {T<:AbstractNode} =
    SDENetwork([T() for i = 1:size(graph)[1]], ϵ, graph, noisemat, coupling)


"""
    $TYPEDSIGNATURES

Returns full noise coupling matrix corresponding to graph. 
"""
function getnoisemat(graph)
    numnodes = nv(graph)
    numedges = ne(graph)
    mat = zeros(numnodes, numedges)
    for (k, edge) in enumerate(edges(graph))
        mat[edge.src, k] = 1
        mat[edge.dst, k] = -1
    end
    mat
end

# Returns the innner coupling matrix corresponding to the fully-state coupled network.
fullstatecoupling(nodes) = diagm(ones(dimension(eltype(nodes))))

"""
    $TYPEDSIGNATURES

Adds `node` to `net`
"""
function addnode!(net::AbstractNetwork, node::AbstractNode = eltype(net.nodes)())
    add_vertex!(net.graph)
    push!(net.nodes, node)
    net
end

"""
    $TYPEDSIGNATURES

Remoes the node with index `idx` from `net`.
"""
function remnode!(net::AbstractNetwork, idx::Int = lastindex(net.nodes))
    deleteat!(net.nodes, idx)
    rem_vertex!(net.graph, idx)
    net
end

"""
    $TYPEDSIGNATURES 

Adds and edge in `net` between the nodes `src` and `dst` with weight `weight`
"""
function addedge!(net::AbstractNetwork, src::Int, dst::Int, weight::Real = 1.0, noisestrength::Real = 0.0)
    add_edge!(net.graph, src, dst, weight)
    is_directed(net.graph) && add_edge!(net.graph, dst, src, weight)
    # Update the noise matrix of the network. Tehe 
    newcol = zeros(size(net.noisemat, 1), 1)
    newcol[src] = noisestrength
    newcol[dst] = -noisestrength
    net.noisemat = [net.noisemat newcol]
    net
end

"""
    $TYPEDSIGNATURES

Removes the edge between `src` and `dst` from `net`
"""
function remedge!(net::AbstractNetwork, src::Int, dst::Int)
    rem_edge!(net.graph, src, dst)
    is_directed(net.graph) && rem_edge!(net.graph, dst, src)
    # Find and delete the column of the noise matrix of the network. 
    rows = net.noisemat[[src, dst], :]
    idx = [!iszero(col[1]) && !iszero(col[2]) for col in eachcol(rows)]
    net.noisemat = net.noisemat[:, .!idx]
    net
end

