multi-puzzle-solver 1.1.8__py3-none-any.whl

puzzle_solver/puzzles/guess/guess.py

@@ -0,0 +1,232 @@
+from collections import Counter, defaultdict
+from itertools import product
+
+import numpy as np
+
+
+class Board:
+    def __init__(self, num_pegs: int = 4, all_colors: tuple[str] = ('R', 'Y', 'G', 'B', 'O', 'P'), show_warnings: bool = True, show_progress: bool = False):
+        assert num_pegs >= 1, 'num_pegs must be at least 1'
+        assert len(all_colors) == len(set(all_colors)), 'all_colors must contain only unique colors'
+        self.previous_guesses = []
+        self.num_pegs = num_pegs
+        self.all_colors = all_colors
+        self.show_progress = show_progress
+        self.tqdm = None
+        if self.show_progress:
+            try:
+                from tqdm import tqdm
+                self.tqdm = tqdm
+            except ImportError:
+                print('tqdm is not installed, so progress bar will not be shown')
+                self.tqdm = None
+        self.possible_triplets = set((i, j, num_pegs-i-j) for i in range(num_pegs+1) for j in range(num_pegs+1-i))
+        int_to_color = {i: c for i, c in enumerate(self.all_colors)}
+        c = len(self.all_colors)**num_pegs
+        if c > 10**5 and show_warnings:
+            print(f'Warning: len(all_colors)**num_pegs is too large (= {c:,}). The solver may take infinitely long to run.')
+        self.all_possible_pegs = tuple({(i, int_to_color[int_]) for i, int_ in enumerate(ints)} for ints in product(range(len(self.all_colors)), repeat=num_pegs))
+
+    def add_guess(self, guess: tuple[tuple[int, str]], guess_result: tuple[int, int, int]):
+        assert len(guess) == self.num_pegs, 'previous guess must have the same number of pegs as the game'
+        assert not set(guess) - set(self.all_colors), f'previous guess must contain only colors in all_colors; invalid colors: {set(guess) - set(self.all_colors)}'
+        assert sum(guess_result) == self.num_pegs, 'guess result must sum to num_pegs'
+        self.previous_guesses.append((guess, guess_result))
+
+    def get_possible_ground_truths(self):
+        """
+        Returns the possible ground truths based on the previous guesses.
+        """
+        previous_guesses = self.previous_guesses
+        all_possible_pegs = self.all_possible_pegs
+        possible_triplets = self.possible_triplets
+        if self.tqdm is not None:
+            previous_guesses = self.tqdm(previous_guesses, desc='Step 1/2: Filtering possible ground truths')
+        # filter possible ground truths based on previous guesses
+        pair_mask = np.full((len(all_possible_pegs), ), True, dtype=bool)
+        for previous_guess, guess_result in previous_guesses:
+            previous_guess = tuple(tuple((i, c) for i, c in enumerate(previous_guess)))
+            pairs = np_information_gain(guess=previous_guess, possible_ground_truths=all_possible_pegs, possible_triplets=possible_triplets, return_pairs=True)
+            mask = np.all(pairs == guess_result, axis=1)
+            pair_mask &= mask
+        possible_ground_truths = tuple(all_possible_pegs[i] for i in range(len(all_possible_pegs)) if pair_mask[i])
+        return possible_ground_truths
+
+    def best_next_guess(
+        self,
+        return_guess_entropy: bool = False,
+        verbose: bool = True,
+    ):
+        """
+        Returns the best next guess that would maximize the Shannon entropy of the next guess.
+        """
+        possible_triplets = self.possible_triplets
+        all_possible_pegs = self.all_possible_pegs
+        possible_ground_truths = self.get_possible_ground_truths()
+        ng = len(possible_ground_truths)  # number of possible ground truths
+        if ng == 0:
+            print('No possible ground truths found. This should not happen in a real game, please check your inputted guesses.')
+            return ng, None
+        elif ng == 1:
+            answer = [c for i, c in sorted(possible_ground_truths[0], key=lambda x: x[0])]
+            if verbose:
+                print(f'Solution found! The solution is: {answer}')
+            return ng, answer
+        if verbose:
+            print(f'out of {len(all_possible_pegs)} possible ground truths, only {ng} are still possible.')
+
+        if self.tqdm is not None:
+            all_possible_pegs = self.tqdm(all_possible_pegs, desc='Step 2/2: Calculating entropy for each guess')
+        guess_entropy = []
+        possible_ground_truths_set = set(tuple((i, c) for i, c in guess) for guess in possible_ground_truths)
+        for guess in all_possible_pegs:
+            entropy = np_information_gain(guess=guess, possible_ground_truths=possible_ground_truths, possible_triplets=possible_triplets)
+            is_possible = tuple(guess) in possible_ground_truths_set
+            guess_entropy.append((guess, entropy, is_possible))
+        guess_entropy = sorted(guess_entropy, key=lambda x: (x[1], x[2]), reverse=True)
+        max_entropy_guess = guess_entropy[0]
+        if verbose:
+            answer = [c for i, c in sorted(max_entropy_guess[0], key=lambda x: x[0])]
+            print(f'max entropy guess is: {answer} with entropy {max_entropy_guess[1]:.4f}')
+        if return_guess_entropy:
+            return ng, max_entropy_guess, guess_entropy
+        else:
+            return ng, max_entropy_guess
+
+
+
+
+def get_triplets(guess, ground_truth, verbose=False):
+    """
+    Returns
+        1. Number of guesses that match the color and location
+        2. Number of guesses that match the color but not the location
+        3. Number of guesses that do not match the color or the location
+    e.g.
+        if guess is ((0, 'Y'), (1, 'R'), (2, 'R'), (3, 'R')) and ground_truth is ((0, 'Y'), (1, 'Y'), (2, 'R'), (3, 'R')), then the triplets are (3, 0, 1)
+        if guess is ((0, 'R'), (1, 'Y'), (2, 'Y'), (3, 'Y')) and ground_truth is ((0, 'Y'), (1, 'Y'), (2, 'R'), (3, 'R')), then the triplets are (1, 2, 1)
+    """
+    color_count = defaultdict(int)
+    for _, color in ground_truth:
+        color_count[color] += 1
+    matching_color_and_location = 0
+    matching_color_but_not_location = 0
+    not_matching = 0
+    done_locs = set()
+    for (loc, color) in guess:
+        if (loc, color) in ground_truth:
+            if verbose:
+                print(f'loc {loc} color {color} matched perfectly')
+            matching_color_and_location += 1
+            color_count[color] -= 1
+            done_locs.add(loc)
+    for (loc, color) in guess:
+        if loc in done_locs:
+            continue
+        if color_count.get(color, 0) > 0:
+            if verbose:
+                print(f'loc {loc} color {color} matched but not in the right location')
+            matching_color_but_not_location += 1
+            color_count[color] -= 1
+        else:
+            not_matching += 1
+    return matching_color_and_location, matching_color_but_not_location, not_matching
+
+def slow_information_gain(guess: set[tuple[int, str]], possible_ground_truths: set[set[tuple[int, str]]], possible_triplets: set[tuple[int, int, int]]):
+    # safe but slow solution used as a reference
+    counts = {triplet: 0 for triplet in possible_triplets}
+    for ground_truth in possible_ground_truths:
+        counts[tuple(get_triplets(guess, ground_truth))] += 1
+    px = {triplet: count / len(possible_ground_truths) for triplet, count in counts.items()}
+    entropy = -sum(px[triplet] * np.log2(px[triplet]) for triplet in possible_triplets if px[triplet] > 0)
+    # print(counts)
+    return entropy
+
+
+def np_information_gain(guess: tuple[tuple[int, str]], possible_ground_truths: tuple[set[tuple[int, str]]], possible_triplets: set[tuple[int, int, int]], return_pairs: bool = False):
+    # my attempt of a vectorized np solution
+    n = len(guess)
+    all_colors = set()
+    for _, color in guess:
+        all_colors.add(color)
+    for gt in possible_ground_truths:
+        for _, color in gt:
+            all_colors.add(color)
+    guess_mask = {c: np.full((n, 1), 0, dtype=np.int8) for c in all_colors}
+    for loc, color in guess:
+        guess_mask[color][loc] = 1
+    guess_mask_repeated = {c: np.repeat(guess_mask[c].T, len(possible_ground_truths), axis=0) for c in all_colors}
+
+    color_matrices = {c: np.full((len(possible_ground_truths), n), 0, dtype=np.int8) for c in all_colors}
+    for i, gt in enumerate(possible_ground_truths):
+        for loc, color in gt:
+            color_matrices[color][i, loc] = 1
+
+    pair_1 = sum(color_matrices[c] @ guess_mask[c] for c in all_colors)
+
+    pair_2_diff = {c: guess_mask_repeated[c] - color_matrices[c] for c in all_colors}
+    pos_mask = {c: pair_2_diff[c] > 0 for c in all_colors}
+    pair_2_extra_guess = {c: pair_2_diff[c].copy() for c in all_colors}
+    pair_2_extra_ground = {c: pair_2_diff[c].copy() for c in all_colors}
+    pair_2 = {}
+    for c in all_colors:
+        pair_2_extra_guess[c][~pos_mask[c]] = 0
+        pair_2_extra_guess[c] = np.sum(pair_2_extra_guess[c], axis=1)
+        pair_2_extra_ground[c][pos_mask[c]] = 0
+        pair_2_extra_ground[c] = np.abs(np.sum(pair_2_extra_ground[c], axis=1))
+        pair_2[c] = np.minimum(pair_2_extra_guess[c], pair_2_extra_ground[c])
+
+    pair_2 = sum(pair_2[c] for c in all_colors)
+    pair_2 = pair_2[:, None]
+
+    pair_3 = n - pair_1 - pair_2
+
+    pair = np.concatenate([pair_1, pair_2, pair_3], axis=1)
+    pair_counter = Counter(tuple(t) for t in pair)
+    counts = {triplet: pair_counter[triplet] for triplet in possible_triplets}
+    px = {triplet: count / len(possible_ground_truths) for triplet, count in counts.items()}
+    entropy = -sum(px[triplet] * np.log2(px[triplet]) for triplet in possible_triplets if px[triplet] > 0)
+    # print(counts)
+    if return_pairs:
+        return pair
+    else:
+        return entropy
+
+
+
+
+
+
+def fast_information_gain(guess: set[tuple[int, str]],
+                          possible_ground_truths: set[set[tuple[int, str]]],
+                          possible_triplets: set[tuple[int, int, int]]):
+    # chatgpt fast solution + many modifications by me
+    counts = {t: 0 for t in possible_triplets}
+
+    for gt in possible_ground_truths:
+        color_count = {}
+        for _, c in gt:
+            color_count[c] = color_count.get(c, 0) + 1
+
+        H = 0
+        for loc, c in guess:
+            if (loc, c) in gt:
+                H += 1
+                color_count[c] -= 1  # safe: gt contributes this occurrence
+
+        color_only = 0
+        for loc, c in guess:
+            if (loc, c) in gt:
+                continue
+            remain = color_count.get(c, 0)
+            if remain > 0:
+                color_only += 1
+                color_count[c] = remain - 1
+
+        triplet = (H, color_only, len(guess) - H - color_only)
+        counts[triplet] += 1
+
+    px = {triplet: count / len(possible_ground_truths) for triplet, count in counts.items()}
+    entropy = -sum(px[triplet] * np.log2(px[triplet]) for triplet in possible_triplets if px[triplet] > 0)
+    # print(counts)
+    return entropy

puzzle_solver/puzzles/heyawake/heyawake.py

@@ -0,0 +1,152 @@
+import numpy as np
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Pos, get_all_pos, get_col_pos, get_neighbors4, get_pos, get_char, get_row_pos
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, force_connected_component
+from puzzle_solver.core.utils_visualizer import combined_function, id_board_to_wall_fn
+
+
+def return_3_consecutives(int_list: list[int]) -> list[tuple[int, int]]:
+    """Given a list of integers (mostly with duplicates), return every consecutive sequence of 3 integer changes.
+    i.e. return a list of (begin_idx, end_idx) tuples where for each r=int_list[begin_idx:end_idx] we have r[0]!=r[1] and r[-2]!=r[-1] and len(r)>=3"""
+    out = []
+    change_indices = [i for i in range(len(int_list) - 1) if int_list[i] != int_list[i+1]]
+    # notice how for every subsequence r, the subsequence beginning index is in change_indices and the ending index - 1 is in change_indices
+    for i in range(len(change_indices) - 1):
+        begin_idx = change_indices[i]
+        end_idx = change_indices[i+1] + 1  # we want to include the first number in the third sequence
+        if end_idx > len(int_list):
+            continue
+        out.append((begin_idx, end_idx))
+    return out
+
+
+def get_diagonal(pos1: Pos, pos2: Pos) -> list[Pos]:
+    assert pos1 != pos2, 'positions must be different'
+    dx = pos1.x - pos2.x
+    dy = pos1.y - pos2.y
+    assert abs(dx) == abs(dy), 'positions must be on a diagonal'
+    sdx = 1 if dx > 0 else -1
+    sdy = 1 if dy > 0 else -1
+    return [get_pos(x=pos2.x + i*sdx, y=pos2.y + i*sdy) for i in range(abs(dx) + 1)]
+
+
+class Board:
+    def __init__(self, board: np.array, region_to_clue: dict[str, int]):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        assert all(str(c.item()).isdecimal() for c in np.nditer(board)), 'board must contain only space or digits'
+        self.board = board
+        self.V, self.H = board.shape
+        self.all_regions: set[int] = {int(c.item()) for c in np.nditer(board)}
+        self.region_to_clue = {int(k): v for k, v in region_to_clue.items()}
+        assert set(self.region_to_clue.keys()).issubset(self.all_regions), f'extra regions in region_to_clue: {set(self.region_to_clue.keys()) - self.all_regions}'
+        self.region_to_pos: dict[int, set[Pos]] = {r: set() for r in self.all_regions}
+        for pos in get_all_pos(self.V, self.H):
+            rid = int(get_char(self.board, pos))
+            self.region_to_pos[rid].add(pos)
+
+        self.model = cp_model.CpModel()
+        self.B: dict[Pos, cp_model.IntVar] = {}
+        self.W: dict[Pos, cp_model.IntVar] = {}
+
+        self.create_vars()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.B[pos] = self.model.NewBoolVar(f'B:{pos}')
+            self.W[pos] = self.model.NewBoolVar(f'W:{pos}')
+            self.model.AddExactlyOne([self.B[pos], self.W[pos]])
+
+    def add_all_constraints(self):
+        # Regions with a number should contain black cells matching the number.
+        for rid, clue in self.region_to_clue.items():
+            self.model.Add(sum([self.B[p] for p in self.region_to_pos[rid]]) == clue)
+        # 2 black cells cannot be adjacent horizontally or vertically.
+        for pos in get_all_pos(self.V, self.H):
+            for neighbor in get_neighbors4(pos, self.V, self.H):
+                self.model.AddBoolOr([self.W[pos], self.W[neighbor]])
+        # All white cells should be connected in a single group.
+        force_connected_component(self.model, self.W)
+        # A straight (orthogonal) line of connected white cells cannot span across more than 2 regions.
+        self.disallow_white_lines_spanning_3_regions()
+        # straight diagonal black lines from side wall to horizontal wall are not allowed; because they would disconnect the white cells
+        self.disallow_full_black_diagonal()
+        # disallow a diagonal black line coming out of a wall of length N then coming back in on the same wall; because it would disconnect the white cells
+        self.disallow_zigzag_on_wall()
+
+    def disallow_white_lines_spanning_3_regions(self):
+        # A straight (orthogonal) line of connected white cells cannot span across more than 2 regions.
+        row_to_region: dict[int, list[int]] = {row: [] for row in range(self.V)}
+        col_to_region: dict[int, list[int]] = {col: [] for col in range(self.H)}
+        for pos in get_all_pos(self.V, self.H):  # must traverse from least to most (both row and col)
+            rid = int(get_char(self.board, pos))
+            row_to_region[pos.y].append(rid)
+            col_to_region[pos.x].append(rid)
+        for row_num, row in row_to_region.items():
+            for begin_idx, end_idx in return_3_consecutives(row):
+                pos_list = [get_pos(x=x, y=row_num) for x in range(begin_idx, end_idx+1)]
+                self.model.AddBoolOr([self.B[p] for p in pos_list])
+        for col_num, col in col_to_region.items():
+            for begin_idx, end_idx in return_3_consecutives(col):
+                pos_list = [get_pos(x=col_num, y=y) for y in range(begin_idx, end_idx+1)]
+                self.model.AddBoolOr([self.B[p] for p in pos_list])
+
+    def disallow_full_black_diagonal(self):
+        corners_dx_dy = [
+            ((0, 0), 1, 1),
+            ((self.H-1, 0), -1, 1),
+            ((0, self.V-1), 1, -1),
+            ((self.H-1, self.V-1), -1, -1),
+        ]
+        for (corner_x, corner_y), dx, dy in corners_dx_dy:
+            for delta in range(1, min(self.V, self.H)):
+                pos1 = get_pos(x=corner_x, y=corner_y + delta*dy)
+                pos2 = get_pos(x=corner_x + delta*dx, y=corner_y)
+                diagonal_line = get_diagonal(pos1, pos2)
+                self.model.AddBoolOr([self.W[p] for p in diagonal_line])
+
+    def disallow_zigzag_on_wall(self):
+        for pos in get_row_pos(0, self.H):  # top line
+            for end_x in range(pos.x + 2, self.H, 2):  # end pos is even distance away from start pos
+                end_pos = get_pos(x=end_x, y=pos.y)
+                dx = end_x - pos.x
+                mid_pos = get_pos(x=pos.x + dx//2, y=pos.y + dx//2)
+                diag_1 = get_diagonal(pos, mid_pos)  # from top wall to bottom triangle tip "\"
+                diag_2 = get_diagonal(end_pos, mid_pos)  # from bottom triangle tip to top wall "/"
+                self.model.AddBoolOr([self.W[p] for p in diag_1 + diag_2])
+        for pos in get_row_pos(self.V-1, self.H):  # bottom line
+            for end_x in range(pos.x + 2, self.H, 2):  # end pos is even distance away from start pos
+                end_pos = get_pos(x=end_x, y=pos.y)
+                dx = end_x - pos.x
+                mid_pos = get_pos(x=pos.x + dx//2, y=pos.y - dx//2)
+                diag_1 = get_diagonal(pos, mid_pos)  # from bottom wall to top triangle tip "/"
+                diag_2 = get_diagonal(end_pos, mid_pos)  # from top triangle tip to bottom wall "\"
+                self.model.AddBoolOr([self.W[p] for p in diag_1 + diag_2])
+        for pos in get_col_pos(0, self.V):  # left line
+            for end_y in range(pos.y + 2, self.V, 2):  # end pos is even distance away from start pos
+                end_pos = get_pos(x=pos.x, y=end_y)
+                dy = end_y - pos.y
+                mid_pos = get_pos(x=pos.x + dy//2, y=pos.y + dy//2)
+                diag_1 = get_diagonal(pos, mid_pos)  # from left wall to right triangle tip "\"
+                diag_2 = get_diagonal(end_pos, mid_pos)  # from right triangle tip to left wall "/"
+                self.model.AddBoolOr([self.W[p] for p in diag_1 + diag_2])
+        for pos in get_col_pos(self.H-1, self.V):  # right line
+            for end_y in range(pos.y + 2, self.V, 2):  # end pos is even distance away from start pos
+                end_pos = get_pos(x=pos.x, y=end_y)
+                dy = end_y - pos.y
+                mid_pos = get_pos(x=pos.x - dy//2, y=pos.y + dy//2)
+                diag_1 = get_diagonal(pos, mid_pos)  # from right wall to left triangle tip "/"
+
+    def solve_and_print(self, verbose: bool = True, max_solutions: int = 20):
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            return SingleSolution(assignment={pos: solver.Value(var) for pos, var in board.B.items()})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H,
+                cell_flags=id_board_to_wall_fn(self.board),
+                is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1,
+                center_char=lambda r, c: self.region_to_clue.get(int(self.board[r, c]), ''),
+                text_on_shaded_cells=False
+            ))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=max_solutions)

puzzle_solver/puzzles/hidden_stars/hidden_stars.py

@@ -0,0 +1,52 @@
+import numpy as np
+from ortools.sat.python import cp_model
+from ortools.sat.python.cp_model import LinearExpr as lxp
+
+from puzzle_solver.core.utils import Pos, get_all_pos, get_char, get_pos, get_row_pos, get_col_pos, Direction8, get_ray
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
+
+
+class Board:
+    def __init__(self, board: np.array, side: np.array, top: np.array):
+        assert (len(side), len(top)) == board.shape, 'side and top must be the same shape as board'
+        assert all(c.item().strip() in ['', 'Q', 'W', 'E', 'A', 'D', 'Z', 'X', 'C'] for c in np.nditer(board)), 'board must contain only space or Q, W, E, A, D, Z, X, C'
+        self.board = board
+        self.side = side
+        self.top = top
+        self.V, self.H = board.shape
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.create_vars()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            if not get_char(self.board, pos).strip():
+                self.model_vars[pos] = self.model.NewBoolVar(f'{pos}')
+
+    def add_all_constraints(self):
+        for row in range(self.V):  # force side counts
+            self.model.Add(lxp.Sum([self.model_vars[pos] for pos in get_row_pos(row, self.H) if pos in self.model_vars]) == self.side[row])
+        for col in range(self.H):  # force top counts
+            self.model.Add(lxp.Sum([self.model_vars[pos] for pos in get_col_pos(col, self.V) if pos in self.model_vars]) == self.top[col])
+        pos_hit_by_ray = set()
+        for pos in get_all_pos(self.V, self.H):  # every arrow must point to a star
+            c = get_char(self.board, pos).strip()
+            if not c:
+                continue
+            d = {'Q': Direction8.UP_LEFT, 'W': Direction8.UP, 'E': Direction8.UP_RIGHT, 'A': Direction8.LEFT, 'D': Direction8.RIGHT, 'Z': Direction8.DOWN_LEFT, 'X': Direction8.DOWN, 'C': Direction8.DOWN_RIGHT}[c]
+            ray = get_ray(pos, d, self.V, self.H)
+            pos_hit_by_ray.update(ray)
+            self.model.AddAtLeastOne([self.model_vars[pos] for pos in ray if pos in self.model_vars])
+        for pos in get_all_pos(self.V, self.H):  # every star must have an arrow pointing to it
+            if pos not in pos_hit_by_ray and pos in self.model_vars:
+                self.model.Add(self.model_vars[pos] == 0)
+
+    def solve_and_print(self, verbose: bool = True):
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            return SingleSolution(assignment={pos: solver.Value(var) for pos, var in board.model_vars.items()})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H, is_shaded=lambda r, c: single_res.assignment.get(get_pos(x=c, y=r), 0) == 1, center_char=lambda r, c: str(self.board[r, c]), text_on_shaded_cells=False))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/hidoku/hidoku.py

@@ -0,0 +1,59 @@
+import numpy as np
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Direction8, Pos, get_all_pos, get_next_pos, get_pos, Direction, in_bounds, get_char
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
+
+
+class Board:
+    def __init__(self, board: np.array):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        assert all((c.item().strip() == '') or str(c.item()).isdecimal() for c in np.nditer(board)), 'board must contain only space or digits'
+        self.board = board
+        self.V, self.H = board.shape
+        self.N = self.V * self.H
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.step_up: dict[tuple[Pos, Direction], cp_model.IntVar] = {}
+        self.is_max: dict[Pos, cp_model.IntVar] = {}
+        self.create_vars()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.model_vars[pos] = self.model.NewIntVar(1, self.N, f'{pos}:num')
+            self.is_max[pos] = self.model.NewBoolVar(f'{pos}:is_max')
+            for direction in Direction8:
+                next_pos = get_next_pos(pos, direction)
+                if not in_bounds(next_pos, self.V, self.H):
+                    continue
+                self.step_up[(pos, direction)] = self.model.NewBoolVar(f'{pos}:{direction}:step_up')
+
+    def add_all_constraints(self):
+        self.model.AddAllDifferent(list(self.model_vars.values()))  # all numbers are unique
+        for pos in get_all_pos(self.V, self.H):
+            c = get_char(self.board, pos).strip()
+            if c:
+                self.model.Add(self.model_vars[pos] == int(c))
+            self.model.Add(self.model_vars[pos] == self.N).OnlyEnforceIf(self.is_max[pos])
+            self.model.Add(self.model_vars[pos] != self.N).OnlyEnforceIf(self.is_max[pos].Not())
+            for direction in Direction8:
+                next_pos = get_next_pos(pos, direction)
+                if not in_bounds(next_pos, self.V, self.H):
+                    continue
+                self.model.Add(self.model_vars[pos] + 1 == self.model_vars[next_pos]).OnlyEnforceIf(self.step_up[(pos, direction)])
+                self.model.Add(self.model_vars[pos] + 1 != self.model_vars[next_pos]).OnlyEnforceIf(self.step_up[(pos, direction)].Not())
+        # each position has a direction that is +1 and a direction that is -1
+        for pos in get_all_pos(self.V, self.H):
+            s_plus_one = [self.step_up[(pos, direction)] for direction in Direction8 if (pos, direction) in self.step_up]
+            self.model.AddBoolOr(s_plus_one + [self.is_max[pos]])
+
+    def solve_and_print(self, verbose: bool = True):
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            return SingleSolution(assignment={pos: solver.Value(var) for pos, var in board.model_vars.items()})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H, center_char=lambda r, c: str(single_res.assignment[get_pos(x=c, y=r)])))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/inertia/inertia.py

@@ -0,0 +1,121 @@
+from collections import defaultdict
+
+import numpy as np
+
+from puzzle_solver.core.utils import Direction8, Pos, get_all_pos, get_char, in_bounds, get_next_pos
+from . import tsp
+
+
+def _jump(board: np.array, pos: Pos, direction: Direction8) -> tuple[Pos, list[Pos]]:
+    # jump from pos in direction, return the next position and the positions of the gems that would be achieved (mostly likely None)
+    initial_pos = pos
+    out = []
+    while True:
+        next_pos = get_next_pos(pos, direction)
+        if not in_bounds(next_pos, board.shape[0], board.shape[1]):
+            break
+        ch = get_char(board, next_pos)
+        if ch == 'W':
+            break
+        if ch == 'O':
+            pos = next_pos
+            break
+        if ch == 'B':  # Note: the ball always starts ontop of an 'O' cell, thus hitting a 'B' is like hitting a 'O'
+            pos = next_pos
+            break
+        if ch == 'M':  # WE HIT A MINE
+            return None, None
+        if ch == 'G':
+            pos = next_pos
+            out.append(next_pos)
+            continue
+        if ch == ' ':
+            pos = next_pos
+            continue
+    if pos == initial_pos:  # we did not move
+        return None, None
+    return pos, out
+
+def parse_nodes_and_edges(board: np.array):
+    "parses the board into a graph where an edge is 1 move, and each gem lists the edges that would get it"
+    assert board.ndim == 2, 'board must be 2d'
+    assert all(c.item() in [' ', 'W', 'O', 'B', 'M', 'G'] for c in np.nditer(board)), 'board must contain only spaces, W, O, B, M, or G'
+    todo_nodes: set[Pos] = set()
+    completed_nodes: set[Pos] = set()
+    edges_to_direction: dict[tuple[Pos, Pos], Direction8] = {}  # edge (u: Pos, v: Pos) -> direction
+    gems_to_edges: dict[Pos, list[tuple[Pos, Pos]]] = defaultdict(list)  # gem position -> list of edges that would get this gem
+    V, H = board.shape
+    start_pos = [p for p in get_all_pos(V, H) if get_char(board, p) == 'B']
+    assert len(start_pos) == 1, 'board must have exactly one start position'
+    start_pos = start_pos[0]
+    todo_nodes.add(start_pos)
+    while todo_nodes:
+        pos = todo_nodes.pop()
+        for direction in Direction8:
+            next_pos, gems = _jump(board, pos, direction)
+            if next_pos is None:
+                continue
+            e = (pos, next_pos)
+            assert e not in edges_to_direction, 'edge already exists'
+            edges_to_direction[e] = direction
+            if len(gems) > 0:
+                for gem in gems:
+                    assert e not in gems_to_edges[gem], 'edge already in gems_to_edges'
+                    gems_to_edges[gem].append(e)
+            if next_pos not in completed_nodes:
+                todo_nodes.add(next_pos)
+        completed_nodes.add(pos)
+    assert len(gems_to_edges) == len([p for p in get_all_pos(V, H) if get_char(board, p) == 'G']), 'some gems are not reachable'
+    edges = set(edges_to_direction.keys())
+    return start_pos, edges, edges_to_direction, gems_to_edges
+
+def get_moves_from_walk(walk: list[tuple[Pos, Pos]], edges_to_direction: dict[tuple[Pos, Pos], Direction8], verbose: bool = True) -> list[str]:
+    direction_to_str = {Direction8.UP: '↑', Direction8.DOWN: '↓', Direction8.LEFT: '←', Direction8.RIGHT: '→', Direction8.UP_LEFT: '↖', Direction8.UP_RIGHT: '↗', Direction8.DOWN_LEFT: '↙', Direction8.DOWN_RIGHT: '↘'}
+    for edge in walk:
+        assert edge in edges_to_direction, f'edge {edge} not valid yet was in walk'
+    walk_directions = [edges_to_direction[edge] for edge in walk]
+    walk_directions_printable = [direction_to_str[x] for x in walk_directions]
+    if verbose:
+        print("number of moves", len(walk_directions))
+        for i, direction in enumerate(walk_directions_printable):
+            print(f"{direction}", end=' ')
+            if i % 10 == 9:
+                print()
+        print()
+    return walk_directions
+
+def simulate_moves(board: np.array, moves: list[str]) -> bool:
+    V, H = board.shape
+    start_pos = [p for p in get_all_pos(V, H) if get_char(board, p) == 'B']
+    assert len(start_pos) == 1, 'board must have exactly one start position'
+    gems_collected_so_far = set()
+    start_pos = start_pos[0]
+    current_pos = start_pos
+    for move in moves:
+        next_pos, gems = _jump(board, current_pos, move)
+        if next_pos is None:
+            print(f'invalid move {move} from {current_pos}. Either hit a wall (considered illegal here) or a mine (dead)')
+            return set()  # Running into a mine is fatal. Even if you picked up the last gem in the same move which then hit a mine, the game will count you as dead rather than victorious.
+        current_pos = next_pos
+        gems_collected_so_far.update(gems)
+    return gems_collected_so_far
+
+
+def is_board_completed(board: np.array, moves: list[str]) -> bool:
+    V, H = board.shape
+    all_gems = set(p for p in get_all_pos(V, H) if get_char(board, p) == 'G')
+    gems_collected = simulate_moves(board, moves)
+    assert gems_collected.issubset(all_gems), f'collected gems that are not on the board??? should not happen, {gems_collected - all_gems}'
+    return gems_collected == all_gems
+
+def solve_optimal_walk(
+    start_pos: Pos,
+    edges: set[tuple[Pos, Pos]],
+    gems_to_edges: defaultdict[Pos, list[tuple[Pos, Pos]]],
+    *,
+    restarts: int = 1,          # try more for harder instances (e.g., 48–128)
+    time_limit_ms: int = 1000,   # per restart
+    seed: int = 0,
+    verbose: bool = False
+) -> list[tuple[Pos, Pos]]:
+    return tsp.solve_optimal_walk(start_pos, edges, gems_to_edges, restarts=restarts, time_limit_ms=time_limit_ms, seed=seed, verbose=verbose)