# Pkg.add.(["LightGraphs", "MetaGraphs", "GraphPlot", "NamedColors", "RowEchelon", Interact", "SymPy"])

using Interact, RowEchelon, LightGraphs, MetaGraphs, GraphPlot, NamedColors, LinearAlgebra
import SymPy
using SymPy: Sym

struct MyGraph
    g::MetaDiGraph
end
Base.copy(mg::MyGraph) = MyGraph(copy(mg.g))
function MyGraph(edges::Pair{<:Integer,<:Integer}...)
    g = SimpleDiGraphFromIterator(Edge(e) for e in edges)
    MyGraph(MetaDiGraph(g))
end

using Random
function deterministic_spring_layout(g::AbstractGraph; seed::Integer=0, kws...)
    rng = MersenneTwister(seed)
    spring_layout(g, 2 .* rand(rng, nv(g)) .- 1.0, 2 .* rand(rng,nv(g)) .- 1.0; kws...)
end
function Base.show(io::IO, m::MIME"image/svg+xml", mg::MyGraph)
    show(io, m, 
         gplot(mg.g, layout=deterministic_spring_layout,
               nodelabel=map(v -> get(MetaGraphs.props(mg.g, v), :label, v), vertices(mg.g)),
               nodefillc=map(v -> get(MetaGraphs.props(mg.g, v), :color, "gray"), vertices(mg.g)),
               edgelabel=map(ie -> get(MetaGraphs.props(mg.g, ie[2]), :label, ie[1]), enumerate(edges(mg.g))),
               edgestrokec=map(e -> get(MetaGraphs.props(mg.g, e), :color, "lightgray"), edges(mg.g)),
         ))
end

function nodecolors!(g::MyGraph, nodes::AbstractVector{<:Integer}, color::String="red")
    for n in nodes
        set_prop!(g.g, n, :color, color)
    end
    g
end
nodecolors(g, nodes, color) = nodecolors!(copy(g), nodes, color)
edgearr(g, e) = e
edgearr(g, e::AbstractVector{<:Integer}) = collect(edges(g))[e]
edgearr(g, e::AbstractVector{<:Pair}) = Edge.(e)
function edgecolors!(g::MyGraph, edges::AbstractVector, color::String="red")
    for e in edgearr(g.g, edges)
        set_prop!(g.g, e, :color, color)
    end
    g
end
edgecolors(g::MyGraph, edges::AbstractVector, color::String="red") = edgecolors!(copy(g), edges, color)

# A little code so that we can label graph nodes/edges with SymPy expressions.
# convert strings like "v_2 - v_0" from SymPy to nicer Unicode strings like "v₂ - v₀"
subchar(d::Integer) = Char(UInt32('₀')+d)
subchar(c::Char) = subchar(UInt32(c)-UInt32('0'))
subchar(s::String) = replace(s, r"_[0-9]" => s -> subchar(s[2]))
labelstring(s::SymPy.Sym) = subchar(repr("text/plain", s))
labelstring(x) = x

function labels!(g::MyGraph; edges=nothing, nodes=nothing)
    if edges !== nothing
        for (e,E) in zip(MetaGraphs.edges(g.g), edges)
            set_prop!(g.g, e, :label, labelstring(E))
        end
    end
    if nodes !== nothing
        for (n,N) in zip(vertices(g.g), nodes)
            set_prop!(g.g, n, :label, labelstring(N))
        end
    end
    g
end
labels(g::MyGraph; kws...) = labels!(copy(g); kws...)

# generate a random graph with a given average #edges per node
function randgraph(numnodes::Integer, edgespernode::Real)
    p = edgespernode/numnodes # probability of each edge
    e = Vector{Pair{Int,Int}}()
    for i = 1:numnodes, j = 1:numnodes
        if i != j && rand() < p
            push!(e, i=>j)
        end
    end
    return MyGraph(e...)
end

# returns the incidence matrix for g
function incidence(g::MyGraph)
    A = zeros(Int, ne(g.g), nv(g.g))
    for (i,e) in enumerate(edges(g.g))
        A[i,e.src] = -1
        A[i,e.dst] = +1
    end
    return A
end

# Find the loops in g by the simplest "textbook" manner:
# get a basis for the left nullspace incidence matrix.
# We do this via the rref form, rather than nullspace(A'), because
# we want a "nice" basis of ±1 and 0 entries.
function leftnullspace(g::MyGraph)
    A = incidence(g)
    R = rref(Matrix(A'))
    m, n = size(R)
    pivots = Int[]
    for i = 1:m
        j = findfirst(!iszero, R[i,:])
        j !== nothing && push!(pivots, j)
    end
    r = length(pivots) # rank
    free = Int[j for j=1:n if j ∉ pivots]
    N = zeros(Int, n, n-r)
    k = 0
    for (k,j) in enumerate(free)
        N[pivots, k] = -R[1:r, j]
        N[j, k] = 1
    end
    return N
end

# color the edges of a spanning tree of g, by the textbook
# method of finding the pivot rows of the incidence matrix
function pivotrows(g::MyGraph)
    A = incidence(g)
    R = rref(Matrix(A'))
    m, n = size(R)
    pivots = Int[]
    for i = 1:m
        j = findfirst(!iszero, R[i,:])
        j !== nothing && push!(pivots, j)
    end
    return pivots
end
colortree(g::MyGraph, color::String="red") = edgecolors(g, pivotrows(g), color)
tree(g::MyGraph) = MyGraph(MetaDiGraph(SimpleDiGraphFromIterator(edgearr(g.g, pivotrows(g)))))

g = MyGraph(1=>4, 4=>5, 5=>6, 6=>3, 3=>2, 2=>1, 2=>6, 4=>6)

A = incidence(g)

A[[4,3,8],:]

A[4,:]' + A[3,:]' + A[8,:]'

[0 0 1 1 0 0 0 1] * A

Matrix{Int}(rref(Matrix(A')))

rank(A)

N = leftnullspace(g)

A' * N

colorloop(g::MyGraph, n::Vector) = edgecolors!(edgecolors(g, findall(n .> 0), "red"), findall(n .< 0), "blue")
function animloops(g::MyGraph)
    L = leftnullspace(g)
    @manipulate for loop in 1:size(L,2)
        colorloop(g, L[:,loop])
    end
end

animloops(g)

colorloop(g, N[:,2] + N[:,3]) # add two loops to make another loop

gbig = randgraph(15, 2)
animloops(gbig)

colortree(g)

tree(g)

colortree(gbig)

tree(gbig)

tree(randgraph(50, 8))

labels(g, edges=[Sym("i_$i") for i = 1:size(A,1)], nodes=[Sym("v_$i") for i = 1:size(A,2)])

v = [Sym("v_$i") for i = 1:size(A,2)]

A

A * v

labels(g, edges=A*v, nodes=v)

d = [Sym("d_$j") for j = 1:size(A,1)]
labels(g, edges=d, nodes=v)

N' * d

Y = diagm(0=>[Sym("Y_$i") for i = 1:size(A,1)])

Y*d

Y*A*v

i = [Sym("i_$j") for j=1:size(A,1)]
A'*i

labels(g, edges=i, nodes=A'*i)

A'

A' * Y * A

nullspace(A)

twodigits(x) = round(x, digits=2)

@manipulate for Y₁=0.1:0.1:10,
                Y₂=0.1:0.1:10,
                Y₃=0.1:0.1:10,
                Y₄=0.1:0.1:10,
                Y₅=0.1:0.1:10,
                Y₆=0.1:0.1:10,
                Y₇=0.1:0.1:10,
                Y₈=0.1:0.1:10
    s = [1,-1,0,0,0,0]
    Y = diagm(0=>[Y₁,Y₂,Y₃,Y₄,Y₅,Y₆,Y₇,Y₈])
    nodecolors!(labels(g, edges=[subchar("i_$j = $i") for (j,i) in enumerate(twodigits.(Y*A*(pinv(A'*Y*A) * s)))]),
                [1,2])
end