module Graf

import Base:getindex, ==, show, iterate, length, rand, filter
import Base.Iterators
using Lazy

export AttribDict, InnerDict, AdjDict, OuterDict, GenericGraf, adj,
add_arc!, add_arcs_from!, add_node!, add_nodes_from!,
rem_arc!, rem_arcs_from!, rem_node!, rem_nodes_from!,
NodeIter, nodes, nn,
ArcIter, arcs, na,
degree, has_node, has_arc

const AttribDict = Dict{Symbol,Any} # attribute
const InnerDict = Dict{Any,AttribDict} # keys of parallel arc
const AdjDict = Dict{Any,InnerDict} # neighbors of node
const OuterDict = Dict{Any,AdjDict} # node container
const NestedDict = [OuterDict,AdjDict,InnerDict,AttribDict]
"""
    GenericGraf
a structure representing a graph.
"""
mutable struct GenericGraf
    na::Integer
    adj::OuterDict
end

function GenericGraf(n::Integer=0)
    adj = OuterDict(i => InnerDict() for i in 1:n)
    return GenericGraf(0, adj)
end

# accessors
na(g::GenericGraf) = g.na
adj(g::GenericGraf) = g.adj

# convenience
show(io::IO, g::GenericGraf) = print(io, "{$(nn(g)),$(na(g))} GenericGraf")
==(g::GenericGraf, h::GenericGraf) = adj(g) == adj(h)

"""
    getindex(g, inds)
an interface to access the data on the GenericGraf using indices such as:
    - G[u,v], G[u,v,k], and G[u,v,k,:attrib]
    - G[u], G[u][v], G[u][v][k], and G[u][v][k][:attrib]
the setter for this type is a set of functions called builders.
there are no straight access such as `setindex!` to directly edit graph.
"""
function getindex(g::GenericGraf, inds::Tuple)
    linds = length(inds)
    if linds <= length(NestedDict)
        current = g.adj
        for i in inds
            current = getindex(current, i)
        end
    else
        msg = KeyError(inds)
        throw(msg)
    end
    return current
end

getindex(g::GenericGraf, c::Colon) = adj(g)
getindex(g::GenericGraf, k1, k2, ks...) = getindex(g, tuple(k1, k2, ks...))

# builders
"""
    add_arc!(g,u,v,key)
add an arc (u,v,key) to graph g. the default for key is nothing.
node existence will be first checked then the keys and then the addition.
"""
function add_arc!(g::GenericGraf, u, v, key=nothing)
    if !haskey(g.adj, u)
        g.adj[u] = AdjDict()
    end
    if !haskey(g.adj, v)
        g.adj[v] = AdjDict()
    end
    if isnothing(key)
        key = new_arc_key(g, u, v)
    end
    if haskey(g.adj[u], v)
        innerdict = g.adj[u][v]
        attribdict = get(innerdict, key, AttribDict())
        innerdict[key] = attribdict
        g.na += 1
    else
        attribdict = AttribDict()
        innerdict = InnerDict()
        innerdict[key] = attribdict
        g.adj[u][v] = innerdict
        g.na += 1
    end
    return key
end

add_arc!(g::GenericGraf, t::Tuple) = add_arc!(g, t...)

"""
    add_arcs_from!(g, arcbunch)
add arcs from an iterable input. abunch is the container of the arc.
each arc can be one of:
    - 2-tuples (u,v)
    - 3-tuples (u,v,k)
default keys will be generated by `new_arc_key()`.
adding an already existing arc will update its list of keys.
all non-compliant arc will be silently not added.
"""
function add_arcs_from!(g::GenericGraf, arcbunch)
    keylist = []
    @inbounds for a in arcbunch
        la = length(a) # arc length
        if la == 2 # 2-tuple
            u, v = a; key = nothing
            key = add_arc!(g, u, v, key) # moneyshot
            push!(keylist, key)
        elseif la == 3 # 3-tuple
            u, v, key = a
            key = add_arc!(g, u, v, key) # moneyshot
            push!(keylist, key)
        end
    end
    return keylist
end

"""
    new_arc_key(g,u,v)
new key generator for arc (u,v). the key will be the lowest integer not used.
the function is internally used for adding arcs so not exported and no check of u,v.
"""
function new_arc_key(g::GenericGraf, u, v)
    keydict = try
        g.adj[u][v] # no arc yet
    catch KeyError
        return 1 # lowest key is 1
    end
    key = length(keydict)
    while key in keys(keydict)
        key += 1
    end # find lowest integer not used
    return key
end

"""
    add_node!(g, u)
add a node u to graph g. u is a not nothing object.
silently doesn't add anything if object is already a node.
"""
function add_node!(g::GenericGraf, u)
    if !haskey(g.adj, u) # operation only for nonexisting
        g.adj[u] = Dict() # moneyshot
        return u # something done
    end
end

"""
    add_nodes_from!(g, nbunch)
add nodes from nbunch to g. nbunch is a container of nodes to add.
"""
function add_nodes_from!(g::GenericGraf, nbunch)
    added = []
    for n in nbunch
        add = add_node!(g, n)
        push!(added, add)
    end
    return added
end

"""
    rem_arc!(g, u, v, key)
remove an arc u,v,key from the graph. u and v are the nodes.
key is used to distinguish multiple arcs between pair of nodes.
"""
function rem_arc!(g::GenericGraf, u, v, key=nothing)
    d = try
        g.adj[u][v]
    catch KeyError
        return false # no key uv on graph
    end
    if isnothing(key)
        pop!(d) # remove element at the back
    else
        if haskey(d, key)
            delete!(d, key) # remove the item key
        else # doesn't have key
            return false
        end
    end
    if length(d) == 0 # no more arc
        delete!(g.adj[u], v) # delete entry uv
    end

    g.na -= 1
    return true # only success operation reach bottom
end

rem_arc!(g::GenericGraf, t::Tuple) = rem_arc!(g, t...)

"""
    rem_arcs_from!(g, abunch)
remove several arcs in abunch from g. the arcs in abunch can be one of:
    - 2-tuples (u,v)
    - 3-tuples (u,v,k)
nonexisting arcs in abunch will silently not be removed.
"""
function rem_arcs_from!(g::GenericGraf, abunch)
    removed = 0
    for a in abunch
        if length(a) >= 2
            remove = try
                rem_arc!(g, a...) # if a doesn't comply,
            catch end # just pass.
            if remove
                removed += 1
            end
        end
    end
    return removed
end

"""
    rem_node!(g, n)
Removes the node n and all adjacent arcs.
Attempting to remove a non-existent node will raise an exception.
"""
function rem_node!(g::GenericGraf, n)
    succs = try
        g.adj[n]
    catch KeyError
        return false
    end
    for u in keys(g.adj)
        if haskey(g.adj[u], n)
            g.na -= length(g.adj[u][n]) # adjust number of arc
            delete!(g.adj[u], n) # delete arc n-u on pred
        end
    end
    g.na -= length(g.adj[n]) # adjust number of arc
    delete!(g.adj, n)

    return true # successful operation returns true
end

"""
    rem_nodes!(g, nbunch)
remove several nodes at once. nbunch are the containers of nodes to remove.
If a node in the container is not in the graph it is silently ignored.
"""
function rem_nodes_from!(g::GenericGraf, nbunch)
    removed = 0
    for n in nbunch
        remove = try
            rem_node!(g, n) # if any key error,
        catch end # just pass.
        if remove
            removed += 1
        end
    end
    return removed
end

function degree(g::GenericGraf, nbunch=nodes(g))
    if length(nbunch) > 1
        res = Dict{Any,Integer}()
        for n in nbunch
            deg = sum(Int64[length(k) for k in values(g.adj[n])])
            push!(res, n => deg)
        end
        return res
    else
        deg = sum(Int64[length(k) for k in values(g.adj[nbunch])])
        return deg
    end
end

has_node(g::GenericGraf, n) = haskey(g.adj, n)

function has_arc(g::GenericGraf, u, v, key=nothing)
    if haskey(g.adj, u)
        if haskey(g.adj[u], v)
            if isnothing(key)
                return true
            else
                return haskey(g.adj[u][v], key)
            end
        else
            return false
        end
    else
        return false
    end
end

# NodeIter
struct NodeIter
    g::GenericGraf
end

# constructor
function iterate(nit::NodeIter, state=1)
    u = state
    n = nit.g.adj

    isempty(n) && return nothing

    res = iterate(n, u)
    if isnothing(res)
        return nothing
    else
        return res[1].first, res[2]
    end
end

# interface
length(nit::NodeIter) = length(nit.g.adj)
Base.rand(x::NodeIter) = Base.rand(collect(x))
Base.rand(x::NodeIter, dims::Integer) = Base.rand(collect(x), dims)
filter(f,nit::NodeIter) = Iterators.filter(f, nit)

# convenience
show(io::IO, nit::NodeIter) = print(io, "NodeIter $(length(nit))")
nodes(g::GenericGraf) = NodeIter(g)
nn(g::GenericGraf) = length(nodes(g))

# ArcIter
struct ArcIter
    g::GenericGraf
end

# constructor
function iterate(ait::ArcIter, state=(1, 1, 1))
    u, v, key = state
    n = ait.g.adj

    isempty(n) && return nothing

    while true
        @label new_u; nbrs_u = iterate(n, u)
        if iterate(nbrs_u[1].second, v) === nothing
            if iterate(n, nbrs_u[2]) === nothing
                return nothing
            else
                u = nbrs_u[2]; v = 1; key = 1
                @goto new_u
            end
        else
            @label new_v; keys_uv = iterate(nbrs_u[1].second, v)
            if iterate(keys_uv[1].second, key) === nothing
                if iterate(nbrs_u[1].second, keys_uv[2]) === nothing
                    if iterate(n, nbrs_u[2]) === nothing
                        return nothing
                    else
                        u = nbrs_u[2]; v = 1; key = 1
                        @goto new_u
                    end
                else
                    v = keys_uv[2]; key = 1
                    @goto new_v
                end
            else
                lastkey = iterate(keys_uv[1].second, key)
                res = nbrs_u[1].first, keys_uv[1].first, lastkey[1].first
                state = u, v, lastkey[2]
                return res, state
            end
        end
    end
end

# interface
length(ait::ArcIter) = na(ait.g)
Base.rand(x::ArcIter) = Base.rand(collect(x))
Base.rand(x::ArcIter, dims::Integer) = Base.rand(collect(x), dims)
filter(f,ait::ArcIter) = Iterators.filter(f, ait)

# convenience
show(io::IO, ait::ArcIter) = print(io, "ArcIter $(length(ait))")
arcs(g::GenericGraf) = ArcIter(g)
function na(g::GenericGraf, u, v)
    arcdata = try
        g.adj[u][v]
    catch
        return 0
    end
    return length(arcdata)
end

end
