• Skynet0's avatar
    Code · 9f5fe1b7
    Skynet0 authored
    9f5fe1b7
masyu.py 1.4 KB
from typing import List, Optional

from cs11puzzles.grid import ORTHO_DIRS, VEC_TO_DIR_NAME, Grid, Point
from cs11puzzles.str_helpers import ortho_dirs_sym, str_matrix_to_text
from genres.masyu.masyu_util import MasyuInstance, load_puzzle
from z3.z3 import Bool, Int, ModelRef, sat


class MasyuCell:
    def __init__(self, p: Point, u: int):
        self.point = p
        self.uniquifier = u

        self.dirs = {
            d: Bool(f"masyu-dir-u{u}-d{VEC_TO_DIR_NAME[d]}-r{p.y}-c{p.x}")
            for d in ORTHO_DIRS
        }
        self.order = Int(f"masyu-order-u{u}-r{p.y}-c{p.x}")

    def cell_to_str(self, m: ModelRef):
        return ortho_dirs_sym(
            **{VEC_TO_DIR_NAME[vec]: bool(m[val]) for vec, val in self.dirs.items()}
        )


def solve_masyu(puzzle: MasyuInstance) -> Optional[List[List[str]]]:
    pearls = puzzle.pearls

    height = len(pearls)
    width = len(pearls[0])
    num_cells = height * width

    g = Grid[MasyuCell](height, width, MasyuCell)

    # TODO: constraints

    if g.solver.check() == sat:

        def label_fn(cell: MasyuCell, m: ModelRef):
            return cell.cell_to_str(m)

        return g.to_str_matrix(label_fn)

    return None


if __name__ == "__main__":
    filename = "tests/masyu/data/masyu_4.txt"
    puzzle = load_puzzle(filename)
    sol = solve_masyu(puzzle)

    if sol:
        print(str_matrix_to_text(sol))
    else:
        print("No solution.")