package edu.uci.ics.jung.algorithms.shortestpath;
   
   import java.util.Collection;
   import java.util.HashSet;
   import java.util.Map;
   import java.util.Set;
   
   import com.google.common.base.Function;
   import com.google.common.base.Functions;
  import com.google.common.base.Supplier;
  
  import edu.uci.ics.jung.graph.Forest;
  import edu.uci.ics.jung.graph.Graph;
  import edu.uci.ics.jung.graph.util.EdgeType;
  import edu.uci.ics.jung.graph.util.Pair;
  
  /**
   * For the input Graph, creates a MinimumSpanningTree
   * using a variation of Prim's algorithm.
   * 
   * @author Tom Nelson - tomnelson@dev.java.net
   *
   * @param <V> the vertex type
   * @param <E> the edge type
   */
  public class MinimumSpanningForest<V,E> {
  	
  	protected Graph<V,E> graph;
  	protected Forest<V,E> forest;
  	protected Function<E, Double> weights;
  	
  	/**
  	 * Creates a Forest from the supplied Graph and supplied Supplier, which
  	 * is used to create a new, empty Forest. If non-null, the supplied root
  	 * will be used as the root of the tree/forest. If the supplied root is
  	 * null, or not present in the Graph, then an arbitrary Graph vertex
  	 * will be selected as the root.
  	 * If the Minimum Spanning Tree does not include all vertices of the
  	 * Graph, then a leftover vertex is selected as a root, and another
  	 * tree is created.
  	 * @param graph the input graph
  	 * @param Supplier the Supplier to use to create the new forest
  	 * @param root the vertex of the graph to be used as the root of the forest 
  	 * @param weights edge weights
  	 */
  	public MinimumSpanningForest(Graph<V, E> graph, Supplier<Forest<V,E>> Supplier, 
  			V root, Map<E, Double> weights) {
  		this(graph, Supplier.get(), root, weights);
  	}
  	
  	/**
  	 * Creates a minimum spanning forest from the supplied graph, populating the
  	 * supplied Forest, which must be empty. 
  	 * If the supplied root is null, or not present in the Graph,
  	 * then an arbitrary Graph vertex will be selected as the root.
  	 * If the Minimum Spanning Tree does not include all vertices of the
  	 * Graph, then a leftover vertex is selected as a root, and another
  	 * tree is created
  	 * @param graph the Graph to find MST in
  	 * @param forest the Forest to populate. Must be empty
  	 * @param root first Tree root, may be null
  	 * @param weights edge weights, may be null
  	 */
  	public MinimumSpanningForest(Graph<V, E> graph, Forest<V,E> forest, 
  			V root, Map<E, Double> weights) {
  		
  		if(forest.getVertexCount() != 0) {
  			throw new IllegalArgumentException("Supplied Forest must be empty");
  		}
  		this.graph = graph;
  		this.forest = forest;
  		if(weights != null) {
  			this.weights = Functions.forMap(weights);
  		}
  		Set<E> unfinishedEdges = new HashSet<E>(graph.getEdges());
  		if(graph.getVertices().contains(root)) {
  			this.forest.addVertex(root);
  		}
  		updateForest(forest.getVertices(), unfinishedEdges);
  	}
  	
      /**
       * Creates a minimum spanning forest from the supplied graph, populating the
       * supplied Forest, which must be empty. 
       * If the supplied root is null, or not present in the Graph,
       * then an arbitrary Graph vertex will be selected as the root.
       * If the Minimum Spanning Tree does not include all vertices of the
       * Graph, then a leftover vertex is selected as a root, and another
       * tree is created
       * @param graph the Graph to find MST in
       * @param forest the Forest to populate. Must be empty
       * @param root first Tree root, may be null
       */
      @SuppressWarnings("unchecked")
  	public MinimumSpanningForest(Graph<V, E> graph, Forest<V,E> forest, 
              V root) {
          
          if(forest.getVertexCount() != 0) {
              throw new IllegalArgumentException("Supplied Forest must be empty");
         }
         this.graph = graph;
         this.forest = forest;
         this.weights = (Function<E, Double>) Functions.constant(1.0);
         Set<E> unfinishedEdges = new HashSet<E>(graph.getEdges());
         if(graph.getVertices().contains(root)) {
             this.forest.addVertex(root);
         }
         updateForest(forest.getVertices(), unfinishedEdges);
     }
 	
 	/**
 	 * @return the generated forest
 	 */
 	public Forest<V,E> getForest() {
 		return forest;
 	}
 	
 	protected void updateForest(Collection<V> tv, Collection<E> unfinishedEdges) {
 		double minCost = Double.MAX_VALUE;
 		E nextEdge = null;
 		V nextVertex = null;
 		V currentVertex = null;
 		for(E e : unfinishedEdges) {
 			
 			if(forest.getEdges().contains(e)) continue;
 			// find the lowest cost edge, get its opposite endpoint,
 			// and then update forest from its Successors
 			Pair<V> endpoints = graph.getEndpoints(e);
 			V first = endpoints.getFirst();
 			V second = endpoints.getSecond();
 			if(tv.contains(first) == true && tv.contains(second) == false) {
 				if(weights.apply(e) < minCost) {
 					minCost = weights.apply(e);
 					nextEdge = e;
 					currentVertex = first;
 					nextVertex = second;
 				}
 			}
 			if(graph.getEdgeType(e) == EdgeType.UNDIRECTED &&
 					tv.contains(second) == true && tv.contains(first) == false) {
 				if(weights.apply(e) < minCost) {
 					minCost = weights.apply(e);
 					nextEdge = e;
 					currentVertex = second;
 					nextVertex = first;
 				}
 			}
 		}
 		
 		if(nextVertex != null && nextEdge != null) {
 			unfinishedEdges.remove(nextEdge);
 			forest.addEdge(nextEdge, currentVertex, nextVertex);
 			updateForest(forest.getVertices(), unfinishedEdges);
 		}
 		Collection<V> leftovers = new HashSet<V>(graph.getVertices());
 		leftovers.removeAll(forest.getVertices());
 		if(leftovers.size() > 0) {
 			V anotherRoot = leftovers.iterator().next();
 			forest.addVertex(anotherRoot);
 			updateForest(forest.getVertices(), unfinishedEdges);
 		}
 	}
 }