///////////////////////////////////////////////////////////////////////////////
// Hungarian.h: Header file for Class HungarianAlgorithm.
// 
// This is a C++ wrapper with slight modification of a hungarian algorithm implementation by Markus Buehren.
// The original implementation is a few mex-functions for use in MATLAB, found here:
// http://www.mathworks.com/matlabcentral/fileexchange/6543-functions-for-the-rectangular-assignment-problem
// 
// Both this code and the orignal code are published under the BSD license.
// by Cong Ma, 2016
// 
// This code was then modified by Justin Silverman, 2019, as a single header library
// for use in R applications. 
// BSD lisence is maintained

#ifndef HUNGARIAN_H
#define HUNGARIAN_H

//#include <iostream>
#include <vector>
#include <stdlib.h>
#include <cfloat> // for DBL_MAX
#include <cmath>  // for fabs()
#include <stdexcept>

using namespace std;


class HungarianAlgorithm
{
public:
  HungarianAlgorithm(){}
  ~HungarianAlgorithm(){}
  
  //********************************************************//
  // A single function wrapper for solving assignment problem.
  //********************************************************//
  double Solve(vector <vector<double> >& DistMatrix, vector<int>& Assignment){
    unsigned int nRows = DistMatrix.size();
    unsigned int nCols = DistMatrix[0].size();
    
    double *distMatrixIn = new double[nRows * nCols];
    int *assignment = new int[nRows];
    double cost = 0.0;
    
    // Fill in the distMatrixIn. Mind the index is "i + nRows * j".
    // Here the cost matrix of size MxN is defined as a double precision array of N*M elements. 
    // In the solving functions matrices are seen to be saved MATLAB-internally in row-order.
    // (i.e. the matrix [1 2; 3 4] will be stored as a vector [1 3 2 4], NOT [1 2 3 4]).
    for (unsigned int i = 0; i < nRows; i++)
      for (unsigned int j = 0; j < nCols; j++)
        distMatrixIn[i + nRows * j] = DistMatrix[i][j];
    
    // call solving function
    assignmentoptimal(assignment, &cost, distMatrixIn, nRows, nCols);
    
    Assignment.clear();
    for (unsigned int r = 0; r < nRows; r++)
      Assignment.push_back(assignment[r]);
    
    delete[] distMatrixIn;
    delete[] assignment;
    return cost;
  }
  
private:
  
  //********************************************************//
  // Solve optimal solution for assignment problem using Munkres algorithm, also known as Hungarian Algorithm.
  //********************************************************//
  void assignmentoptimal(int *assignment, double *cost, double *distMatrixIn, int nOfRows, int nOfColumns){
    double *distMatrix, *distMatrixTemp, *distMatrixEnd, *columnEnd, value, minValue;
    bool *coveredColumns, *coveredRows, *starMatrix, *newStarMatrix, *primeMatrix;
    int nOfElements, minDim, row, col;
    
    /* initialization */
    *cost = 0;
    for (row = 0; row<nOfRows; row++)
      assignment[row] = -1;
    
    /* generate working copy of distance Matrix */
    /* check if all matrix elements are positive */
    nOfElements = nOfRows * nOfColumns;
    distMatrix = (double *)malloc(nOfElements * sizeof(double));
    distMatrixEnd = distMatrix + nOfElements;
    
    for (row = 0; row<nOfElements; row++)
    {
      value = distMatrixIn[row];
      // if (value < 0)
      //   throw std::runtime_error("All matrix elements have to be non-negative.");
      //   //cerr << "All matrix elements have to be non-negative."  << endl;
      distMatrix[row] = value;
    }
    
    
    /* memory allocation */
    coveredColumns = (bool *)calloc(nOfColumns, sizeof(bool));
    coveredRows = (bool *)calloc(nOfRows, sizeof(bool));
    starMatrix = (bool *)calloc(nOfElements, sizeof(bool));
    primeMatrix = (bool *)calloc(nOfElements, sizeof(bool));
    newStarMatrix = (bool *)calloc(nOfElements, sizeof(bool)); /* used in step4 */
    
    /* preliminary steps */
    if (nOfRows <= nOfColumns)
    {
      minDim = nOfRows;
      
      for (row = 0; row<nOfRows; row++)
      {
        /* find the smallest element in the row */
        distMatrixTemp = distMatrix + row;
        minValue = *distMatrixTemp;
        distMatrixTemp += nOfRows;
        while (distMatrixTemp < distMatrixEnd)
        {
          value = *distMatrixTemp;
          if (value < minValue)
            minValue = value;
          distMatrixTemp += nOfRows;
        }
        
        /* subtract the smallest element from each element of the row */
        distMatrixTemp = distMatrix + row;
        while (distMatrixTemp < distMatrixEnd)
        {
          *distMatrixTemp -= minValue;
          distMatrixTemp += nOfRows;
        }
      }
      
      /* Steps 1 and 2a */
      for (row = 0; row<nOfRows; row++)
        for (col = 0; col<nOfColumns; col++)
          if (fabs(distMatrix[row + nOfRows*col]) < DBL_EPSILON)
            if (!coveredColumns[col])
            {
              starMatrix[row + nOfRows*col] = true;
              coveredColumns[col] = true;
              break;
            }
    }
    else /* if(nOfRows > nOfColumns) */
    {
      minDim = nOfColumns;
      
      for (col = 0; col<nOfColumns; col++)
      {
        /* find the smallest element in the column */
        distMatrixTemp = distMatrix + nOfRows*col;
        columnEnd = distMatrixTemp + nOfRows;
        
        minValue = *distMatrixTemp++;
        while (distMatrixTemp < columnEnd)
        {
          value = *distMatrixTemp++;
          if (value < minValue)
            minValue = value;
        }
        
        /* subtract the smallest element from each element of the column */
        distMatrixTemp = distMatrix + nOfRows*col;
        while (distMatrixTemp < columnEnd)
          *distMatrixTemp++ -= minValue;
      }
      
      /* Steps 1 and 2a */
      for (col = 0; col<nOfColumns; col++)
        for (row = 0; row<nOfRows; row++)
          if (fabs(distMatrix[row + nOfRows*col]) < DBL_EPSILON)
            if (!coveredRows[row])
            {
              starMatrix[row + nOfRows*col] = true;
              coveredColumns[col] = true;
              coveredRows[row] = true;
              break;
            }
            for (row = 0; row<nOfRows; row++)
              coveredRows[row] = false;
            
    }
    
    /* move to step 2b */
    step2b(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
    
    /* compute cost and remove invalid assignments */
    computeassignmentcost(assignment, cost, distMatrixIn, nOfRows);
    
    /* free allocated memory */
    free(distMatrix);
    free(coveredColumns);
    free(coveredRows);
    free(starMatrix);
    free(primeMatrix);
    free(newStarMatrix);
    
    return;
  }
  
  //********************************************************/
  void buildassignmentvector(int *assignment, bool *starMatrix, int nOfRows, int nOfColumns){
    int row, col;
    
    for (row = 0; row<nOfRows; row++)
      for (col = 0; col<nOfColumns; col++)
        if (starMatrix[row + nOfRows*col])
        { 
          assignment[row] = col;
          break;
        }
  }
  
  //********************************************************//
  void computeassignmentcost(int *assignment, double *cost, double *distMatrix, int nOfRows){
    int row, col;
    
    for (row = 0; row<nOfRows; row++)
    {
      col = assignment[row];
      if (col >= 0)
        *cost += distMatrix[row + nOfRows*col];
    }
  }
  
  
  //********************************************************//
  void step2a(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim){
    bool *starMatrixTemp, *columnEnd;
    int col;
    
    /* cover every column containing a starred zero */
    for (col = 0; col<nOfColumns; col++)
    {
      starMatrixTemp = starMatrix + nOfRows*col;
      columnEnd = starMatrixTemp + nOfRows;
      while (starMatrixTemp < columnEnd){
        if (*starMatrixTemp++)
        {
          coveredColumns[col] = true;
          break;
        }
      }
    }
    
    /* move to step 3 */
    step2b(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
  }
  
  //********************************************************//
  void step2b(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim){
    int col, nOfCoveredColumns;
    
    /* count covered columns */
    nOfCoveredColumns = 0;
    for (col = 0; col<nOfColumns; col++)
      if (coveredColumns[col])
        nOfCoveredColumns++;
      
      if (nOfCoveredColumns == minDim)
      {
        /* algorithm finished */
        buildassignmentvector(assignment, starMatrix, nOfRows, nOfColumns);
      }
      else
      {
        /* move to step 3 */
        step3(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
      }
  }
  
  //********************************************************//
  void step3(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim){
    bool zerosFound;
    int row, col, starCol;
    
    zerosFound = true;
    while (zerosFound)
    {
      zerosFound = false;
      for (col = 0; col<nOfColumns; col++)
        if (!coveredColumns[col])
          for (row = 0; row<nOfRows; row++)
            if ((!coveredRows[row]) && (fabs(distMatrix[row + nOfRows*col]) < DBL_EPSILON))
            {
              /* prime zero */
              primeMatrix[row + nOfRows*col] = true;
              
              /* find starred zero in current row */
              for (starCol = 0; starCol<nOfColumns; starCol++)
                if (starMatrix[row + nOfRows*starCol])
                  break;
                
                if (starCol == nOfColumns) /* no starred zero found */
                {
                  /* move to step 4 */
                  step4(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim, row, col);
                  return;
                }
                else
                {
                  coveredRows[row] = true;
                  coveredColumns[starCol] = false;
                  zerosFound = true;
                  break;
                }
            }
    }
    
    /* move to step 5 */
    step5(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
  }
  
  //********************************************************//
  void step4(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim, int row, int col){
    int n, starRow, starCol, primeRow, primeCol;
    int nOfElements = nOfRows*nOfColumns;
    
    /* generate temporary copy of starMatrix */
    for (n = 0; n<nOfElements; n++)
      newStarMatrix[n] = starMatrix[n];
    
    /* star current zero */
    newStarMatrix[row + nOfRows*col] = true;
    
    /* find starred zero in current column */
    starCol = col;
    for (starRow = 0; starRow<nOfRows; starRow++)
      if (starMatrix[starRow + nOfRows*starCol])
        break;
      
      while (starRow<nOfRows)
      {
        /* unstar the starred zero */
        newStarMatrix[starRow + nOfRows*starCol] = false;
        
        /* find primed zero in current row */
        primeRow = starRow;
        for (primeCol = 0; primeCol<nOfColumns; primeCol++)
          if (primeMatrix[primeRow + nOfRows*primeCol])
            break;
          
          /* star the primed zero */
          newStarMatrix[primeRow + nOfRows*primeCol] = true;
          
          /* find starred zero in current column */
          starCol = primeCol;
          for (starRow = 0; starRow<nOfRows; starRow++)
            if (starMatrix[starRow + nOfRows*starCol])
              break;
      }
      
      /* use temporary copy as new starMatrix */
      /* delete all primes, uncover all rows */
      for (n = 0; n<nOfElements; n++)
      {
        primeMatrix[n] = false;
        starMatrix[n] = newStarMatrix[n];
      }
      for (n = 0; n<nOfRows; n++)
        coveredRows[n] = false;
      
      /* move to step 2a */
      step2a(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
  }
  
  //********************************************************//
  void step5(int *assignment, double *distMatrix, bool *starMatrix, bool *newStarMatrix, bool *primeMatrix, bool *coveredColumns, bool *coveredRows, int nOfRows, int nOfColumns, int minDim){
    double h, value;
    int row, col;
    
    /* find smallest uncovered element h */
    h = DBL_MAX;
    for (row = 0; row<nOfRows; row++)
      if (!coveredRows[row])
        for (col = 0; col<nOfColumns; col++)
          if (!coveredColumns[col])
          {
            value = distMatrix[row + nOfRows*col];
            if (value < h)
              h = value;
          }
          
          /* add h to each covered row */
          for (row = 0; row<nOfRows; row++)
            if (coveredRows[row])
              for (col = 0; col<nOfColumns; col++)
                distMatrix[row + nOfRows*col] += h;
            
            /* subtract h from each uncovered column */
            for (col = 0; col<nOfColumns; col++)
              if (!coveredColumns[col])
                for (row = 0; row<nOfRows; row++)
                  distMatrix[row + nOfRows*col] -= h;
              
              /* move to step 3 */
              step3(assignment, distMatrix, starMatrix, newStarMatrix, primeMatrix, coveredColumns, coveredRows, nOfRows, nOfColumns, minDim);
  }
  
};


#endif
