magic_cube_solver.py

import numpy as np
from collections.abc import Sized, Iterable, Iterator
from itertools import product, count, combinations, permutations
from collections import defaultdict
from typing import Optional
import cycosat

BlocksColor = [
    '#ffb631',
    '#f0ea4b',
    '#7a1b7d',
    '#8ad6b3',
    '#2c8dda',
    '#96e2fa',
    '#ff6961',
]

BlocksCubePosition = [
    [[0, 0, 0], [ 0, 1, 0], [ 1, 0, 0]],
    [[0, 0, 0], [-1, 0, 0], [ 1, 1, 0], [0, 1, 0]],
    [[0, 0, 0], [ 1, 0, 0], [-1, 0, 0], [0, 1, 0]],
    [[0, 0, 0], [ 1, 0, 0], [ 2, 0, 0], [0, 1, 0]],
    [[0, 0, 0], [ 0, 1, 0], [ 1, 0, 0], [0, 0, 1]],
    [[0, 0, 0], [ 0, 1, 0], [ 1, 0, 0], [0, 1, 1]],
    [[0, 0, 0], [ 0, 1, 0], [ 1, 0, 0], [1, 0, 1]],
]

BlockNames = list("VZTLY<>")


def rotate(points, axis:str, angle:float):
    angle = np.deg2rad(angle)
    c = np.cos(angle)
    s = np.sin(angle)

    if axis == "x":
        m = [[1, 0, 0], [0, c, -s], [0, s, c]]
    elif axis == "y":
        m = [[c, 0, s], [0, 1, 0], [-s, 0, c]]
    elif axis == "z":
        m = [[c, -s, 0], [s, c, 0], [0, 0, 1]]

    m = np.array(m)
    return (m @ points.T).T


def create_all_rotations():
    rotations = []
    for axis, r1, r2 in product("xy", range(0, 360, 90), range(0, 360, 90)):
        rotation = []
        if axis == "y" and r1 in (0, 180):
            continue
        if r1 != 0:
            rotation.append((axis, r1))
        if r2 != 0:
            rotation.append(("z", r2))
        rotations.append(rotation)
    return rotations


def get_unique_blocks(blocks):
    res = []
    for block in blocks:
        block = block - np.min(block, axis=0, keepdims=True)
        block = tuple(sorted([tuple(point) for point in block.tolist()]))
        res.append(block)
    return list(set(res))    


def get_rotated_blocks(block):
    rotations = create_all_rotations()
    block = np.asarray(block, dtype=np.float)
    blocks = []
    for rotation in rotations:
        rotated_block = block.copy()
        for axis, angle in rotation:
            rotated_block = rotate(rotated_block, axis, angle)

        arr = np.round(rotated_block).astype(np.int)
        blocks.append(arr)
    return get_unique_blocks(blocks)


def get_placed_blocks(block, target):
    block = np.asarray(block)
    first_cube = block[0]
    blocks = []
    for target_cube in target:
        delta = target_cube - first_cube
        offseted_block = block + delta
        offseted_block = tuple(tuple(cube) for cube in offseted_block.tolist())
        offseted_block_set = set(offseted_block)
        if target.contains(offseted_block_set):
            blocks.append(offseted_block)
    return blocks


def exact_one(variables):
    clauses = [list(variables)]
    for v1, v2 in combinations(variables, 2):
        clauses.append([-v1, -v2])
    return clauses


class Block:
    def __init__(self, block, block_id):
        self.id = block_id
        self.blocks = get_rotated_blocks(block)
        
    def place(self, target):
        self.placed_blocks = []
        for block in self.blocks:
            self.placed_blocks.extend(get_placed_blocks(block, target))
        
    def create_variables(self, variables):
        self.variables = []
        for block in self.placed_blocks:
            v = next(variables)
            self.variables.append(v)
            
    def create_clauses(self, target, all_block_ids):
        self.clauses = []
        self.clauses.extend(exact_one(self.variables))
        for block, v in zip(self.placed_blocks, self.variables):
            for id_ in all_block_ids:
                if id_ == self.id:
                    continue
                target_variables = target.get_block_variables(block, id_)
                for v2 in target_variables:
                    self.clauses.append([-v, -v2])   


class Target:
    def __init__(self, cubes):
        self.cubes = np.asarray(cubes)
        self.cubes_set = set(tuple(cube) for cube in self.cubes.tolist())

    @classmethod
    def make_3x3x3_target(class_):
        return class_(list(product(range(3), range(3), range(3))))

    @classmethod
    def from_command(class_, command):
        stack = []
        x = y = z = 0
        points = [(x, y, z)]
        for c in command:
            if c == '[':
                stack.append((x, y, z))
                continue
            elif c == ']':
                x, y, z = stack.pop()
                continue
            elif c == 'x':
                x += 1
            elif c == "X":
                x -= 1
            elif c == 'y':
                y += 1
            elif c == 'Y':
                y -= 1
            elif c == 'z':
                z += 1
            elif c == 'Z':
                z -= 1
            points.append((x, y, z))
        return class_(points)
        
    def contains(self, cubes_set):
        return self.cubes_set.issuperset(cubes_set)
    
    def __getitem__(self, index):
        return self.cubes[index]
    
    def __len__(self):
        return len(self.cubes)
    
    def get_block_variables(self, block, block_id):
        return [self.variables[cube, block_id] for cube in block]
    
    def create_variables(self, blocks, variables):
        self.variables = {}
        self.cube_variables = defaultdict(dict)
        
        for cube in self.cubes_set:
            for block in blocks:
                v = next(variables)
                self.variables[cube, block.id] = v
                self.cube_variables[cube][block.id] = v
                
    def create_clauses(self):
        self.clauses = []
        for key, val in self.cube_variables.items():
            self.clauses.extend(exact_one(list(val.values())))
            
    def view(self):
        import pyvista as pv

        a = 0.48
        box = pv.Box(bounds=(-a, a, -a, a, -a, a))
        plotter = pv.Plotter()
        
        if hasattr(self, "solution"):
            for bid, block in self.solution.items():
                block = np.array(block)
                ds = pv.PolyData(block)
                plotter.add_mesh(ds.glyph(geom=box), color=BlocksColor[bid])
        else:
            ds = pv.PolyData(self.cubes)        
            plotter.add_mesh(ds.glyph(geom=box))
            
        plotter.show()

    def solve(self, block_names:Optional[str]=None):
        if block_names is None:
            block_names = BlockNames
        
        blocks = [Block(block, i) 
                    for i, (name, block) in enumerate(zip(BlockNames, BlocksCubePosition)) 
                    if name in block_names]

        all_block_ids = [block.id for block in blocks]

        variables = count(1, 1)

        for block in blocks:
            block.place(self)

        self.create_variables(blocks, variables)
        for block in blocks:
            block.create_variables(variables)

        self.create_clauses()
        for block in blocks:
            block.create_clauses(self, all_block_ids)

        sat = cycosat.CoSAT()
        sat.add_clauses(target.clauses)
        for block in blocks:
            sat.add_clauses(block.clauses)
        solution = sat.solve()

        self.solution = defaultdict(list)
        for (cube, bid), vid in self.variables.items():
            v = solution[vid - 1]
            if v > 0:
                self.solution[bid].append(cube)
        self.solution = dict(self.solution)


def solve(target, block_names:Optional[str]=None):
    if not isinstance(target):
        target = Target(target)

    target.solve(block_names)
    return target


Questions = {
    29:("yyxx[Y]ZXXYx", "VTL"),
    107:("xxZxZZYz[z]XXXyx[x]zXXYZZy[z]xxx", None),
    101:("Zxy[xxx]XXYYZxxxy[xx]XXXyxxxxx", None),
    102:("xZXZxxZZYYzyXX[zx]YxxZXXyyzx[x]ZY", None),
    105:("yZZyyyZYYYZ[Y]yyyx[z]Y[z]YY[z]x[Y]yy", None),
    103:("ZxxZZYXYXXyyx[Y]xzXXYYzy[y]xZx", None),
    104:("yZxYxYYXX[Y]y[y]xZxyyXXY[x]YY[Y]xx", None)
}

if __name__ == "__main__":
    import pprint
    #target = Target.make_3x3x3_target()
    command, block_names = Questions[104]
    target = Target.from_command(command)
    target.solve(block_names=block_names)
    target.view()
    #target.solve()
    #pprint.pprint(target.solution)