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));

        }

}