text: #!/usr/bin/python -u
  # -*- coding: latin-1 -*-
  # 
  # Magic square in Z3
  #
  # Variant using an IntVector of size n*n instead of a could of Ints
  # (compare with magic_square.py).
  # 
  # This Z3 model was written by Hakan Kjellerstrand (hakank@gmail.com)
  # See also my Z3 page: http://hakank.org/z3/
  # 
  # 
  from z3_utils_hakank import *
  import time
  
  def print_square(mod,x,n):
      for i in range(n):
          for j in range(n):
              print("%2s" % mod.eval(x[(i, j)]),end=" ")
          print()
      
  
  def magic_square(n,all=0):
      start = time.clock()
      
      sol = Solver()
      nn = n*n
      # Here we define and IntVector instead
      # But almost everything else is the same as magic_square.py
      xi = IntVector("xxxz3", nn)
      
      x = {}
      for i in range(n):
          for j in range(n):
              x[(i, j)] = xi[(i-1)*n+j] # connect xi and x
              sol.add(x[(i,j)] >= 1,  x[(i,j)] <= n*n)
          
      s = Int("s")
      sol.add(s >= 1, s <= n*n*n)
      sol.add(Distinct([xi[i] for i in range(nn)])) # flattened
      [sol.add(Sum([x[(i, j)] for j in range(n)]) == s) for i in range(n)]
      [sol.add(Sum([x[(i, j)] for i in range(n)]) == s) for j in range(n)]
  
      sol.add(Sum([x[(i, i)] for i in range(n)]) == s)  # diag 1
      sol.add(Sum([x[(i, n - i - 1)] for i in range(n)]) == s)  # diag 2
  
      # symmetry breaking
      # sol.add(x[(0,0)] == 1)
  
      print(sol.check())
      end = time.clock()
      value = end - start
      print("Time to sat: ", value)
  
      if sol.check() == sat:
          mod = sol.model()
          print("s:", mod.eval(s))
          # print_square(mod,x,n)
      else:
          print("No solution found!")
  
      end = time.clock()
      value = end - start
      print("Time all: ", value)
  
  
  for n in [3,4,5]:
      print("Testing ", n)
      magic_square(n)
      print()
  