Small model checker for WFC "Even Simpler Tiled Model" (3 values for i*j world tiles)

model_checker.py

from sympy.logic.inference import satisfiable
from sympy.logic.boolalg import Implies,  Exclusive, to_cnf, Xor,Equivalent
from sympy import Symbol
from sympy import Or,And, Not



rule_dict = {
    'L': {
        (1, 0): {'C', 'L'},
        (0, 1): {'C', 'L'},
        (0, -1): {'C', 'L'},
        (-1, 0): {'L'}
        }, 
    'C': {
        (-1, 0): {'L'},
        (0, -1): {'C', 'S', 'L'},
        (1, 0): {'S'},
        (0, 1): {'C', 'S', 'L'}
        }, 
    'S': {
        (-1, 0): {'C', 'S'},
        (0, -1): {'C', 'S'},
        (1, 0): {'S'},
        (0, 1): {'C', 'S'}
        }
    }

rule_dict2 = {
    'K': {
        (1, 0): {'S'}, 
        (0, 1): {'B'}
        }, 
    'B': {
        (0, -1): {'K', 'S'}, 
        (1, 0): {'B'}, 
        (-1, 0): {'B'}
        }, 
    'S': {
        (-1, 0): {'K'}, 
        (0, 1): {'B'}
        }
    }

def get_neighbours(i,j,height,width):
    return [(x,y) for x, y in [
            (i-1, j),   # n
            (i, j-1),   # w
            (i+1, j),   # s
            (i, j+1),   # e
        ]
        if 0 <= x < height and 0 <= y < width]

def sort_f(s):
    return s[-2:]

def format_models(models):
    m = []
    sorted_keys = []
    for model in models:
        if model:
            sorted_dict = dict(sorted(model.items(), key=lambda i: str(i[0]).split('_')[::-1]))
            sorted_values = list(sorted_dict.values())
            m.append(sorted_values)
            sorted_keys = [str(key) for key, value in sorted_dict.items()]
        else:
            print("UNSAT")
    return m, sorted_keys
        

def check_sat(map_size, literals, rule_dict):
    indexed_literals = []
    indexed_ruleset = dict()
    uniqueness_rules = []
    height,width = map_size
    for literal in literals:
        for i in range(height):
            for j in range(width):
                indexed_literal = Symbol(f"{literal}_{i}{j}")
                indexed_literals.append(indexed_literal)
                uniqueness_rules.append([Symbol(f"{u}_{i}{j}") for u in literals])
                neighbours = get_neighbours(i,j,height,width)
                or_rule = []
                for n in neighbours:
                    nx,ny = n
                    rx,ry = nx - i, ny - j
                    if (rx,ry) in rule_dict[literal].keys():
                        vinculi = rule_dict[literal][(rx,ry)]
                        indexed_vinculi = [Symbol(f"{v}_{nx}{ny}") for v in vinculi]
                        to_exclude = set(literals).difference(vinculi)                        
                        not_vinculi = [Not(Symbol(f"{e}_{nx}{ny}")) for e in to_exclude]
                        if not_vinculi:
                            v_or = And(Or(*indexed_vinculi),*not_vinculi)
                        else:
                            v_or = Or(*indexed_vinculi)
                    else:
                        # no value for tile is admissible here
                        no_tile_vinculi = [Not(Symbol(f"{l}_{nx}{ny}")) for l in literals]
                        v_or = And(*no_tile_vinculi)
                    or_rule.append(v_or)                           
                if or_rule:
                   indexed_ruleset[indexed_literal] = And(*or_rule)
                
    iif_ruleset = []
    for k,v in indexed_ruleset.items():
        iif_ruleset.append(Equivalent(k,v))

    for u_rule in uniqueness_rules:
        iif_ruleset.append(Or(*u_rule))
    model_ruleset = And(*iif_ruleset)
    models = satisfiable(model_ruleset, all_models=True)
    return indexed_ruleset, model_ruleset, models

map_size = (2,2)

literals = ['C',
            'L',
            'S',
            ]
literals2 = ['B',
            'K',
            'S',
            ]

indexed_ruleset, model_ruleset, models = check_sat(map_size=map_size, literals=literals2, rule_dict=rule_dict2)
print('*** INDEXED RULESET ***')
print(indexed_ruleset)
print('*** MODELS FOR RULESET ***')
print(model_ruleset)
m, sorted_keys = format_models(models)
print(",".join(sorted_keys))
for mod in m:
    print(mod)
# print('*** CNF ***')
# print(to_cnf(model_ruleset))
# a_model = satisfiable(model_ruleset)
# print('*** A MODEL ***')
# print(a_model)