Edge-Matching CSP

edge-matching.py

pieces = [
    ('B', 'P', 'O', 'Y'),
    ('W', 'B', 'P', 'R'),
    ('Y', 'B', 'G', 'W'),
    ('Y', 'G', 'B', 'R')
]

variables = [
    (0,0), (0,1),
    (1,0), (1,1)
]

domains = { v: [(p,n) for p in pieces for n in range(4)] for v in variables }
neighbors = { x: [y for y in variables if y != x] for x in variables}

def constraints(A,a,B,b):
    def rotate(l, n): # -n is clockwise
        return l[-n:] + l[:-n]
    
    piece_a, rotate_a = a
    piece_b, rotate_b = b
    
    if piece_a == piece_b:
        return False

    x,y = A[0]-B[0],A[1]-B[1]
    
    checks = {
        (1,0): (0,2),
        (0,-1): (1,3),
        (-1,0): (2,0),
        (0,1): (3,1)
    }
    
    if (x,y) in checks:    
        i,j = checks[(x,y)]
        piece_a = rotate(piece_a,rotate_a)
        piece_b = rotate(piece_b,rotate_b)   
        return piece_a[i] == piece_b[j]
    
    return True
    
scramble = CSP(variables,domains,neighbors,constraints)