from z3 import *

# 3x3-sudoku instance. '0' means empty
instance = ((0,0,0,0,9,4,0,3,0),
            (0,0,0,5,1,0,0,0,7),
            (0,8,9,0,0,0,0,4,0),
            (0,0,0,0,0,0,2,0,8),
            (0,6,0,2,0,1,0,5,0),
            (1,0,2,0,0,0,0,0,0),
            (0,7,0,0,0,0,5,2,0),
            (9,0,0,0,6,5,0,0,0),
            (0,4,0,9,7,0,0,0,0))

i = 0
def fresh_variable():
  global i
  i = i + 1
  return Int("x" + str(i))

Var = [ [ fresh_variable() for _ in range(9) ]
        for _ in range(9) ]

the_instance_rule = [
    Var[i][j] == instance[i][j]
   for i in range(9)
  for j in range(9)
  if instance[i][j] != 0
  ]

one_through_nine_rule = [
    And(0 < Var[i][j], Var[i][j] < 10)
  for i in range(9)
  for j in range(9)
  ]

row_rule = [
    Distinct( Var[i] ) for i in range(9)
  ]

column_rule = [
    Distinct( [ Var[i][j] for i in range(9) ] )
  for j in range(9)
  ]

small_square_rule = [
    Distinct( [ Var[i0*3 + i][j0*3 + j]
                for i in range(3)
                for j in range(3)])
              for i0 in range(3)
              for j0 in range(3)
  ]

def run(constraints):
  s = Solver()
  s.add(constraints)
  if s.check() == sat:
      return s.model()
  else:
      raise Exception("Could not solve your instance")

model = run(the_instance_rule + one_through_nine_rule + row_rule + column_rule + small_square_rule)

print_matrix(
  [ [ model.evaluate(Var[i][j]) for j in range(9) ]
    for i in range(9)]
  )