graph.go

package gorgonia

import (
	"bytes"
	"fmt"
	"unsafe"

	"github.com/awalterschulze/gographviz"
	"github.com/gonum/graph"
)

// a dag is just a holding structure for a directed acyclic graph (of expressions). It's based on the
// reference implementation of gonum's simple example.
//
// The main difference is that the dag is add-only, and there is no removal of nodes
type ExprGraph struct {
	name string

	all Nodes

	byHash map[uint32]*Node
	evac   map[uint32]Nodes
	to     map[*Node]Nodes

	leaves    Nodes
	constants Nodes
	roots     Nodes
}

type graphconopt func(g *ExprGraph)

// WithGraphName is a ExprGraph construction option that provides a name.
func WithGraphName(name string) graphconopt {
	f := func(g *ExprGraph) {
		g.name = name
	}
	return f
}

// NewGraph creates a new graph. Duh
func NewGraph(opts ...graphconopt) *ExprGraph {
	g := &ExprGraph{
		byHash: make(map[uint32]*Node),
		evac:   make(map[uint32]Nodes),
		to:     make(map[*Node]Nodes),

		leaves:    make(Nodes, 0),
		constants: make(Nodes, 0),
	}

	for _, opt := range opts {
		opt(g)
	}

	return g
}

// AddNode adds n to the graph. It panics if the added node ID matches an existing node ID.
func (g *ExprGraph) AddNode(n *Node) (retVal *Node) {
	defer func() {
		if _, ok := g.to[retVal]; !ok {
			g.to[retVal] = nil
		}
	}()

	hash := n.Hashcode()
	if existing, ok := g.byHash[hash]; ok {
		if existing == nil {
			// this means that there has been previous collisions
			// so look at evac map
			for _, e := range g.evac[hash] {
				if nodeEq(n, e) {
					return e
				}
			}
			g.evac[hash] = append(g.evac[hash], n)
			g.addToAll(n)
			incrCC() // collision counter
			return n
		}

		if !nodeEq(n, existing) {
			g.evac[hash] = Nodes{existing, n}
			g.byHash[hash] = nil // to signal that it's collided
			g.addToAll(n)
			incrCC()
			return n
		}
		incrEC() // expected collision (they're the same node!)
		return existing
	}

	if n.isConstant() {
		n = n.clone()
		n.g = g
	}

	g.addToAll(n)
	g.byHash[hash] = n
	return n
}

func (g *ExprGraph) addToAll(n *Node) {
	if n == nil {
		panic("HELP! trying to add nil")
	}
	g.all = append(g.all, n)
}

// RemoveNode removes n from the graph, as well as any edges attached to it. If the node
// is not in the graph it is a no-op.
func (g *ExprGraph) RemoveNode(node graph.Node) {
	n := node.(*Node)
	hash := n.Hashcode()

	delete(g.byHash, hash)
	delete(g.to, n)
	g.evac[hash] = g.evac[hash].remove(n)
	g.all = g.all.remove(n)
}

// SetEdge adds e, an edge from one node to another. If the nodes do not exist, they are added.
// It will panic if the IDs of the e.From and e.To are equal.
func (g *ExprGraph) SetEdge(e graph.Edge) {
	from := e.From().(*Node)
	to := e.To().(*Node)

	if from == to {
		panic(fmt.Sprintf("cannot add self edge: from %v to %v", from, to))
	}

	if !g.Has(from) {
		from = g.AddNode(from)
	}

	if !g.Has(to) {
		to = g.AddNode(to)
	}

	// g.to[to] = g.to[to].Add(from)
	g.to[to] = append(g.to[to], from)
}

// Roots returns a list of nodes that are not children of any other nodes
func (g *ExprGraph) Roots() (retVal Nodes) {
	// handle subgraph
	if g.roots != nil {
		return g.roots
	}

	for n, tos := range g.to {
		if len(tos) == 0 {
			retVal = append(retVal, n)
		}
	}
	return retVal
}

// Inputs returns a list of nodes which are inputs (that is to say, the user is required to set a value in it)
func (g *ExprGraph) Inputs() (retVal Nodes) {
	for _, n := range g.all {
		if n.isInput() {
			retVal = append(retVal, n)
		}
	}
	return
}

// ByName returns nodes that have the name provided.
// Bear in mind that the name that is compared to is the internal name,
// not the result of calling node.Name(). The reason for doing this is
// for ease of finding only names that are user-supplied, instead of autogenerated names
func (g *ExprGraph) ByName(name string) (retVal Nodes) {
	for _, n := range g.byHash {
		if n.name == name {
			retVal = append(retVal, n)
		}
	}
	return
}

func (g *ExprGraph) String() string {
	var buf bytes.Buffer
	buf.WriteString("Graph: [\n")
	for _, n := range g.byHash {
		fmt.Fprintf(&buf, "\t%d: %s\n", n.Hashcode(), n)
	}
	buf.WriteString("]")
	return buf.String()
}

// ToDot generates the graph in graphviz format. The use of this is to generate for the entire graph
// which may have multiple trees with different roots
// TODO: This is getting unwieldly. Perhaps refactor out into a ToDot(...Opt)?
func (g *ExprGraph) ToDot() string {
	gv := gographviz.NewEscape()
	gv.SetName(fullGraphName)
	gv.SetDir(true)

	gv.AddAttr(fullGraphName, "nodesep", "1")
	gv.AddAttr(fullGraphName, "ranksep", "1.5 equally")
	gv.AddAttr(fullGraphName, "rankdir", "TB")
	if len(g.byHash) > 100 {
		gv.AddAttr(fullGraphName, "nslimit", "3") // numiter=3*len(nodes)
		// gv.AddAttr(fullGraphName, "splines", "line") // ugly as sin.
	}

	groups := make(map[string]struct{})
	for h, n := range g.byHash {
		if n != nil {
			group := n.dotCluster()
			groups[group] = struct{}{}
			continue
		}
		// other wise it'se a clash of hash
		for _, n := range g.evac[h] {
			group := n.dotCluster()
			groups[group] = struct{}{}

		}
	}

	for grp := range groups {
		attrs := gographviz.NewAttrs()
		attrs.Add("label", grp)

		parentGraph := fullGraphName
		if grp == inputsClust || grp == constantsClust {
			parentGraph = inputConsts
			if !gv.IsSubGraph(inputConsts) {
				groupAttrs := gographviz.NewAttrs()
				groupAttrs.Add("rank", "max")
				gv.AddSubGraph(fullGraphName, inputConsts, groupAttrs)
			}
		}
		gv.AddSubGraph(parentGraph, "cluster_"+grp, attrs)
	}

	// for _, n := range g.byHash {
	for _, n := range g.all {
		group := n.dotCluster()
		n.dotString(gv, "cluster_"+group)
	}

	// for _, from := range g.byHash {
	for _, from := range g.all {
		for i, child := range from.children {
			if ok := g.all.Contains(child); !ok {
				// not in graph, so ignore it...
				continue
			}
			fromID := fmt.Sprintf("Node_%p", from)
			toID := fmt.Sprintf("Node_%p", child)

			edgeAttrs := gographviz.NewAttrs()
			edgeAttrs.Add("taillabel", fmt.Sprintf(" %d ", i))
			edgeAttrs.Add("labelfloat", "false")

			// we invert the from and to nodes for gradients, As the expressionGraph builds upwards from bottom, the gradient builds downwards.
			if from.group == gradClust && child.group == gradClust {
				edgeAttrs.Add("dir", "back")
				gv.AddPortEdge(toID, toID+":anchor:s", fromID, fromID+":anchor:n", true, edgeAttrs)
			} else {
				gv.AddPortEdge(fromID, fromID+":anchor:s", toID, toID+":anchor:n", true, edgeAttrs)
			}
		}
	}

	// draw deriv lines
	if debugDerives {
		edgeAttrs := gographviz.NewAttrs()
		edgeAttrs.Add("style", "dashed")
		edgeAttrs.Add("constraint", "false")
		edgeAttrs.Add("weight", "999")

		for _, n := range g.byHash {
			if n.derivOf != nil {
				id := fmt.Sprintf("Node_%p", n)
				for _, derivOf := range n.derivOf {
					if _, ok := g.to[derivOf]; !ok {
						continue
					}
					ofID := fmt.Sprintf("Node_%p", derivOf)
					// gv.AddPortEdge(id, ":anchor:w", ofID, ofID+":anchor:e", true, edgeAttrs)
					gv.AddEdge(id, ofID, true, edgeAttrs)
				}
			}
		}
	}

	// stupid invisible nodes to keep expressiongraph on the left
	subGAttrs := gographviz.NewAttrs()
	// subGAttrs.Add("rank", "max")
	gv.AddSubGraph(fullGraphName, outsideSubG, subGAttrs)

	attrs := gographviz.NewAttrs()
	attrs.Add("style", "invis")
	gv.AddNode(outsideSubG, outsideRoot, attrs)

	outsides := []string{outsideRoot}
	var insides []string

	// build the inside and outside list
	if _, hasInputs := groups[inputsClust]; hasInputs {
		insides = append(insides, insideInputs)
		gv.AddNode("cluster_inputs", insideInputs, attrs)
	}

	if _, hasConst := groups[constantsClust]; hasConst {
		if len(insides) > 0 {
			outsides = append(outsides, outsideConsts)
			gv.AddNode(outsideSubG, outsideConsts, attrs)
		}
		insides = append(insides, insideConsts)
		gv.AddNode("cluster_constants", insideConsts, attrs)
	}

	if len(insides) > 0 {
		outsides = append(outsides, outsideExprG)
		gv.AddNode(outsideSubG, outsideExprG, attrs)
	}
	insides = append(insides, insideExprG)
	gv.AddNode("cluster_expressionGraph", insideExprG, attrs)

	for group := range groups {
		if group == exprgraphClust || group == constantsClust || group == inputsClust {
			continue
		}
		inside := "inside_" + group
		outside := "outside_" + group
		insides = append(insides, inside)
		outsides = append(outsides, outside)

		gv.AddNode(outsideSubG, outside, attrs)
		gv.AddNode("cluster_"+group, inside, attrs)
	}

	edgeAttrs := gographviz.NewAttrs()
	edgeAttrs.Add("style", "invis")
	edgeAttrs.Add("weight", "999")
	edgeAttrs.Add("constraint", "false")
	for i, o := range outsides {
		// outside-inside
		gv.AddEdge(o, insides[i], true, edgeAttrs)

		if i > 0 {
			// outside-outside
			gv.AddEdge(outsides[i-1], o, true, edgeAttrs)

			// inside-inside
			gv.AddEdge(insides[i-1], insides[i], true, edgeAttrs)
		}
	}
	return gv.String()
}

// other private methods

func (g *ExprGraph) removeAllEdgesFrom(n *Node) {
	for k, ns := range g.to {
		g.to[k] = ns.remove(n)
	}
}

/* Graph interface */

// Node returns the node in the graph with the given ID.
func (g *ExprGraph) Node(id int) graph.Node {
	n := (*Node)(unsafe.Pointer(uintptr(id)))
	if _, ok := g.to[n]; ok {
		return n
	}
	return nil
}

// Has returns whether the node exists within the graph.
func (g *ExprGraph) Has(node graph.Node) bool {
	n := node.(*Node)
	_, ok := g.to[n]
	return ok
}

// Nodes returns all the nodes in the graph.
func (g *ExprGraph) Nodes() []graph.Node {
	// nodes := make([]graph.Node, len(g.from))
	ns := g.AllNodes()

	nodes := nodeToGraphNode(ns)
	return nodes
}

// AllNodes is like Nodes, but returns Nodes instead of []graph.Node.
// Nodes() has been reserved for the graph.Directed interface, so this one is named AllNodes instead
func (g *ExprGraph) AllNodes() Nodes {
	return g.all
}

// From returns all nodes in g that can be reached directly from n.
func (g *ExprGraph) From(node graph.Node) []graph.Node {
	n := node.(*Node)
	return nodeToGraphNode(n.children)
}

// HasEdgeBetween returns whether an edge exists between nodes x and y without
// considering direction.
func (g *ExprGraph) HasEdgeBetween(x, y graph.Node) bool {
	xid := x.(*Node)
	yid := y.(*Node)

	return xid.children.Contains(yid) || yid.children.Contains(xid)
}

// Edge returns the edge from u to v if such an edge exists and nil otherwise.
// The node v must be directly reachable from u as defined by the From method.
func (g *ExprGraph) Edge(u, v graph.Node) graph.Edge {
	uid := u.(*Node)
	vid := v.(*Node)

	if !uid.children.Contains(vid) {
		return nil
	}
	e := edge{from: uid, to: vid}
	return e
}

/* Directed interface */

// HasEdgeFromTo returns whether an edge exists in the graph from u to v.
func (g *ExprGraph) HasEdgeFromTo(u, v graph.Node) bool {
	uid := u.(*Node)
	vid := v.(*Node)

	return uid.children.Contains(vid)
}

// To returns all nodes in g that can reach directly to n.
func (g *ExprGraph) To(n graph.Node) []graph.Node {
	ns := g.to[n.(*Node)]
	ns = ns.Set()
	g.to[n.(*Node)] = ns
	return nodeToGraphNode(ns)
}

// subgraph is basically a subset of nodes. This is useful for compiling sub sections of the graph
func (g *ExprGraph) subgraph(ns Nodes, opts ...Nodes) *ExprGraph {
	ns = ns.Set()
	allset := ns.mapSet()

	// uniquify the froms and at the same time build a new roots
	var roots Nodes

	if len(opts) == 0 {
		for _, n := range ns {
			if len(g.to[n]) == 0 {
				if n.isStmt {
					roots = append(roots, n.children[0])
				} else {
					roots = append(roots, n)
				}
				continue
			}

			var hasParent bool
			for _, parent := range g.to[n] {
				if allset.Contains(parent) {
					hasParent = true
					break
				}
			}
			if !hasParent {
				roots = append(roots, n)
			}
		}
	} else {
		rs := opts[0]
		roots = make(Nodes, len(rs))
		for i, n := range rs {
			if n.isStmt {
				roots[i] = n.children[0]
				continue
			}
			roots[i] = n
		}
	}
	roots = roots.Set()

	retVal := &ExprGraph{
		all:    ns,
		byHash: g.byHash,
		evac:   g.evac,
		to:     g.to,

		leaves:    g.leaves,
		constants: g.constants,
		roots:     roots,
	}

	return retVal
}

// Subgraph is a function with overloaded meanings. If only one node is passed in, it assumes that the one node is the root,
// otherwise, it treats ns as the subset of nodes to be included in the subgraph
func (g *ExprGraph) Subgraph(ns ...*Node) *ExprGraph {
	if len(ns) == 1 {
		g.SubgraphRoots(ns[0])
	}
	return g.subgraph(ns)
}

func (g *ExprGraph) SubgraphRoots(ns ...*Node) *ExprGraph {
	sub := make(Nodes, len(ns))
	copy(sub, ns)

	walked := NewNodeSet()
	for _, n := range ns {
		ch := make(chan *Node)
		go func(ch chan *Node) {
			defer close(ch)
			walkGraph(n, ch, walked)
		}(ch)

		for node := range ch {
			sub = append(sub, node)
		}
	}

	return g.subgraph(sub, ns)
}

type edge struct {
	from, to graph.Node
	weight   float64
}

func (e edge) From() graph.Node { return e.from }
func (e edge) To() graph.Node   { return e.to }
func (e edge) Weight() float64  { return e.weight }