#ifndef DLX_DLX_HPP_
#define DLX_DLX_HPP_

#include <cstddef>
#include <functional>
#include <vector>

namespace dlx {

class Dlx {

  struct Cell {
    Cell *U, *D, *L, *R;
    std::size_t n;
    union {
      Cell* c;
      std::size_t s;
    };

    auto LR_self() -> void { L = R = this; }
    auto UD_self() -> void { U = D = this; }
    auto LR_delete() -> void { L->R = R; R->L = L; }
    auto UD_delete() -> void { U->D = D; D->U = U; }
    auto UD_restore() -> void { U->D = D->U = this; }
    auto LR_restore() -> void { L->R = R->L = this; }
    auto LR_insert(Cell* k) -> void { L = k->L; R = k; k->L = k->L->R = this; }
    auto UD_insert(Cell* k) -> void { U = k->U, D = k, k->U = k->U->D = this; }

    static Cell* ColNew() {
      auto c = new Cell;
      c->UD_self();
      c->s = 0;
      return c;
    }

    auto CoverCol() -> void {
      LR_delete();
      for (auto i = D; i != this; i = i->D) {
        for (auto j = i->R; j != i; j = j->R) {
          j->UD_delete();
          j->c->s--;
        }
      }
    }

    auto UncoverCol() -> void {
      for (auto i = U; i != this; i = i->U) {
        for (auto j = i->L; j != i; j = j->L) {
          j->c->s++;
          j->UD_restore();
        }
      }
      LR_restore();
    }

  };

  std::vector<Cell*> ctab_;
  std::vector<Cell*> rtab_;
  Cell* root_;

  auto AddCol() -> void;
  auto AddRow() -> void;
  auto AllocCol(std::size_t n) -> void;
  auto AllocRow(std::size_t n) -> void;

public:
  Dlx();
  ~Dlx();

  // Returns number of rows.
  auto Rows() const -> std::size_t;

  // Returns number of columns.
  auto Cols() const -> std::size_t;

  // Places a 1 in the given row and column.
  // Increases the number of rows and columns if necessary.
  auto Set(std::size_t row, std::size_t col) -> void;

// Marks a column as optional: a solution need not cover the given column,
// but it still must respect the constraints it entails.
  auto MarkOptional(std::size_t col) -> void;

  // Removes a row from consideration. Returns 0 on success, -1 otherwise.
  // Should only be called after all dlx_set() calls.
  auto RemoveRow(std::size_t row) -> int;

  // Picks a row to be part of the solution. Returns 0 on success, -1 otherwise.
  // Should only be called after all dlx_set() calls and dlx_remove_row() calls.
  // TODO: Check the row can be legally chosen.
  auto PickRow(std::size_t row) -> int;

  auto Solve(
    std::function<void(std::size_t, std::size_t, std::size_t)> try_cb,
    std::function<void()> undo_cb,
    std::function<void()> found_cb,
    std::function<void(std::size_t)> stuck_cb) -> void;
};

auto ForallCover(Dlx& dlx, auto cb) {
  std::vector<std::size_t> sol;
  dlx.Solve(
    [&](std::size_t c, std::size_t s, std::size_t r) { sol.push_back(r); },
    [&](){ sol.pop_back(); },
    [&](){ cb(sol); },
    [](std::size_t){}
  );
}

} // namespace

#endif