/* Copyright (c) 1998-2016
   Ewgenij Gawrilow, Michael Joswig (Technische Universitaet Berlin, Germany)
   http://www.polymake.org

   This program is free software; you can redistribute it and/or modify it
   under the terms of the GNU General Public License as published by the
   Free Software Foundation; either version 2, or (at your option) any
   later version: http://www.gnu.org/licenses/gpl.txt.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.
--------------------------------------------------------------------------------
*/

#ifndef POLYMAKE_HUNGARIAN_METHOD_H
#define POLYMAKE_HUNGARIAN_METHOD_H

#include "polymake/client.h"
#include "polymake/Graph.h"
#include "polymake/Set.h"
#include "polymake/Vector.h"
#include "polymake/Matrix.h"
#include "polymake/graph/graph_iterators.h"

namespace polymake { namespace graph {

      // The implementation is adapted to Figure 11-2 in
      // Papadimitriou, Christos H.; Steiglitz, Kenneth
      // Combinatorial optimization: algorithms and complexity. Prentice-Hall, Inc., Englewood Cliffs, N.J., 1982. xvi+496 pp. ISBN: 0-13-152462-3
      //
      // Corrections of the algorithms in this book can be found at
      // www.cs.princeton.edu/~ken/latest.pdf
      //
      // We fixed the algorithm independently of those errata.

template <typename E>
class HungarianMethod;

class TreeGrowVisitor : public NodeVisitor<> {
   using base_t = NodeVisitor<>;
   template <typename> friend class HungarianMethod;

   // label encodes the path in the hungarian trees from an exposed point to the selected node
   std::vector<int> label;
   int leaf;
   const int dim;
   const Graph<Directed>* H;
   Set<int> start_nodes;
public:
   TreeGrowVisitor() = default;

   TreeGrowVisitor(const Graph<Directed>& G)
      : base_t(G)
      , label(G.top().dim(), -1)
      , leaf(-1)
      , dim((G.top().dim() + 1)/2)
      , H(&G)
   {}

   void clear(const Graph<Directed>&) {}

   // This method provides two functionalities: If the search had found a augmenting path and hence the equality subgraph was modified, it initializes a new search. Otherwise only another starting node for the growing of an hungarian tree is set.
   bool operator()(const int start_node)
   {
      if (start_nodes.contains(start_node) || leaf >= 0) {
         start_nodes.clear();
         std::fill(label.begin(), label.end(), -1);
         visited.clear();
         leaf = -1;
      }
      label[start_node] = start_node;
      visited += start_node;
      start_nodes += start_node;
      return true;
   }

   bool operator() (int n_from, int n_to)
   {
      if (visited.contains(n_to)) return false;

      visited += n_to;
      // here the label is set during a BFS step in the equality subgraph
      if (H->edge_exists(n_from, n_to)) {
         label[n_to] = n_from;
      }
      // checking, if n is a leaf and is on the right side and so the path to n is augmenting
      if ((n_to >= dim) && (H->out_degree(n_to) == 0)) leaf = n_to;
      return true;
   }

   const int& operator[] (int n) const { return label[n]; }
};

template <typename E>
class HungarianMethod {
protected:
   Matrix<E> weights;
   const int dim;

   Vector<E> a, b, slack, labeledColMin;
   Graph<Directed> equality_subgraph;
   Set<int> exposed_points;
   Set<int>::const_iterator r;  // an iterator over the exposed points
   int start_node;
   BFSiterator<Graph<Directed>, VisitorTag<TreeGrowVisitor>> it; // helps growing the hungarian trees from a start node
   Graph<Directed> test_graph;
   Matrix<E> wmatrix;
   Set<int> labeled_points;
   Array<int> matching;
   const E INF = std::numeric_limits<E>::infinity();
   bool finished;
   bool inf_matching;

public:
   HungarianMethod() = default;

   explicit HungarianMethod(const Matrix<E>& input_weights)
      : dim(input_weights.cols())
      , a(dim, 0)
      , b(dim)
      , slack(dim, -1)
      , labeledColMin(dim, -1)
      , equality_subgraph(2*dim)
      , exposed_points(sequence(0,dim))
      , r(entire(exposed_points))
      , start_node(*r)
      , it(equality_subgraph, start_node)
      , matching(dim)
      , finished(false)
      , inf_matching(false)
   {

      // ensure that the weights are non-negative
      Vector<E> minvector(dim);
      auto minv = minvector.begin();
      for (auto wt = entire(rows(weights)); !wt.at_end(); ++wt, ++minv) {
         *minv = accumulate(*wt, operations::min());
      }
      weights = input_weights - repeat_col(minvector, dim);

      // Initialisation of vectors for rows and columns (dual variables)
      for (int j = 0; j < dim; ++j) {
         b[j] = accumulate(weights.col(j), operations::min());
      }

      // Build equality subgraph; it is a bipartite graph with vertex sets {0, ..., n-1} and {n, ..., 2n-1}
      for (int j = 0; j < dim; ++j) {
         for (int i = 0; i < dim; ++i) {
            if ((a[i] + b[j]) == weights[i][j]) {
               equality_subgraph.add_edge(i, dim + j);
            }
         }
      }
   }

   // removes in the directed graph G the directed edge with starting point 'start' and end point 'end' and adds the reverse edge
   void reverse_edge (int start, int end)
   {
      assert(start >= 0 && end >= 0);
      equality_subgraph.delete_edge(start,end);
      equality_subgraph.add_edge(end,start);
   }


   /* searches in the equality subgraph for an augmenting path starting with the exposed point 'start_node'
      returns true, if such a path is found, and augments the matching by reversing edges in the equality subgraph. */
   bool augment()
   {
      int node = it.node_visitor().leaf;
      int predecessor;
      // Going backwards in the hungarian tree from the leaf which was just found.
      while (node != start_node) {
         predecessor = it.node_visitor()[node];
         // modifies the equality_subgraph, so that a new matching is encoded.
         reverse_edge(predecessor, node);
         node = predecessor;
      }
      // remove the start_node from exposed_points since it is matched now.
      exposed_points -= start_node;
      // reset iterator over exposed_points
      r = entire(exposed_points);
      // reset slack to -1 since the equality subgraph has changed
      std::fill(slack.begin(), slack.end(), -1);
      std::fill(labeledColMin.begin(), labeledColMin.end(), -1);

      // initialize the growing of hungarian trees if there are still exposed points left
      if (!r.at_end()) {
         it.reset(*r);
         return false;
      } else {
         return true;
      }
   }

   // checks for every right node (dual variable b) if the corresponding slack should be adjusted with the left node 'index' (dual variable a); the lowest value is chosen for it.
   void compare_slack(int index)
   {
      E sl;
      for (int k = 0; k < b.dim(); k++) {
         sl = weights[index][k] - a[index] - b[k];
         if((sl < slack[k] || slack[k] == -1 || slack[k] == 0)) {
            if(sl > 0) { //this excludes the case where sl == 0
               slack[k] = sl;
               if (labeledColMin[k] != 0) {
                  labeledColMin[k] = sl;
               }
            }
         }
         if(sl == 0) labeledColMin[k] = 0; // records that (k +dim) is indeed a matched node
      }
   }

   // auxiliary method for compare_slack
   void change_slack(int n)
   {
      if (n == start_node) compare_slack(n);
      // sets new slack if n is a node on the right shore
      if  (n >= dim) {
         for (auto e=entire(equality_subgraph.out_edges(n)); !e.at_end(); ++e)
            compare_slack(e.to_node());
      }
   }

   // auxiliary method to find an infty entry in weights to create a minimizing permutation in the case that it is infty
   std::pair<int, int> inf_entry()
   {
      for (const auto i : sequence(0, dim)) {
         for (const auto j : sequence(0, dim)) {
            if (weights[i][j] == INF)
               return std::make_pair(i,j);
         }
      }
      throw std::runtime_error("no inf entry found but slack is inf; this happened due to an implementation error");
   }

   // here the dual variables a and b are changed in case that the equality subgraph does not contain enough edges and so a maximal matching is not perfect
   void modify()
   {
      E theta = -1;
      // theta is the lowest positive value of the weights[i][j], where i corresponds to labeled left nodes and j runs from 0 to dim-1
      for (int k = 0; k < dim; k++) {
         if ((labeledColMin[k] > 0) && (slack[k] > 0) && ((slack[k] < theta) || (theta == -1)))
            theta = slack[k];
      }

      if (theta == INF) {
         finished = true;
         inf_matching = true;
         return;
      }

      for (int k = 0; k < dim; ++k) {
         if (it.node_visitor()[k] != -1)
            a[k] = a[k] + theta;
      }

      for (int k = 0; k < dim; ++k) {
         if (it.node_visitor()[k+dim] != -1) //(labeledColMin[k] == 0 )
            b[k] = b[k] - theta;
      }

      for (int k = 0; k < dim; ++k) {
         for (int j = 0; j < dim; ++j) {
            if (a[j] + b[k] != weights[j][k]) {
               equality_subgraph.delete_edge(j, k + dim);
               equality_subgraph.delete_edge(k + dim,j);
            }
         }
      }

      for (int k = 0; k < dim; k++) {
         if (labeledColMin[k] > 0) { // slack could also be -1 -- this is a symbolic infinity
            slack[k] = slack[k] - theta;
            if (slack[k] == 0) { // at least one new edge has been created at this point
               for (int j = 0; j < dim; j++) {
                  if (a[j] + b[k] == weights[j][k]) {
                     equality_subgraph.delete_edge(j,dim+k); //ensures that there are no multiple edges
                     equality_subgraph.add_edge(j,dim+k);
                  }
               }
            }
            if (labeledColMin[k] > 0) labeledColMin[k] = slack[k];
         }
      }
      std::fill(slack.begin(), slack.end(), -1);
      std::fill(labeledColMin.begin(), labeledColMin.end(), -1);
      r = entire(exposed_points);
   }

   // initializes a bfs with start node start_node
   int growTree ()
   {
      it.reset(start_node);
      // search stops, if there is no further edge to go or a leaf of the hungarian tree is found, so that one can augment
      while (!it.at_end() && (it.node_visitor().leaf == -1) ) {
         // the vector slack is adjusted, so that it contains the lowest values in the labeled rows when it is needed in the modification step
         change_slack(*it);
         ++it;
      }
      return it.node_visitor().leaf;
   }

   // repeats the process of growing trees from exposed nodes until a perfect matching is found
   void stage ()
   {
      if (dim != 0) {
         while (!finished) {
            /* print_stuff(0); */
            while(!r.at_end()) {
               start_node = *r;
               if (growTree() != -1)
                  finished = augment();
               else ++r;
            }
            if (!finished) {
               modify();
               if (!inf_matching) {
                  it.reset(start_node);
                  it.reset(*r);
               }
            }
         }

         if(!inf_matching) {
            for (int k = 0; k < dim; k++) {
               matching[k] = equality_subgraph.in_adjacent_nodes(k).front() - dim;
            }
         } else {
            // inf matching recognized
            std::pair<int,int> trans = inf_entry();
            matching = Array<int>(sequence(0,dim));
            matching[trans.first] = trans.second;
            matching[trans.second] = trans.first;
         }
      }
   }

   void print_stuff(int flag)
   {
      cout << "a: " << a << endl;
      cout << "b: " << b << endl;
      /* cout << "slack: " << slack << endl; */
      /* cout << "labeledcolmin: " << labeledColMin << endl; */
      cout << "matching: " << matching << endl;
      cout << "exposed: " << exposed_points << endl;
      if (flag) {
         cout << "weights:\n" << weights << endl;
         cout << "eq_subgraph:\n" << equality_subgraph << endl;
      }
   }

   // TODO genericvector?
   void dynamic_stage(int index, const Vector<E>& column)
   {

      /* print_stuff(1); */

      // insert new column:
      weights.col(index) = column;
      /* for (int i = 0; i < dim; i++) */
      /*    weights[i][index] = column[i]; */

      // TODO man muss mit der non-negative-machung aufpassen! (glaube ich)
      // recalculate dual variable at 'index':
      b[index] = accumulate(column - a, operations::min());

      // adjust the eq-subgraph and exposed point(s):
      for (int i = 0; i < dim; i++) { // TODO iterate over equality_subgraph.adjacent_nodes() instead?
         equality_subgraph.delete_edge(i, dim + index);
         equality_subgraph.delete_edge(dim + index, i);
         if (a[i] + b[index] == weights[i][index])
            equality_subgraph.add_edge(i, dim + index);
         if (matching[i] == index)
            exposed_points += i;
      }

      // reset iterator over exposed point(s):
      r = entire(exposed_points);

      /* print_stuff(1); */

      // one final stage of the hungarian algo:
      finished = false;
      stage();

      /* print_stuff(1); */

   }

   Array<int> get_matching()
   {
      return matching;
   }

   Graph<Directed> get_equality_subgraph()
   {
      return equality_subgraph;
   }

   E get_value()
   {
      if (inf_matching)
         return INF;
      else
         return (accumulate(a, operations::add()) + accumulate(b, operations::add()));
   }
};

} //end graph namespace
} //end polymake namespace

#endif // POLYMAKE_HUNGARIAN_METHOD_H

// Local Variables:
// mode:C++
// c-basic-offset:3
// indent-tabs-mode:nil
// End:
// vim:shiftwidth=3:softtabstop=3:
