LAC-Practical-Assignments-WS2024/Lecture2/maestria-in-class.py

# coding: utf-8
import os, sys
from z3 import *
import string
# get the playground information
directory = os.path.dirname(os.path.realpath(__file__))
fname = directory + "/" + "test0.txt"

if len(sys.argv) == 2:
    fname = sys.argv[1]

with open(fname) as f:
    playground = f.read()
playground = [[c for c in row] for row in playground.strip().split("\n")]

# get the playground size
size_y = len(playground)
assert(size_y != 0)
size_x = len(playground[0])
assert(size_x != 0)
for row in playground:
    assert(len(row) == size_x)

bells = {}
for i in range(size_y):
    for j in range(size_x):
        c = playground[i][j]
        if ord(c) in range(ord("1"), ord("9") + 1):
            assert(int(c) not in bells.keys())
            bells[int(c)] = (j, i)

################################### Maestria ###################################


print(f"Playground size : {size_x} x {size_y}")
print("Playground:")
print("\n".join(["".join(i) for i in playground]))
print("Bell positions:")
print(bells)
print("\n")

# create the solver
solver = Solver()

fudge_x, fudge_y = Ints("fudge_x fudge_y")


# position must be in the board
solver.add(fudge_x >= 0, fudge_x < size_x, fudge_y >= 0, fudge_y < size_y)

# position must be on walkable cell
for i in range(size_y):
    for j in range(size_x):
        if playground[i][j] != "_":
            solver.add(Not(And(fudge_y == i, fudge_x == j)))

# encode distances in formula
distances = {bell: None for bell in bells.keys()}
for bell, (x,y) in bells.items():
    distances[bell] = (fudge_x - x) * (fudge_x - x) + (fudge_y - y) * (fudge_y - y)
#
# enforce ordering of distances
distances_sorted = [distances[b] for b in sorted(bells.keys())]
for i in range(len(distances_sorted) - 1):
    near = distances_sorted[i]
    far = distances_sorted[i+1]
    solver.add(near < far)
# call the solver and check satisfiability
while solver.check() == sat:
#result = solver.check()
#if result == sat:
    m = solver.model()
    if fudge_x is not None and fudge_y is not None:
        pos_x = m[fudge_x].as_long()
        pos_y = m[fudge_y].as_long()
        playground[pos_y][pos_x] = "F"
        solver.add(Or(fudge_x != pos_x, fudge_y != pos_y))

################################################################################

from colorama import Fore, Back, Style

cols = [Fore.BLUE, Fore.GREEN, Fore.RED, Fore.YELLOW, Fore.MAGENTA, Fore.CYAN]
tones = [Style.BRIGHT, Style.DIM]
cols = sum([[c + t for c in cols] for t in tones], [])

if sys.stdout.isatty():
    for bell in sorted(list(bells.keys())):
        x, y = bells[bell]
        playground[y][x] = Back.BLACK + cols[bell - 1] + playground[y][x] + Style.RESET_ALL
    for y in range(size_y):
        for x in range(size_x):
            c = playground[y][x]
            if c == "_": playground[y][x] = Fore.BLACK + Back.WHITE + c + Style.RESET_ALL
            if c == "X": playground[y][x] = Fore.WHITE + Back.BLACK + c + Style.RESET_ALL
            if c == "F": playground[y][x] = Fore.BLACK + Back.WHITE + c + Style.RESET_ALL

text = "\n".join(["".join([c for c in row]) for row in playground])
print(text)