APS tetris solver

# Adapted from https://stackoverflow.com/a/47934736
import math
import copy
import subprocess
import time
import numpy as np
import matplotlib.pyplot as plt      # plotting-only
import seaborn as sns                # plotting-only
np.set_printoptions(linewidth=120)   # more nice console-output

# 3x3 breaks but whatever
M, N = 29, 29
# T-shaped tetromino (rotations)
T0 = np.array([[1, 1, 1],
               [0, 1, 0]])
T90 = np.array([[0, 1],
                [1, 1],
                [0, 1]])
T180 = np.array([[0, 1, 0],
                 [1, 1, 1]])
T270 = np.array([[1, 0],
                 [1, 1],
                 [1, 0]])
polyominos = [T0, T90, T180, T270]
def create_inverted_circle(N):
    if N % 2 == 0:
        raise ValueError("N must be an odd integer")

    grid = [[0 for _ in range(N)] for _ in range(N)]
    center =(int) ((N - 1) / 2.0)   # Grid center coordinate
    radius = N / 2.0   # Adjust radius to ensure symmetry

    for x in range(N):
        for y in range(N):
            # Calculate distance from cell center to grid center
            distance = math.sqrt((x - center)**2 + (y - center)**2)
            if distance >= radius:
                grid[x][y] = 1
    grid[center][center] = 1
    return np.array(grid)

inverted_circle_grid = create_inverted_circle(N)
""" Preprocessing
        Calculate:
        A: possible placements
        B: covered positions
        C: collisions between placements
"""
placements = []
covered = []
for p_ind, p in enumerate(polyominos):
    mP, nP = p.shape
    for x in range(M):
        for y in range(N):
            if x + mP <= M and y + nP <= N:
                # Check if the placement is valid (not overlapping with cells set to 1)
                placement_valid = True
                for i in range(mP):
                    for j in range(nP):
                        if p[i,j] == 1 and inverted_circle_grid[x+i,y+j] == 1:
                            placement_valid = False
                            break
                    if not placement_valid:
                        break
                if placement_valid:
                    placements.append((p_ind, x, y))
                    cover = np.zeros((M, N), dtype=bool)
                    cover[x:x+mP, y:y+nP] = p
                    covered.append(cover)
covered = np.array(covered)

collisions = []
for m in range(M):
    for n in range(N):
        collision_set = np.flatnonzero(covered[:, m, n])
        collisions.append(collision_set)

print(f"Placements: {len(placements)}, Collision sets: {len(collisions)}")

def parse_solution(output):
    # assumes there is one
    vars = []
    for line in output.split("\n"):
        if line:
            if line[0] == 'v':
                line_vars = list(map(lambda x: int(x), line.split()[1:]))
                vars.extend(line_vars)
    return vars

def solve(CNF):
    p = subprocess.Popen(["cryptominisat5.exe"], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, encoding='utf-8')
    result = p.communicate(input=CNF)[0]

    sat_line = result.find('s SATISFIABLE')
    if sat_line != -1:
        # solution found!
        vars = parse_solution(result)
        return True, vars
    else:
        return False, None

""" Helper-function: Cardinality constraints """

# K-ARY CONSTRAINT GENERATION
# ###########################
# SINZ, Carsten. Towards an optimal CNF encoding of boolean cardinality constraints.
# CP, 2005, 3709. Jg., S. 827-831.

def next_var_index(start):
    next_var = start
    while(True):
        yield next_var
        next_var += 1

class s_index():
    def __init__(self, start_index):
        self.firstEnvVar = start_index

    def next(self,i,j,k):
        return self.firstEnvVar + i*k +j

def gen_seq_circuit(k, input_indices, next_var_index_gen):
    cnf_string = ''
    s_index_gen = s_index(next(next_var_index_gen))

    if len(input_indices) < 3:
        # Handle cases with fewer than 3 input indices
        for i in range(len(input_indices)):
            cnf_string += (str(-input_indices[i]) + ' ' + str(s_index_gen.next(0, i, k)) + ' 0\n')
        return (cnf_string, len(input_indices) * k, 2 * len(input_indices) * k + len(input_indices) - 3 * k - 1)

    # write clauses of first partial sum (i.e. i=0)
    cnf_string += (str(-input_indices[0]) + ' ' + str(s_index_gen.next(0, 0, k)) + ' 0\n')
    for i in range(1, k):
        cnf_string += (str(-s_index_gen.next(0, i, k)) + ' 0\n')

    # write clauses for general case (i.e. 0 < i < n-1)
    for i in range(1, len(input_indices) - 1):
        cnf_string += (str(-input_indices[i]) + ' ' + str(s_index_gen.next(i, 0, k)) + ' 0\n')
        cnf_string += (str(-s_index_gen.next(i - 1, 0, k)) + ' ' + str(s_index_gen.next(i, 0, k)) + ' 0\n')
        for u in range(1, k):
            cnf_string += (str(-input_indices[i]) + ' ' + str(-s_index_gen.next(i - 1, u - 1, k)) + ' ' + str(s_index_gen.next(i, u, k)) + ' 0\n')
            cnf_string += (str(-s_index_gen.next(i - 1, u, k)) + ' ' + str(s_index_gen.next(i, u, k)) + ' 0\n')
        cnf_string += (str(-input_indices[i]) + ' ' + str(-s_index_gen.next(i - 1, k - 1, k)) + ' 0\n')

    # last clause for last variable
    cnf_string += (str(-input_indices[-1]) + ' ' + str(-s_index_gen.next(len(input_indices) - 2, k - 1, k)) + ' 0\n')

    return (cnf_string, (len(input_indices) - 1) * k, 2 * len(input_indices) * k + len(input_indices) - 3 * k - 1)

def gen_at_most_n_constraints(vars, start_var, n):
    constraint_string = ''
    used_clauses = 0
    used_vars = 0
    index_gen = next_var_index(start_var)
    circuit = gen_seq_circuit(n, vars, index_gen)
    constraint_string += circuit[0]
    used_clauses += circuit[2]
    used_vars += circuit[1]
    start_var += circuit[1]

    return [constraint_string, used_clauses, used_vars, start_var]

""" SAT-CNF: BASE """
X = np.arange(1, len(placements)+1)                                     # decision-vars
                                                                        # 1-index for CNF
Y = np.arange(len(placements)+1, len(placements)+1 + M*N).reshape(M,N)
next_var = len(placements)+1 + M*N                                      # aux-var gen
n_clauses = 0

cnf = ''                                                                # slow string appends
                                                                        # int-based would be better
# <= 1 for each collision-set
for cset in collisions:
    constraint_string, used_clauses, used_vars, next_var = \
        gen_at_most_n_constraints(X[cset].tolist(), next_var, 1)
    n_clauses += used_clauses
    cnf += constraint_string

# if field marked: one of covering placements active
for x in range(M):
    for y in range(N):
        covering_placements = X[np.flatnonzero(covered[:, x, y])]  # could reuse collisions
        clause = str(-Y[x,y])
        for i in covering_placements:
            clause += ' ' + str(i)
        clause += ' 0\n'
        cnf += clause
        n_clauses += 1

print('BASE CNF size')
print('clauses: ', n_clauses)
print('vars: ', next_var - 1)



""" SOLVE in loop -> decrease number of placed-fields until SAT """
print('CORE LOOP')
# Calculate area inside circle (non-inverted cells)
circle_area = np.sum(inverted_circle_grid == 0)

# Calculate minimum polyomino area
min_poly_area = min(np.sum(p) for p in polyominos)

# Floor to nearest multiple of minimum polyomino area
N_FIELD_HIT = (circle_area // min_poly_area) * min_poly_area
max_possible = (int)(N_FIELD_HIT/min_poly_area)



time_start = time.time()
while True:
    print(' N_FIELDS >= ', N_FIELD_HIT)
    # sum(y) >= N_FIELD_HIT
    # == sum(not y) <= M*N - N_FIELD_HIT
    cnf_final = copy.copy(cnf)
    n_clauses_final = n_clauses

    if N_FIELD_HIT == M*N:  # awkward special case
        constraint_string = ''.join([str(y) + ' 0\n' for y in Y.ravel()])
        n_clauses_final += N_FIELD_HIT
    else:
        constraint_string, used_clauses, used_vars, next_var = \
            gen_at_most_n_constraints((-Y).ravel().tolist(), next_var, M*N - N_FIELD_HIT)
        n_clauses_final += used_clauses

    n_vars_final = next_var - 1
    cnf_final += constraint_string
    cnf_final = 'p cnf ' + str(n_vars_final) + ' ' + str(n_clauses) + \
        ' \n' + cnf_final  # header

    status, sol = solve(cnf_final)
    if status:
        print(' SOL found: ', N_FIELD_HIT)

        """ Print sol """
        res = np.full((M, N), -1)  # Initialize with -1 for 'X'
        # Replace -1 with 0 where inverted_circle_grid is 0
        res[inverted_circle_grid == 0] = 0

        # Create mapping of polyomino types to list of available numbers
        poly_numbers = {}
        current_num = 1

        import random
        # Count placements and create a shuffled range of numbers
        n_placements = sum(1 for v in sol[:X.shape[0]] if v > 0)
        numbers = list(range(1, n_placements + 1))
        random.shuffle(numbers)

        # Create mapping of placement index to number
        number_map = {}
        number_idx = 0

        # First pass: assign numbers to placements
        for i, v in enumerate(sol[:X.shape[0]]):
            if v > 0:
                number_map[i] = numbers[number_idx]
                number_idx += 1

        # Second pass: fill the result grid
        counter = n_placements + 1
        for i, v in enumerate(sol[:X.shape[0]]):
            if v > 0:
                p, x, y = placements[i]
                pM, pN = polyominos[p].shape
                poly_nnz = np.where(polyominos[p] != 0)
                x_inds, y_inds = x+poly_nnz[0], y+poly_nnz[1]
                res[x_inds, y_inds] = number_map[i]
        print(res)

        """ Plot """
        ax1 = plt.subplot2grid((5, 12), (0, 0), colspan=11, rowspan=5)
        mask = (res==0)
        # Create custom colormap with black for -1, white for 0, and colors for others
        from matplotlib.colors import ListedColormap
        colors = ['black', 'white'] + sns.color_palette('gist_rainbow', n_colors=counter)
        custom_cmap = ListedColormap(colors)
        sns.heatmap(res, cmap=custom_cmap, mask=mask, cbar=False, square=True, linewidths=.1, ax=ax1, vmin=-1, vmax=counter-1)
        elapsed_time = time.time() - time_start
        plt.title(f"{M}x{N} Placements: {sum(1 for v in sol[:X.shape[0]] if v > 0)}/{max_possible} ({elapsed_time:.2f}s)")
        plt.tight_layout()
        plt.show()
        break
    with open(f"output_{M}X{N}.txt",mode="w") as file:
        file.write(cnf_final)

    N_FIELD_HIT -= min_poly_area  # binary-search could be viable in some cases
                      # but beware the empirical asymmetry in SAT-solvers:
                      #    finding solution vs. proving there is none!