"""
Python program to solve a logic puzzle.
Full blog post: https://fliegendewurst.github.io/z3-logic-puzzle-solving.html
Puzzle description: https://www.zeit.de/zeit-magazin/2022/32/logelei
🄯 FliegendeWurst 2022
This work is licensed under a Creative Commons Attribution 4.0 International License.
http://creativecommons.org/licenses/by/4.0/
"""

from z3 import *
import math

s = Solver()

# n*n grid
n = 6

grid = [[FreshInt() for x in range(n)] for y in range(n)]

def print_solution(m, print_grid=True):
    for k in range(num_sums):
        print(" " * (3*num_sums + 2), end="")
        for x in range(n):
            val = m[sums_v[x][k]].as_long()
            if val == -1:
                print("  ", end=" ")
            else:
                print(f"{val: >2}", end=" ")
        print()
    print()
    print(" " * (3*num_sums + 1), end="")
    print("-" * (3*n))
    y = 0
    for row in grid:
        for sum_h in sums_h[y]:
            val = m[sum_h].as_long()
            if val == -1:
                print("  ", end=" ")
            else:
                print(f"{val: >2}", end=" ")
        print("|", end=" ")
        x = 0
        for cell in row:
            val = m[cell].as_long()
            if val == 0 or not print_grid:
                val = "  "
            print(f"{val: >2}", end=" ")
            x += 1
        print()
        y += 1

def distinct_if_nonzero(x):
    for i in range(len(x)):
        for j in range(i+1, len(x)):
            cell_i = x[i]
            cell_j = x[j]
            s.add(If(And(cell_i != 0, cell_j != 0), cell_i != cell_j, True))

# empty grid cell = 0
# other numbers have to be in [1,2,3,4,5]
# no duplicate digits in any row / column
for row in grid:
    for cell in row:
        s.add(cell >= 0)
        s.add(cell <= 5)
    distinct_if_nonzero(row)

for column in [[grid[y][x] for y in range(n)] for x in range(n)]:
    distinct_if_nonzero(column)

# each row and column has at most math.ceil(n/2) segments
# if there are less, they are indicated by -1
num_sums = math.ceil(n / 2)
sums_h = [[FreshInt() for i in range(num_sums)] for y in range(n)]
sums_v = [[FreshInt() for i in range(num_sums)] for x in range(n)]
# for each pattern of used digits in a row, determine the sums
# iterate over each row of the grid
for y in range(n):
    # iterate over each binary pattern
    for p in range(2**n):
        # match condition of this pattern
        mc = True
	# will contain a list of variables for each segment
        segments = []
        start_next_segment = True
        # iterate over each grid cell
        for x in range(n):
            # check that binary digit in the pattern
            if (p >> x) & 1 == 1:
                mc = And(mc, grid[y][x] != 0)
                if start_next_segment:
                    segments.append([grid[y][x]])
                    start_next_segment = False
                else:
                    segments[-1].append(grid[y][x])
            else:
                mc = And(mc, grid[y][x] == 0)
                start_next_segment = True
        # specify sums value
        for k in range(num_sums):
            if k < len(segments):
                off_by_one = Or(
                    sums_h[y][k] == sum(segments[k]) + 1,
                    sums_h[y][k] == sum(segments[k]) - 1
                )
                s.add(If(mc, off_by_one, True))
            else:
                s.add(If(mc, sums_h[y][k] == -1, True))

# for each pattern of used digits in a columns, determine the sums
for x in range(n):
    for p in range(2**n):
        # match condition
        mc = True
        segments = []
        start_next_segment = True
        for y in range(n):
            if (p >> y) & 1 == 1:
                mc = And(mc, grid[y][x] != 0)
                if start_next_segment:
                    segments.append([grid[y][x]])
                    start_next_segment = False
                else:
                    segments[-1].append(grid[y][x])
            else:
                mc = And(mc, grid[y][x] == 0)
                start_next_segment = True
        # specify sums value
        #print(f"{p:b}"[::-1], segments)
        for k in range(num_sums):
            if k < len(segments):
                off_by_one = Or(
                    sums_v[x][k] == sum(segments[k]) + 1,
                    sums_v[x][k] == sum(segments[k]) - 1
                )
                s.add(If(mc, off_by_one, True))
            else:
                s.add(If(mc, sums_v[x][k] == -1, True))

res = s.check()
print(res)

def add_spec(spec_h, spec_v):
    for a, b in zip(spec_h, sums_h):
        for x, y in zip(a, b):
            s.add(x == y)
    for a, b in zip(spec_v, sums_v):
        for x, y in zip(a, b):
            s.add(x == y)
spec_h = [
    [2, 13, -1],
    [5, 3, 6],
    [8, 7, -1],
    [4, 4, 1],
    [5, 5, -1],
    [9, 0, -1]
]

spec_v = [
    [11, -1, -1],
    [4, 5, -1],
    [7, 7, -1],
    [2, 7, 3],
    [4, 4, 5],
    [10, 2, -1]
]
s.push()
add_spec(spec_h, spec_v)

res = s.check()
print(res)

while res == sat:
    m = s.model()
    print_solution(m)
    # find another solution, if possible
    c = False
    for row in grid:
        for cell in row:
            c = Or(c, cell != m[cell].as_long())
    s.add(c)
    res = s.check()
print("no other solutions")
s.pop()

# try to find another puzzle with a unique solution
while True:
    s.push()
    print("new push")
    print(s.check())
    m = s.model()
    # condition c: keep found off-by-one sums configuration
    # condition c2: to ban this sums configuration
    c = True
    c2 = False
    for l in [*sums_h, *sums_v]:
        for x in l:
            c = And(c, x == m[x].as_long())
            c2 = Or(c2, x != m[x].as_long())
    s.add(c)
    # try to find another solution
    c = False
    for row in grid:
        for cell in row:
            c = Or(c, cell != m[cell].as_long())
    s.add(c)
    res = s.check()
    print(res)
    if res == unsat:
        print("found another puzzle with unique solution")
        print_solution(m, print_grid=False)
        input("continue? [press enter]")
    s.pop()
    s.add(c2)