#
# Survo puzzle in Elixir.
#
# http://en.wikipedia.org/wiki/Survo_Puzzle
# """
# Survo puzzle is a kind of logic puzzle presented (in April 2006) and studied
# by Seppo Mustonen. The name of the puzzle is associated to Mustonen's
# Survo system which is a general environment for statistical computing and
# related areas.

# In a Survo puzzle the task is to fill an m * n table by integers 1,2,...,m*n so
# that each of these numbers appears only once and their row and column sums are
# equal to integers given on the bottom and the right side of the table.
# Often some of the integers are given readily in the table in order to
# guarantee uniqueness of the solution and/or for making the task easier.
# """
#
# See also
# http://www.survo.fi/english/index.html
# http://www.survo.fi/puzzles/index.html
#
#
# This program was created by Hakan Kjellerstrand, hakank@gmail.com
# See also my Elixir page: http://www.hakank.org/elxir/
#
defmodule CPSolver.Examples.SurvoPuzzle do
  alias CPSolver.IntVariable
  alias CPSolver.Constraint.Sum
  alias CPSolver.Constraint.AllDifferent
  alias CPSolver.Model

  @doc """
  mat_at(m,i,j)

  Returns the value (`i`,`j`) of the 2d matrix `mat`.

  ##Examples##

       iex> [[1,2,3],[4,5,6],[7,8,9]] |> mat_at(1,2)
       6
  """
  def mat_at(m, i, j) do
    m |> Enum.at(i) |> Enum.at(j)
  end

  @doc """
  transpose(m)

  Returns a transposed version of `m`.

  ##Example##

       iex> [[1,2,3],[4,5,6],[7,8,9]] |> transpose
       [[1, 4, 7], [2, 5, 8], [3, 6, 9]]
  """
  def transpose(m) do
    Enum.zip_with(m, &Function.identity/1)
  end

  def puzzle(1) do
    rowsums = [30, 18, 30]
    colsums = [27, 16, 10, 25]
    # 0 is unknown -> to be decided
    problem = [[0, 6, 0, 0], [8, 0, 0, 0], [0, 0, 3, 0]]
    [rowsums, colsums, problem]
  end

  def main() do
    [rowsums, colsums, problem] = puzzle(1)

    rows = length(rowsums)
    cols = length(colsums)
    dom = 1..(rows * cols)

    # Decision variables
    x_matrix =
      for i <- 0..(rows - 1) do
        for j <- 0..(cols - 1) do
          v = mat_at(problem, i, j)

          if v > 0 do
            IntVariable.new(v, name: "x[#{i},#{j}]")
          else
            IntVariable.new(dom, name: "x[#{i},#{j}]")
          end
        end
      end

    # Row constraints
    row_constraints =
      for {s, row} <- Enum.zip(rowsums, x_matrix) do
        Sum.new(s, row)
      end

    # Column constraints
    col_constraints =
      for {s, row} <- Enum.zip(colsums, transpose(x_matrix)) do
        Sum.new(s, row)
      end

    x = List.flatten(x_matrix)
    constraints = [AllDifferent.new(x) | row_constraints ++ col_constraints]

    model =
      Model.new(
        x,
        constraints
      )

    {:ok, _result} =
      CPSolver.solve_sync(model,
        search: {:input_order, :indomain_random},
        # stop_on: {:max_solutions, 1},
        timeout: :infinity,
        space_threads: 12
      )
  end

  def check_solution(solution, row_sums, col_sums) do
    table_content =
      Enum.take(solution, length(row_sums) * length(col_sums))

    table = Enum.chunk_every(table_content, length(col_sums))

    row_sums_correct? =
      table
      |> Enum.zip(row_sums)
      |> Enum.all?(fn {row, sum} -> sum == Enum.sum(row) end)

    col_sums_correct? =
      table
      |> transpose()
      |> Enum.zip(col_sums)
      |> Enum.all?(fn {col, sum} -> sum == Enum.sum(col) end)

    all_different? =
      length(row_sums) * length(col_sums) == MapSet.new(table_content) |> MapSet.size()

    col_sums_correct? && row_sums_correct? && all_different?
  end
end