witness_bruteforce.py

import itertools
import os
from multiprocessing.pool import Pool

import more_itertools

shapes = [
    [
        0b1100011000000000000000000,
        0b1110000000000000000000000,
    ],
    [
        0b1110000100000000000000000,
        0b0100001000110000000000000,
        # 0b1000011100000000000000000,
        0b1100010000100000000000000,
        0b1110010000000000000000000,
        0b1100001000010000000000000,
        0b0010011100000000000000000,
        0b1000010000110000000000000,
    ],
    [
        0b0100011000000000000000000,
        0b1000011000000000000000000,
    ],
    [
        0b1100010000000000000000000,
        0b1110000100000000000000000,
    ],
]

def get_shape_rows(shape):
    if not (shape & 0b111111111111111):
        if not (shape & 0b11111111111111111111):
            return 1
        return 2
    return 3

def get_row(shape, row):
    return (shape >> ((4-row)*5)) & 0b11111

def get_cell(shape, col, row):
    return (get_row(shape, row) >> (4-col)) & 1

def get_shape_cols(shape):
    rows = get_shape_rows(shape)
    max_cols = 0
    for row in range(rows):
        row = get_row(shape, row)
        row_cols = 5
        while not (row & 1):
            row_cols -= 1
            row >>= 1
        max_cols = max(max_cols, row_cols)
    return max_cols

def get_shape_cells(shape):
    cells = 0
    while shape:
        if shape & 1:
            cells += 1
        shape >>= 1
    return cells

def calculate(values):
    results = []
    for combination, x_values, y_values, total_cells in values:
        is_valid = True
        field = 0
        for shape, x, y in zip(combination, x_values, y_values):
            field |= shape >> (x + (5 * y))
        cells_filled = get_shape_cells(field)
        if cells_filled != total_cells:
            is_valid = False
        else:
            for row_idx in range(5):
                row_value = get_row(field, row_idx)
                if row_value not in ([0b11111, 0b01110, 0b00100] if row_idx not in [1, 3] else [0b11111, 0b01110]):
                    is_valid = False
        if is_valid:
            results.append((field, combination))
    return results

def main():
    shape_rows = {
        shape: get_shape_rows(shape) for shape in itertools.chain(*shapes)
    }
    shape_cols = {
        shape: get_shape_cols(shape) for shape in itertools.chain(*shapes)
    }
    shape_cells = {
        shape: get_shape_cells(shape) for shape in itertools.chain(*shapes)
    }

    # for shape in itertools.chain(*shapes):
    #     print(f'Shape {bin(shape)}: {shape_cols[shape]}x{shape_rows[shape]}, cells: {shape_cells[shape]}')


    valid_fields = []

    shape_combinations = list(itertools.product(*shapes))

    param_combinations = []

    for combination in shape_combinations:
        x_range = []
        y_range = []
        combination_cells = []
        for shape in combination:
            x_range.append(5 - shape_cols[shape] + 1)
            y_range.append(5 - shape_rows[shape] + 1)
            combination_cells.append(shape_cells[shape])

        total_cells = sum(combination_cells)

        for x_values in itertools.product(list(range(x_range[0])), list(range(x_range[1])), list(range(x_range[2])), list(range(x_range[3]))):
            for y_values in itertools.product(list(range(y_range[0])), list(range(y_range[1])), list(range(y_range[2])), list(range(y_range[3]))):
                param_combinations.append((combination, x_values, y_values, total_cells))

    pool = Pool()
    with pool:
        valid_fields = itertools.chain(*pool.map(calculate, more_itertools.divide(os.cpu_count(), param_combinations)))

    for field, combination in valid_fields:
        print('┌─────┐')
        for row_idx in range(5):
            row_value = get_row(field, row_idx)
            print('│', end='')
            for i in range(5):
                print('X' if (row_value >> (4 - i)) & 1 else ' ', end='')
            print('│  ', end='')

            for shape in combination:
                row_value = get_row(shape, row_idx)
                for i in range(5):
                    print('X' if (row_value >> (4 - i)) & 1 else ' ', end='')
            print()
        print('└─────┘')
        print('\n' + '==================================\n')

if __name__ == '__main__':
    main()