# Copyright 2021 Hakan Kjellerstrand hakank@gmail.com
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
  A puzzle in OR-tools CP-SAT Solver.

  From "God plays dice"
  "A puzzle"
  http://gottwurfelt.wordpress.com/2012/02/22/a-puzzle/
  And the sequel "Answer to a puzzle"
  http://gottwurfelt.wordpress.com/2012/02/24/an-answer-to-a-puzzle/

  This problem instance was taken from the latter blog post.
  '''
  8809 = 6
  7111 = 0
  2172 = 0
  6666 = 4
  1111 = 0
  3213 = 0
  7662 = 2
  9312 = 1
  0000 = 4
  2222 = 0
  3333 = 0
  5555 = 0
  8193 = 3
  8096 = 5
  7777 = 0
  9999 = 4
  7756 = 1
  6855 = 3
  9881 = 5
  5531 = 0

  2581 = ?
  '''

  Note: 
  This instance yields 10 solutions, since x4 is not 
  restricted in the constraints. 
  All solutions has x assigned to the correct result. 


  The problem stated in "A puzzle" (the first blog post)
  http://gottwurfelt.wordpress.com/2012/02/22/a-puzzle/
  is
  '''
  8809 = 6
  7662 = 2
  9312 = 1
  8193 = 3
  8096 = 5
  7756 = 1
  6855 = 3
  9881 = 5

  2581 = ?
  '''
  This problem instance - using the same principle - yields 
  two different solutions of x, one is the same (correct) as 
  for the above problem instance, and one is not.
  It's because here both x4 and x1 are underdefined.


  This model was written by Hakan Kjellerstrand (hakank@gmail.com)
  See also my OR-tools page: http://hakank.org/or_tools/
"""
from __future__ import print_function
from ortools.sat.python import cp_model as cp
from cp_sat_utils import SimpleSolutionPrinter2


def puzzle(problem=1):
  
  print("Problem", problem)
    
  model = cp.CpModel()

  n = 10

  # variables
  all = [model.NewIntVar(0,9,f"all[{i}]") for i in range(n)]
  [x0,x1,x2,x3,_x4,x5,x6,x7,x8,x9] = all

  x = model.NewIntVar(0,9,"x")

  if problem==1:
    model.Add(x8+x8+x0+x9 == 6)
    model.Add(x7+x1+x1+x1 == 0)
    model.Add(x2+x1+x7+x2 == 0)
    model.Add(x6+x6+x6+x6 == 4)
    model.Add(x1+x1+x1+x1 == 0)
    model.Add(x3+x2+x1+x3 == 0)
    model.Add(x7+x6+x6+x2 == 2)
    model.Add(x9+x3+x1+x2 == 1)
    model.Add(x0+x0+x0+x0 == 4)
    model.Add(x2+x2+x2+x2 == 0)
    model.Add(x3+x3+x3+x3 == 0)
    model.Add(x5+x5+x5+x5 == 0)
    model.Add(x8+x1+x9+x3 == 3)
    model.Add(x8+x0+x9+x6 == 5)
    model.Add(x7+x7+x7+x7 == 0)
    model.Add(x9+x9+x9+x9 == 4)
    model.Add(x7+x7+x5+x6 == 1)
    model.Add(x6+x8+x5+x5 == 3)
    model.Add(x9+x8+x8+x1 == 5)
    model.Add(x5+x5+x3+x1 == 0)
    
    model.Add(x2+x5+x8+x1 == x)

  else:
    # This yields two different values for x
    model.Add(x8+x8+x0+x9 == 6)
    model.Add(x7+x6+x6+x2 == 2)
    model.Add(x9+x3+x1+x2 == 1)
    model.Add(x8+x1+x9+x3 == 3)
    model.Add(x8+x0+x9+x6 == 5)
    model.Add(x7+x7+x5+x6 == 1)
    model.Add(x6+x8+x5+x5 == 3)
    model.Add(x9+x8+x8+x1 == 5)
    
    model.Add(x2+x5+x8+x1 == x)

  solver = cp.CpSolver()
  solution_printer = SimpleSolutionPrinter2([x,all])
  status = solver.SearchForAllSolutions(model,solution_printer)

  if not status in [cp.OPTIMAL,cp.FEASIBLE]:
    print("No solution!")

  print()
  print("NumSolutions:", solution_printer.SolutionCount())  
  print("NumConflicts:", solver.NumConflicts())
  print("NumBranches:", solver.NumBranches())
  print("WallTime:", solver.WallTime())


if __name__ == '__main__':
  puzzle(1)
  print()
  puzzle(2)