package org.javaan.graph;

import java.util.HashSet;
import java.util.Set;
import java.util.Stack;

import org.jgrapht.DirectedGraph;
import org.jgrapht.event.TraversalListener;
import org.jgrapht.graph.EdgeReversedGraph;
import org.jgrapht.graph.GraphDelegator;
import org.jgrapht.traverse.BreadthFirstIterator;
import org.jgrapht.traverse.DepthFirstIterator;
import org.jgrapht.traverse.GraphIterator;

public class TraversalDirectedGraph<V, E> extends GraphDelegator<V, E> implements DirectedGraph<V, E> {

	private static final long serialVersionUID = 1L;

	public TraversalDirectedGraph(DirectedGraph<V, E> g) {
		super(g);
	}

	/**
	 * Add sourceVertex and targetVertex to graph and create an edge between them.
	 */
	@Override
	public E addEdge(V sourceVertex, V targetVertex) {
		addVertex(sourceVertex);
		addVertex(targetVertex);
		return super.addEdge(sourceVertex, targetVertex);
	}
	
	@Override
	public boolean addEdge(V sourceVertex, V targetVertex, E edge) {
		addVertex(sourceVertex);
		addVertex(targetVertex);
		return super.addEdge(sourceVertex, targetVertex, edge);
	}

	private Set<V> getTargetSet(Set<E> edges) {
		Set<V> targets = new HashSet<V>();
		for (E e : edges) {
			targets.add(getEdgeTarget(e));
		}
		return targets;
	}

	private Set<V> getSourceSet(Set<E> edges) {
		Set<V> sources = new HashSet<V>();
		for (E e : edges) {
			sources.add(getEdgeSource(e));
		}
		return sources;
	}

	public Set<V> targetVerticesOf(V vertex) {
		return getTargetSet(outgoingEdgesOf(vertex));
	}

	public Set<V> sourceVerticesOf(V vertex) {
		return getSourceSet(incomingEdgesOf(vertex));
	}

	public Set<V> successorsOf(V vertex) {
		Set<V> successors = new HashSet<V>();
		Stack<V> successorStack = new Stack<V>();
		successorStack.addAll(targetVerticesOf(vertex));
		while(!successorStack.isEmpty()) {
			V successor = successorStack.pop();
			// detect cycle, ignore existing nodes
			if (!successors.contains(successor)) {
				successors.add(successor);
				successorStack.addAll(targetVerticesOf(successor));
			}
		}
		return successors;
	}

	public Set<V> predecessorsOf(V vertex) {
		Set<V> predecessors = new HashSet<V>();
		Stack<V> predecessorStack = new Stack<V>();
		predecessorStack.addAll(sourceVerticesOf(vertex));
		while(!predecessorStack.isEmpty()) {
			V predecessor = predecessorStack.pop();
			// detect cycle, ignore existing nodes
			if (!predecessors.contains(predecessor)) {
				predecessors.add(predecessor);
				predecessorStack.addAll(sourceVerticesOf(predecessor));
			}
		}
		return predecessors;
	}

	private void traverseGraph(DirectedGraph<V, E> graph, V startVertex, GraphVisitor<V, E> visitor, boolean depthFirst) {
		if (!containsVertex(startVertex)) {
			return;
		}
		TraversalListener<V, E> listener = null;
		GraphIterator<V, E> iterator = null;
		if (depthFirst) {
			listener = new DepthFirstTraversalListener<V, E>(graph, visitor);
			iterator = new DepthFirstIterator<V, E>(graph, startVertex);
		} else {
			listener = new BreadthFirstTraversalListener<V, E>(visitor);
			iterator = new BreadthFirstIterator<V, E>(graph, startVertex);
		}
		iterator.addTraversalListener(listener);
		while (iterator.hasNext() && !visitor.finished()) {
			iterator.next();
		}
	}

	public void traverseSuccessorsDepthFirst(V startVertex, GraphVisitor<V, E> visitor) {
		traverseGraph(this, startVertex, visitor, true);
	}

	public void traversePredecessorsDepthFirst(V startVertex, GraphVisitor<V, E> visitor) {
		DirectedGraph<V, E> graph = new EdgeReversedGraph<V, E>(this);
		traverseGraph(graph, startVertex, visitor, true);
	}

	public void traverseSuccessorsBreadthFirst(V startVertex, GraphVisitor<V, E> visitor) {
		traverseGraph(this, startVertex, visitor, false);
	}

	public void traversePredecessorsBreadthFirst(V startVertex, GraphVisitor<V, E> visitor) {
		DirectedGraph<V, E> graph = new EdgeReversedGraph<V, E>(this);
		traverseGraph(graph, startVertex, visitor, false);
	}

	private Set<V> collectLeafVertices(DirectedGraph<V, E> graph, V startVertex) {
		Set<V> leaves = new HashSet<V>();
		if (graph.containsVertex(startVertex)) {
			GraphIterator<V, E> iterator = new DepthFirstIterator<V, E>(graph, startVertex);
			iterator.next();
			while (iterator.hasNext()) {
				V vertex = iterator.next();
				if (graph.outDegreeOf(vertex) == 0) {
					leaves.add(vertex);
				}
			}
		}
		return leaves;
	}

	public Set<V> getLeafSuccessors(V startVertex) {
		return collectLeafVertices(this, startVertex);
	}

	public Set<V> getLeafPredecessors(V startVertex) {
		DirectedGraph<V, E> graph = new EdgeReversedGraph<V, E>(this);
		return collectLeafVertices(graph, startVertex);
	}

}

@@ -2669,16 +2669,39 @@
 %7B%0A%09%09if (
+startVertex != null && 
 !contain
@@ -3240,24 +3240,135 @@
 ();%0A%09%09%7D%0A%09%7D%0A%0A
+%09public void traverseDepthFirst(GraphVisitor%3CV, E%3E visitor) %7B%0A%09%09traverseGraph(this, null, visitor, true);%0A%09%7D%0A%09%0A
 %09public void
@@ -3699,32 +3699,145 @@
 tor, true);%0A%09%7D%0A%0A
+%09public void traverseBreadthFirst(GraphVisitor%3CV, E%3E visitor) %7B%0A%09%09traverseGraph(this, null, visitor, false);%0A%09%7D%0A%0A
 %09public void tra

support graph traversal without start vertex