package algs44;
import  stdlib.*;

/* ***********************************************************************
 *  Compilation:  javac AssignmentProblem.java
 *  Execution:    java AssignmentProblem N
 *  Dependencies: DijkstraSP.java DirectedEdge.java
 *
 *  Solve an N-by-N assignment problem in N^3 log N time using the
 *  successive shortest path algorithm.
 *
 *  Remark: could use dense version of Dijsktra's algorithm for
 *  improved theoretical efficiency of N^3, but it doesn't seem to
 *  help in practice.
 *
 *  Assumes N-by-N cost matrix is nonnegative.
 *
 *
 *********************************************************************/

public class AssignmentProblem {
  private static final int UNMATCHED = -1;

  private int N;              // number of rows and columns
  private double[][] weight;  // the N-by-N cost matrix
  private double[] px;        // px[i] = dual variable for row i
  private double[] py;        // py[j] = dual variable for col j
  private int[] xy;           // xy[i] = j means i-j is a match
  private int[] yx;           // yx[j] = i means i-j is a match


  public AssignmentProblem(double[][] weight) {
    N = weight.length;
    this.weight = new double[N][N];
    for (int i = 0; i < N; i++)
      for (int j = 0; j < N; j++)
        this.weight[i][j] = weight[i][j];

    // dual variables
    px = new double[N];
    py = new double[N];

    // initial matching is empty
    xy = new int[N];
    yx = new int[N];
    for (int i = 0; i < N; i++) xy[i] = UNMATCHED;
    for (int j = 0; j < N; j++) yx[j] = UNMATCHED;

    // add N edges to matching
    for (int k = 0; k < N; k++) {
      assert isDualFeasible();
      assert isComplementarySlack();
      augment();
    }
    assert check();
  }

  // find shortest augmenting path and upate
  private void augment() {

    // build residual graph
    EdgeWeightedDigraph G = new EdgeWeightedDigraph(2*N+2);
    int s = 2*N, t = 2*N+1;
    for (int i = 0; i < N; i++) {
      if (xy[i] == UNMATCHED) G.addEdge(new DirectedEdge(s, i, 0.0));
    }
    for (int j = 0; j < N; j++) {
      if (yx[j] == UNMATCHED) G.addEdge(new DirectedEdge(N+j, t, py[j]));
    }
    for (int i = 0; i < N; i++) {
      for (int j = 0; j < N; j++) {
        if (xy[i] == j) G.addEdge(new DirectedEdge(N+j, i, 0.0));
        else            G.addEdge(new DirectedEdge(i, N+j, reduced(i, j)));
      }
    }

    // compute shortest path from s to every other vertex
    DijkstraSP spt = new DijkstraSP(G, s);

    // augment along alternating path
    for (DirectedEdge e : spt.pathTo(t)) {
      int i = e.from(), j = e.to() - N;
      if (i < N) {
        xy[i] = j;
        yx[j] = i;
      }
    }

    // update dual variables
    for (int i = 0; i < N; i++) px[i] += spt.distTo(i);
    for (int j = 0; j < N; j++) py[j] += spt.distTo(N+j);
  }

  // reduced cost of i-j
  private double reduced(int i, int j) {
    return weight[i][j] + px[i] - py[j];
  }

  // dual variable for row i
  public double dualRow(int i) {
    return px[i];
  }

  // dual variable for column j
  public double dualCol(int j) {
    return py[j];
  }

  // total weight of min weight perfect matching
  public double weight() {
    double total = 0.0;
    for (int i = 0; i < N; i++) {
      if (xy[i] != UNMATCHED)
        total += weight[i][xy[i]];
    }
    return total;
  }

  public int sol(int i) {
    return xy[i];
  }

  // check that dual variables are feasible
  private boolean isDualFeasible() {
    // check that all edges have >= 0 reduced cost
    for (int i = 0; i < N; i++) {
      for (int j = 0; j < N; j++) {
        if (reduced(i, j) < 0) {
          StdOut.println("Dual variables are not feasible");
          return false;
        }
      }
    }
    return true;
  }

  // check that primal and dual variables are complementary slack
  private boolean isComplementarySlack() {

    // check that all matched edges have 0-reduced cost
    for (int i = 0; i < N; i++) {
      if ((xy[i] != UNMATCHED) && (reduced(i, xy[i]) != 0)) {
        StdOut.println("Primal and dual variables are not complementary slack");
        return false;
      }
    }
    return true;
  }

  // check that primal variables are a perfect matching
  private boolean isPerfectMatching() {

    // check that xy[] is a perfect matching
    boolean[] perm = new boolean[N];
    for (int i = 0; i < N; i++) {
      if (perm[xy[i]]) {
        StdOut.println("Not a perfect matching");
        return false;
      }
      perm[xy[i]] = true;
    }

    // check that xy[] and yx[] are inverses
    for (int j = 0; j < N; j++) {
      if (xy[yx[j]] != j) {
        StdOut.println("xy[] and yx[] are not inverses");
        return false;
      }
    }
    for (int i = 0; i < N; i++) {
      if (yx[xy[i]] != i) {
        StdOut.println("xy[] and yx[] are not inverses");
        return false;
      }
    }

    return true;
  }


  // check optimality conditions
  private boolean check() {
    return isPerfectMatching() && isDualFeasible() && isComplementarySlack();
  }

  public static void main(String[] args) {
    In in = new In(args[0]);
    int N = in.readInt();
    double[][] weight = new double[N][N];
    for (int i = 0; i < N; i++) {
      for (int j = 0; j < N; j++) {
        weight[i][j] = in.readDouble();
      }
    }

    AssignmentProblem assignment = new AssignmentProblem(weight);
    StdOut.println("weight = " + assignment.weight());
    for (int i = 0; i < N; i++)
      StdOut.println(i + "-" + assignment.sol(i) + "' " + weight[i][assignment.sol(i)]);

    for (int i = 0; i < N; i++)
      StdOut.println("px[" + i + "] = " + assignment.dualRow(i));
    for (int j = 0; j < N; j++)
      StdOut.println("py[" + j + "] = " + assignment.dualCol(j));
    for (int i = 0; i < N; i++)
      for (int j = 0; j < N; j++)
        StdOut.println("reduced[" + i + "-" + j + "] = " + assignment.reduced(i, j));

  }

}