LAC-Practical-Assignments-2024/trees_and_tents.py

# coding: utf-8
import os, sys
from z3 import *

# get the playground information
if len(sys.argv) != 2:
    print("Usage: python3 trees_and_tents.py <forest-file>")

EMPTY = "?"
TREE = 2

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

with open(fname) as f:
    playground = f.read()
    cols = [col for col in playground.strip().split("\n")[0][1:]]
    rows = [row[0] for row in playground.strip().split("\n")[1:]]
playground = [[EMPTY if x == EMPTY else x for x in row[1:]]
        for row in playground.strip().split("\n")[1:]]

# get the playground size
size_y = len(rows)
assert(size_y != 0)
size_x = len(cols)
assert(size_x != 0)

def print_result(model):
    print("  " + ''.join(cols) + " "*5 + ''.join(cols))
    for i in range(size_y):
        print(rows[i] + " " + ''.join(playground[i]) + " >>> ", end='')
        for j in range(size_x):
            cell = model[cells[i][j]].as_long()
            if cell != TREE:
                print(cell, end='')
            elif cell == TREE:
                print("T", end='')
        print()

def print_forest():
    print("  " + ''.join(cols))
    for i in range(size_y):
        print(rows[i] + " ", end='')
        for j in range(size_x):
            cell = playground[i][j]
            if cell != TREE:
                print(cell, end='')
            elif cell == TREE:
                print("T", end='')
        print()

################################# Trees and Tents ##################################
# create the solver
cells = [[None for j in range(size_x)] for i in range(size_y)]
solver = Solver()

def get_possible_tents_in_col(col):
    possible_tent_positions = []
    for row in range(size_y):
        if playground[row][col] != "T":
           possible_tent_positions.append(cells[row][col])
    return possible_tent_positions

def get_possible_tents_in_row(row):
    possible_tent_positions = []
    for col in range(size_x):
        if playground[row][col] != "T":
           possible_tent_positions.append(cells[row][col])
    return possible_tent_positions

def get_neighbours(i, j):
    cs = []
    for p in range(max(i-1, 0), min(i+2, size_y)):
        for q in range(max(j-1, 0), min(j+2, size_x)):
            if p == i and q == j: continue
            cs.append(cells[p][q])
    return cs

def get_tree_neighbours(i, j):
    tree_pos = []
    for p in [max(i-1, 0), min(i+1, size_y - 1)]:
        if p == i: continue
        tree_pos.append(cells[p][j])
    for q in [max(j-1, 0), min(j+1, size_x - 1)]:
        if q == j: continue
        tree_pos.append(cells[i][q])
    return tree_pos


# TODO Cell entries need to be represented by a Z3 variable
# TODO Bound/restrict the values of the cells according to the input

# TODO The sums of tents per row/column must match the amounts given in the input

for i in range(size_y):
    for j in range(size_x):
        if playground[i][j] == "T":
            # TODO A tent needs to be next to a tree
            # Incrementally build up the tree_constraint ...
            tree_constraint = False
            for possible_tent in get_tree_neighbours(i, j):
                pass
        else:
            neighbours = get_neighbours(i, j)
            # TODO A tent must not be next to another tent

res = solver.check()
if res == unsat:
    print("UNSAT")
    print_forest()
else:
    m = solver.model()
    print(m)
    print_result(m)