/*
 * testSudoku.cpp
 * @brief develop code for Sudoku CSP solver
 * @date Jan 29, 2012
 * @author Frank Dellaert
 */

#include <CppUnitLite/TestHarness.h>
#include <gtsam_unstable/discrete/CSP.h>

#include <boost/assign/std/map.hpp>
using boost::assign::insert;
#include <stdarg.h>

#include <iostream>
#include <sstream>

using namespace std;
using namespace gtsam;

#define PRINT false

class Sudoku : public CSP {
  /// sudoku size
  size_t n_;

  /// discrete keys
  typedef std::pair<size_t, size_t> IJ;
  std::map<IJ, DiscreteKey> dkeys_;

 public:
  /// return DiscreteKey for cell(i,j)
  const DiscreteKey& dkey(size_t i, size_t j) const {
    return dkeys_.at(IJ(i, j));
  }

  /// return Key for cell(i,j)
  Key key(size_t i, size_t j) const { return dkey(i, j).first; }

  /// Constructor
  Sudoku(size_t n, ...) : n_(n) {
    // Create variables, ordering, and unary constraints
    va_list ap;
    va_start(ap, n);
    Key k = 0;
    for (size_t i = 0; i < n; ++i) {
      for (size_t j = 0; j < n; ++j, ++k) {
        // create the key
        IJ ij(i, j);
        dkeys_[ij] = DiscreteKey(k, n);
        // get the unary constraint, if any
        int value = va_arg(ap, int);
        // cout << value << " ";
        if (value != 0) addSingleValue(dkeys_[ij], value - 1);
      }
      // cout << endl;
    }
    va_end(ap);

    // add row constraints
    for (size_t i = 0; i < n; i++) {
      DiscreteKeys dkeys;
      for (size_t j = 0; j < n; j++) dkeys += dkey(i, j);
      addAllDiff(dkeys);
    }

    // add col constraints
    for (size_t j = 0; j < n; j++) {
      DiscreteKeys dkeys;
      for (size_t i = 0; i < n; i++) dkeys += dkey(i, j);
      addAllDiff(dkeys);
    }

    // add box constraints
    size_t N = (size_t)sqrt(double(n)), i0 = 0;
    for (size_t I = 0; I < N; I++) {
      size_t j0 = 0;
      for (size_t J = 0; J < N; J++) {
        // Box I,J
        DiscreteKeys dkeys;
        for (size_t i = i0; i < i0 + N; i++)
          for (size_t j = j0; j < j0 + N; j++) dkeys += dkey(i, j);
        addAllDiff(dkeys);
        j0 += N;
      }
      i0 += N;
    }
  }

  /// Print readable form of assignment
  void printAssignment(DiscreteFactor::sharedValues assignment) const {
    for (size_t i = 0; i < n_; i++) {
      for (size_t j = 0; j < n_; j++) {
        Key k = key(i, j);
        cout << 1 + assignment->at(k) << " ";
      }
      cout << endl;
    }
  }

  /// solve and print solution
  void printSolution() {
    DiscreteFactor::sharedValues MPE = optimalAssignment();
    printAssignment(MPE);
  }
};

/* ************************************************************************* */
TEST_UNSAFE(Sudoku, small) {
  Sudoku csp(4, 1, 0, 0, 4, 0, 0, 0, 0,

             4, 0, 2, 0, 0, 1, 0, 0);

  // Do BP
  csp.runArcConsistency(4, 10, PRINT);

  // optimize and check
  CSP::sharedValues solution = csp.optimalAssignment();
  CSP::Values expected;
  insert(expected)(csp.key(0, 0), 0)(csp.key(0, 1), 1)(csp.key(0, 2), 2)(
      csp.key(0, 3), 3)(csp.key(1, 0), 2)(csp.key(1, 1), 3)(csp.key(1, 2), 0)(
      csp.key(1, 3), 1)(csp.key(2, 0), 3)(csp.key(2, 1), 2)(csp.key(2, 2), 1)(
      csp.key(2, 3), 0)(csp.key(3, 0), 1)(csp.key(3, 1), 0)(csp.key(3, 2), 3)(
      csp.key(3, 3), 2);
  EXPECT(assert_equal(expected, *solution));
  // csp.printAssignment(solution);
}

/* ************************************************************************* */
TEST_UNSAFE(Sudoku, easy) {
  Sudoku sudoku(9, 0, 0, 5, 0, 9, 0, 0, 0, 1, 0, 0, 0, 0, 0, 2, 0, 7, 3, 7, 6,
                0, 0, 0, 8, 2, 0, 0,

                0, 1, 2, 0, 0, 9, 0, 0, 4, 0, 0, 0, 2, 0, 3, 0, 0, 0, 3, 0, 0,
                1, 0, 0, 9, 6, 0,

                0, 0, 1, 9, 0, 0, 0, 5, 8, 9, 7, 0, 5, 0, 0, 0, 0, 0, 5, 0, 0,
                0, 3, 0, 7, 0, 0);

  // Do BP
  sudoku.runArcConsistency(4, 10, PRINT);

  // sudoku.printSolution(); // don't do it
}

/* ************************************************************************* */
TEST_UNSAFE(Sudoku, extreme) {
  Sudoku sudoku(9, 0, 0, 9, 7, 4, 8, 0, 0, 0, 7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2,
                0, 1, 0, 9, 0, 0, 0,

                0, 0, 7, 0, 0, 0, 2, 4, 0, 0, 6, 4, 0, 1, 0, 5, 9, 0, 0, 9, 8,
                0, 0, 0, 3, 0, 0,

                0, 0, 0, 8, 0, 3, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0,
                2, 7, 5, 9, 0, 0);

  // Do BP
  sudoku.runArcConsistency(9, 10, PRINT);

#ifdef METIS
  VariableIndexOrdered index(sudoku);
  index.print("index");
  ofstream os("/Users/dellaert/src/hmetis-1.5-osx-i686/extreme-dual.txt");
  index.outputMetisFormat(os);
#endif

  // sudoku.printSolution(); // don't do it
}

/* ************************************************************************* */
TEST_UNSAFE(Sudoku, AJC_3star_Feb8_2012) {
  Sudoku sudoku(9, 9, 5, 0, 0, 0, 6, 0, 0, 0, 0, 8, 4, 0, 7, 0, 0, 0, 0, 6, 2,
                0, 5, 0, 0, 4, 0, 0,

                0, 0, 0, 2, 9, 0, 6, 0, 0, 0, 9, 0, 0, 0, 0, 0, 2, 0, 0, 0, 2,
                0, 6, 3, 0, 0, 0,

                0, 0, 9, 0, 0, 7, 0, 6, 8, 0, 0, 0, 0, 3, 0, 2, 9, 0, 0, 0, 0,
                1, 0, 0, 0, 3, 7);

  // Do BP
  sudoku.runArcConsistency(9, 10, PRINT);

  // sudoku.printSolution(); // don't do it
}

/* ************************************************************************* */
int main() {
  TestResult tr;
  return TestRegistry::runAllTests(tr);
}
/* ************************************************************************* */