// Exercise 4.3.21 4.3.22 (Solution published at http://algs4.cs.princeton.edu/)
package algs43;
import stdlib.*;
import algs13.Queue;
import algs15.WeightedUF;
import algs24.IndexMinPQ;
/* ****************************************************************************
 *  Compilation:  javac PrimMST.java
 *  Execution:    java PrimMST filename.txt
 *  Dependencies: EdgeWeightedGraph.java Edge.java Queue.java
 *                IndexMinPQ.java UF.java In.java StdOut.java
 *  Data files:   http://algs4.cs.princeton.edu/43mst/tinyEWG.txt
 *                http://algs4.cs.princeton.edu/43mst/mediumEWG.txt
 *                http://algs4.cs.princeton.edu/43mst/largeEWG.txt
 *
 *  Compute a minimum spanning forest using Prim's algorithm.
 *
 *  %  java PrimMST tinyEWG.txt
 *  1-7 0.19000
 *  0-2 0.26000
 *  2-3 0.17000
 *  4-5 0.35000
 *  5-7 0.28000
 *  6-2 0.40000
 *  0-7 0.16000
 *  1.81000
 *
 *  % java PrimMST mediumEWG.txt
 *  1-72   0.06506
 *  2-86   0.05980
 *  3-67   0.09725
 *  4-55   0.06425
 *  5-102  0.03834
 *  6-129  0.05363
 *  7-157  0.00516
 *  ...
 *  10.46351
 *
 *  % java PrimMST largeEWG.txt
 *  ...
 *  647.66307
 *
 ******************************************************************************/

public class PrimMST {
  private final Edge[] edgeTo;        // edgeTo[v] = shortest edge from tree vertex to non-tree vertex
  private final double[] distTo;      // distTo[v] = weight of shortest such edge
  private final boolean[] marked;     // marked[v] = true if v on tree, false otherwise
  private final IndexMinPQ<Double> pq;

  public PrimMST(EdgeWeightedGraph G) {
    edgeTo = new Edge[G.V()];
    distTo = new double[G.V()];
    marked = new boolean[G.V()];
    pq = new IndexMinPQ<>(G.V());
    for (int v = 0; v < G.V(); v++) distTo[v] = Double.POSITIVE_INFINITY;

    for (int v = 0; v < G.V(); v++)      // run from each vertex to find
      if (!marked[v]) prim(G, v);      // minimum spanning forest

    // check optimality conditions
    assert check(G);
  }

  // run Prim's algorithm in graph G, starting from vertex s
  private void prim(EdgeWeightedGraph G, int s) {
    distTo[s] = 0.0;
    pq.insert(s, distTo[s]);
    while (!pq.isEmpty()) {
      int v = pq.delMin();
      scan(G, v);
    }
  }

  // scan vertex v
  private void scan(EdgeWeightedGraph G, int v) {
    marked[v] = true;
    for (Edge e : G.adj(v)) {
      int w = e.other(v);
      if (marked[w]) continue;         // v-w is obsolete edge
      if (e.weight() < distTo[w]) {
        distTo[w] = e.weight();
        edgeTo[w] = e;
        if (pq.contains(w)) pq.decreaseKey(w, distTo[w]);
        else                pq.insert(w, distTo[w]);
      }
    }
  }

  // return iterator of edges in MST
  public Iterable<Edge> edges() {
    Queue<Edge> mst = new Queue<>();
    for (Edge e : edgeTo) {
      if (e != null) {
        mst.enqueue(e);
      }
    }
    return mst;
  }


  // return weight of MST
  public double weight() {
    double weight = 0.0;
    for (Edge e : edges())
      weight += e.weight();
    return weight;
  }


  // check optimality conditions (takes time proportional to E V lg* V)
  private boolean check(EdgeWeightedGraph G) {

    // check weight
    double totalWeight = 0.0;
    for (Edge e : edges()) {
      totalWeight += e.weight();
    }
    double EPSILON = 1E-12;
    if (Math.abs(totalWeight - weight()) > EPSILON) {
      System.err.format("Weight of edges does not equal weight(): %f vs. %f\n", totalWeight, weight());
      return false;
    }

    // check that it is acyclic
    WeightedUF uf = new WeightedUF(G.V());
    for (Edge e : edges()) {
      int v = e.either(), w = e.other(v);
      if (uf.connected(v, w)) {
        System.err.println("Not a forest");
        return false;
      }
      uf.union(v, w);
    }

    // check that it is a spanning forest
    for (Edge e : edges()) {
      int v = e.either(), w = e.other(v);
      if (!uf.connected(v, w)) {
        System.err.println("Not a spanning forest");
        return false;
      }
    }

    // check that it is a minimal spanning forest (cut optimality conditions)
    for (Edge e : edges()) {
      int v = e.either(), w = e.other(v);

      // all edges in MST except e
      uf = new WeightedUF(G.V());
      for (Edge f : edges()) {
        int x = f.either(), y = f.other(x);
        if (f != e) uf.union(x, y);
      }

      // check that e is min weight edge in crossing cut
      for (Edge f : G.edges()) {
        int x = f.either(), y = f.other(x);
        if (!uf.connected(x, y)) {
          if (f.weight() < e.weight()) {
            System.err.println("Edge " + f + " violates cut optimality conditions");
            return false;
          }
        }
      }

    }

    return true;
  }


  public static void main(String[] args) {
    args = new String[] { "data/10000EWG.txt" };
    //args = new String[] { "data/mediumEWG.txt" };
    //args = new String[] { "data/largeEWG.txt" };
    In in = new In(args[0]);
    EdgeWeightedGraph G = new EdgeWeightedGraph(in);
    PrimMST mst = new PrimMST(G);
    for (Edge e : mst.edges()) {
      StdOut.println(e);
    }
    StdOut.format("%.5f\n", mst.weight());
  }


}