multi-puzzle-solver 0.9.13__py3-none-any.whl → 0.9.14__py3-none-any.whl

puzzle_solver/puzzles/black_box/black_box.py

@@ -0,0 +1,313 @@
+from dataclasses import dataclass
+from typing import Optional, Union
+import json
+
+import numpy as np
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Pos, get_all_pos, get_row_pos, get_col_pos, in_bounds, get_opposite_direction, get_next_pos, Direction
+from puzzle_solver.core.utils_ortools import generic_solve_all
+
+
+@dataclass(frozen=True)
+class SingleSolution:
+    assignment: dict[Pos, Union[str, int]]
+    beam_assignments: dict[Pos, list[tuple[int, Pos, Direction]]]
+
+    def get_hashable_solution(self) -> str:
+        result = []
+        for pos, v in self.assignment.items():
+            result.append((pos.x, pos.y, v))
+        return json.dumps(result, sort_keys=True)
+
+
+class Board:
+    def __init__(self, top: list, right: list, bottom: list, left: list, ball_count: Optional[tuple[int, int]] = None, max_travel_steps: Optional[int] = None):
+        assert len(top) == len(bottom), 'top and bottom must be the same length'
+        assert len(left) == len(right), 'left and right must be the same length'
+        self.K = len(top) + len(right) + len(bottom) + len(left)  # side count
+        self.H = len(top)
+        self.V = len(left)
+        if max_travel_steps is None:
+            self.T = self.V * self.H  # maximum travel steps for a beam that bounces an undefined number of times
+        else:
+            self.T = max_travel_steps
+        self.ball_count = ball_count
+
+        # top and bottom entry cells are at -1 and V
+        self.top_cells = set(get_row_pos(row_idx=-1, H=self.H))
+        self.bottom_cells = set(get_row_pos(row_idx=self.V, H=self.H))
+        # left and right entry cells are at -1 and H
+        self.left_cells = set(get_col_pos(col_idx=-1, V=self.V))
+        self.right_cells = set(get_col_pos(col_idx=self.H, V=self.V))
+        # self.top_cells = set([Pos(x=1, y=-1)])
+        # self.bottom_cells = set()
+        # self.left_cells = set()
+        # self.right_cells = set()
+        # print(self.top_cells)
+
+        self.top_values = top
+        self.right_values = right
+        self.bottom_values = bottom
+        self.left_values = left
+        
+        self.model = cp_model.CpModel()
+        self.ball_states: dict[Pos, cp_model.IntVar] = {}
+        # (entry_pos, T, cell_pos, direction) -> True if the beam that entered from the board at "entry_pos" is present in "cell_pos" and is going in the direction "direction" at time T
+        self.beam_states: dict[tuple[Pos, int, Pos, Direction], cp_model.IntVar] = {}
+        # self.beam_states_ending_at: dict[Pos, cp_model.IntVar] = {}
+        self.beam_states_at_t: dict[int, dict[Pos, dict[tuple[Pos, Direction], cp_model.IntVar]]] = {}
+        self.create_vars()
+        print('Total number of variables:', len(self.ball_states), len(self.beam_states))
+        self.add_all_constraints()
+        print('Solving...')
+
+    def get_outside_border(self):
+        top_border = tuple(get_row_pos(row_idx=-1, H=self.H))
+        bottom_border = tuple(get_row_pos(row_idx=self.V, H=self.H))
+        left_border = tuple(get_col_pos(col_idx=-1, V=self.V))
+        right_border = tuple(get_col_pos(col_idx=self.H, V=self.V))
+        return (*top_border, *bottom_border, *left_border, *right_border)
+
+    def get_all_pos_extended(self):
+        return (*self.get_outside_border(), *get_all_pos(self.V, self.H))
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.ball_states[pos] = self.model.NewBoolVar(f'ball_at:{pos}')
+        for pos in self.get_all_pos_extended():  # NxN board + 4 edges
+            if pos not in self.ball_states:  # pos is not in the board -> its on the edge
+                self.ball_states[pos] = None  # balls cant be on the edge
+
+        for entry_pos in (self.top_cells | self.right_cells | self.bottom_cells | self.left_cells):
+
+            for t in range(self.T):
+                self.beam_states[(entry_pos, t, 'HIT', 'HIT')] = self.model.NewBoolVar(f'beam:{entry_pos}:{t}:HIT:HIT')
+                for cell in self.get_all_pos_extended():
+                    for direction in Direction:
+                        self.beam_states[(entry_pos, t, cell, direction)] = self.model.NewBoolVar(f'beam:{entry_pos}:{t}:{cell}:{direction}')
+        
+        for (entry_pos, t, cell, direction) in self.beam_states.keys():
+            if t not in self.beam_states_at_t:
+                self.beam_states_at_t[t] = {}
+            if entry_pos not in self.beam_states_at_t[t]:
+                self.beam_states_at_t[t][entry_pos] = {}
+            self.beam_states_at_t[t][entry_pos][(cell, direction)] = self.beam_states[(entry_pos, t, cell, direction)]
+
+    def add_all_constraints(self):
+        self.init_beams()
+        self.constrain_beam_movement()
+        self.constrain_final_beam_states()
+        if self.ball_count is not None:
+            s = sum([b for b in self.ball_states.values() if b is not None])
+            b_min, b_max = self.ball_count
+            self.model.Add(s >= b_min)
+            self.model.Add(s <= b_max)
+
+    def init_beams(self):
+        beam_ids = []
+        beam_ids.extend((beam_id, Direction.DOWN) for beam_id in self.top_cells)
+        beam_ids.extend((beam_id, Direction.LEFT) for beam_id in self.right_cells)
+        beam_ids.extend((beam_id, Direction.UP) for beam_id in self.bottom_cells)
+        beam_ids.extend((beam_id, Direction.RIGHT) for beam_id in self.left_cells)
+        
+        for (beam_id, direction) in beam_ids:
+            # beam at t=0 is present at beam_id and facing direction
+            self.model.Add(self.beam_states[(beam_id, 0, beam_id, direction)] == 1)
+            for p in self.get_all_pos_extended():
+                # print(f'beam can be at {p}')
+                for direction in Direction:
+                    if (p, direction) != (beam_id, direction):
+                        self.model.Add(self.beam_states[(beam_id, 0, p, direction)] == 0)
+        # for p in self.get_all_pos_extended():
+        #     print(f'beam can be at {p}')
+
+
+    def constrain_beam_movement(self):
+        for t in range(self.T):
+            for entry_pos in self.beam_states_at_t[t].keys():
+                next_state_dict = self.beam_states_at_t[t][entry_pos]
+                self.model.AddExactlyOne(list(next_state_dict.values()))
+                # print('add exactly one constraint for beam id', entry_pos, 'at time', t, next_state_dict.keys(), '\n')
+                if t == self.T - 1:
+                    continue
+                for (cell, direction), prev_state in next_state_dict.items():
+                    # print(f'for beam id {entry_pos}, time {t}\nif its at {cell} facing {direction}')
+                    self.constrain_next_beam_state(entry_pos, t+1, cell, direction, prev_state)
+
+
+    def constrain_next_beam_state(self, entry_pos: Pos, t: int, cur_pos: Pos, direction: Direction, prev_state: cp_model.IntVar):
+        # print(f"constraining next beam state for {entry_pos}, {t}, {cur_pos}, {direction}")
+        if cur_pos == "HIT":  # a beam that was "HIT" stays "HIT"
+            # print(f'                                        HIT -> stays HIT')
+            self.model.Add(self.beam_states[(entry_pos, t, "HIT", "HIT")] == 1).OnlyEnforceIf(prev_state)
+            return
+        # if a beam is outside the board AND it is not facing the board -> it maintains its state
+        pos_ahead = get_next_pos(cur_pos, direction)
+        if not in_bounds(pos_ahead, self.V, self.H) and not in_bounds(cur_pos, self.V, self.H):
+            # print(f'                                        OUTSIDE BOARD -> beam stays in the same state')
+            self.model.Add(self.beam_states[(entry_pos, t, cur_pos, direction)] == 1).OnlyEnforceIf(prev_state)
+            return
+
+        # look at the 3 balls ahead of the beam: thus 8 possible scenarios
+        # A beam with no balls ahead of it -> moves forward in the same direction (1 scenario)
+        # A beam that hits a ball head-on -> beam is "HIT" (4 scenarios)
+        # A beam with a ball in its front-left square and no ball ahead of it -> gets deflected 90 degrees to the right (1 scenario)
+        # A beam with a ball in its front-right square and no ball ahead of it -> gets similarly deflected to the left (1 scenario)
+        # A beam that would in its front-left AND front-right squares -> is reflected (1 scenarios)
+
+        direction_left = {
+            Direction.UP: Direction.LEFT,
+            Direction.LEFT: Direction.DOWN,
+            Direction.DOWN: Direction.RIGHT,
+            Direction.RIGHT: Direction.UP,
+        }[direction]
+        direction_right = {
+            Direction.UP: Direction.RIGHT,
+            Direction.RIGHT: Direction.DOWN,
+            Direction.DOWN: Direction.LEFT,
+            Direction.LEFT: Direction.UP,
+        }[direction]
+        reflected = get_opposite_direction(direction)
+        ball_left_pos = get_next_pos(pos_ahead, direction_left)
+        ball_right_pos = get_next_pos(pos_ahead, direction_right)
+        if in_bounds(pos_ahead, self.V, self.H):
+            ball_ahead = self.ball_states[pos_ahead]
+            ball_ahead_not = ball_ahead.Not()
+        else:
+            ball_ahead = False
+            ball_ahead_not = True
+        if in_bounds(ball_left_pos, self.V, self.H):
+            ball_left = self.ball_states[ball_left_pos]
+            ball_left_not = ball_left.Not()
+        else:
+            ball_left = False
+            ball_left_not = True
+        if in_bounds(ball_right_pos, self.V, self.H):
+            ball_right = self.ball_states[ball_right_pos]
+            ball_right_not = ball_right.Not()
+        else:
+            ball_right = False
+            ball_right_not = True
+        
+        pos_left = get_next_pos(cur_pos, direction_left)
+        pos_right = get_next_pos(cur_pos, direction_right)
+        pos_reflected = get_next_pos(cur_pos, reflected)
+        if not in_bounds(pos_left, self.V, self.H):
+            pos_left = cur_pos
+        if not in_bounds(pos_right, self.V, self.H):
+            pos_right = cur_pos
+        if not in_bounds(pos_reflected, self.V, self.H):
+            pos_reflected = cur_pos
+
+        # debug_states = {
+        #     'if ball head': (entry_pos, t, "HIT", "HIT"),
+        #     'if ball in front-left': (entry_pos, t, get_next_pos(cur_pos, direction_right), direction_right),
+        #     'if ball in front-right': (entry_pos, t, get_next_pos(cur_pos, direction_left), direction_left),
+        #     'if ball in front-left and front-right': (entry_pos, t, get_next_pos(cur_pos, reflected), reflected),
+        #     'if no ball ahead': (entry_pos, t, pos_ahead, direction),
+        # }
+        # for k,v in debug_states.items():
+        #     print(f'                                        {k} -> {v}')
+
+        # ball head-on -> beam is "HIT"
+        self.model.Add(self.beam_states[(entry_pos, t, "HIT", "HIT")] == 1).OnlyEnforceIf([
+            ball_ahead,
+            prev_state,
+        ])
+        # ball in front-left -> beam is deflected right
+        self.model.Add(self.beam_states[(entry_pos, t, pos_right, direction_right)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left,
+            ball_right_not,
+            prev_state,
+        ])
+        # ball in front-right -> beam is deflected left
+        self.model.Add(self.beam_states[(entry_pos, t, pos_left, direction_left)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left_not,
+            ball_right,
+            prev_state,
+        ])
+        # ball in front-left and front-right -> beam is reflected
+        self.model.Add(self.beam_states[(entry_pos, t, pos_reflected, reflected)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left,
+            ball_right,
+            prev_state,
+        ])
+        # no ball ahead -> beam moves forward in the same direction
+        self.model.Add(self.beam_states[(entry_pos, t, pos_ahead, direction)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left_not,
+            ball_right_not,
+            prev_state,
+        ])
+
+    def constrain_final_beam_states(self):
+        all_values = []
+        all_values.extend([(Pos(x=c, y=-1), top_value) for c, top_value in enumerate(self.top_values)])
+        all_values.extend([(Pos(x=self.H, y=c), right_value) for c, right_value in enumerate(self.right_values)])
+        all_values.extend([(Pos(x=c, y=self.V), bottom_value) for c, bottom_value in enumerate(self.bottom_values)])
+        all_values.extend([(Pos(x=-1, y=c), left_value) for c, left_value in enumerate(self.left_values)])
+        digits = {}
+        hits = []
+        reflects = []
+        for pos, value in all_values:
+            value = str(value)
+            if value.isdecimal():
+                digits.setdefault(value, []).append(pos)
+            elif value == 'H':
+                hits.append(pos)
+            elif value == 'R':
+                reflects.append(pos)
+            else:
+                raise ValueError(f'Invalid value: {value}')
+        for digit, pos_list in digits.items():
+            assert len(pos_list) == 2, f'digit {digit} has {len(pos_list)} positions: {pos_list}'
+            p1, p2 = pos_list
+            self.model.AddExactlyOne([self.beam_states[(p1, self.T-1, p2, direction)] for direction in Direction])
+            self.model.AddExactlyOne([self.beam_states[(p2, self.T-1, p1, direction)] for direction in Direction])
+        for hit in hits:
+            self.model.AddExactlyOne([self.beam_states[(hit, self.T-1, 'HIT', 'HIT')]])
+        for reflect in reflects:
+            self.model.AddExactlyOne([self.beam_states[(reflect, self.T-1, reflect, direction)] for direction in Direction])
+
+
+    def solve_and_print(self, verbose: bool = True):
+        count = 0
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            nonlocal count
+            count += 1
+            # if count > 99:
+            #     import sys
+            #     sys.exit(0)
+            assignment = {}
+            for pos in get_all_pos(self.V, self.H):
+                assignment[pos] = solver.value(self.ball_states[pos])
+            beam_assignments = {}
+            # for (entry_pos, t, cell, direction), v in self.beam_states.items():
+            #     if entry_pos not in beam_assignments:
+            #         beam_assignments[entry_pos] = []
+            #     if solver.value(v):  # for every beam it can only be present in one state at a time
+            #         beam_assignments[entry_pos].append((t, cell, direction))
+            # for k,v in beam_assignments.items():
+            #     beam_assignments[k] = sorted(v, key=lambda x: x[0])
+            #     print(k, beam_assignments[k], '\n')
+            # print(beam_assignments)
+            return SingleSolution(assignment=assignment, beam_assignments=beam_assignments)
+
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            res = np.full((self.V, self.H), ' ', dtype=object)
+            for pos in get_all_pos(self.V, self.H):
+                ball_state = 'O' if single_res.assignment[pos] else ' '
+                res[pos.y][pos.x] = ball_state
+            print(res)
+        r = generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)
+        # print('non unique count:', count)
+
+
+
+
+

puzzle_solver/puzzles/filling/filling.py

@@ -1,192 +1,117 @@
-import numpy as np
-from ortools.sat.python import cp_model
-
-from puzzle_solver.core.utils import Pos, get_all_pos, get_char, set_char, get_neighbors4, get_all_pos_to_idx_dict
-from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, and_constraint
-
-
-class Board:
-    """
-    Per cell p:
-      val[p]    ∈ {1..9} (respect givens)
-      region[p] ∈ {0..N-1}   # region id is the linear index of the region's root
-      is_root[p] <=> (region[p] == idx[p])
-      # NOTE: root is the minimum index among its members via region[p] ≤ idx[p]
-
-    Per (p, k) where k is a root index (0..N-1):
-      in_region[p,k]  <=> (region[p] == k)
-      dist[p,k] ∈ {0..INF}
-        - If in_region[p,k] = 0  ⇒  dist[p,k] = INF
-        - If p == pos_of(k) and is_root[pos_of(k)] = 1  ⇒  dist[p,k] = 0
-        - If in_region[p,k] = 1 and p != pos_of(k)  ⇒
-              dist[p,k] = 1 + min_n masked_dist[n,k]
-          where masked_dist[n,k] = dist[n,k] + 1 if in_region[n,k] else INF
-
-    Edge (u,v):
-      same-digit neighbors must be in same region.
-
-    Region sizes:
-      For each k: size[k] == #{p : region[p] == k}
-      If is_root[pos_of(k)] → size[k] == val[pos_of(k)]
-      Else size[k] == 0
-    """
-
-    def __init__(self, board: np.ndarray):
-        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
-        self.board = board
-        self.V, self.H = board.shape
-        assert all((c == ' ') or (str(c).isdecimal() and 0 <= int(c) <= 9)
-                   for c in np.nditer(board)), "board must contain space or digits 0..9"
-
-        self.N = self.V * self.H
-        self.INF = self.N + 1  # a safe "infinity" upper bound for distances
-
-        # Linear index maps (keyed by Pos; do NOT construct tuples)
-        self.idx_of: dict[Pos, int] = get_all_pos_to_idx_dict(self.V, self.H)
-        self.pos_of: list[Pos] = [None] * self.N
-        for pos, idx in self.idx_of.items():
-            self.pos_of[idx] = pos
-
-        m = self.model = cp_model.CpModel()
-
-        # Variables
-        self.val = {}          # val[p]
-        self.region = {}       # region[p]
-        self.same_region = {}  # same_region[(p,q)]
-        self.is_root = {}      # is_root[p]
-        self.is_val = {}       # is_val[(p,k)]  (k=1..9)
-        self.in_region = {}    # in_region[(p,k)]  (k = 0..N-1)
-        self.dist = {}         # dist[(p,k)] ∈ [0..INF]
-
-        # Per-cell vars and givens
-        for p in get_all_pos(self.V, self.H):
-            idx = self.idx_of[p]
-
-            v = m.NewIntVar(1, 9, f'val[{idx}]')
-            ch = get_char(self.board, p)
-            if str(ch).isdecimal():
-                m.Add(v == int(ch))
-            self.val[p] = v
-
-            r = m.NewIntVar(0, self.N - 1, f'region[{idx}]')
-            self.region[p] = r
-
-            b = m.NewBoolVar(f'is_root[{idx}]')
-            self.is_root[p] = b
-            m.Add(r == idx).OnlyEnforceIf(b)
-            m.Add(r != idx).OnlyEnforceIf(b.Not())
-
-            # is_val indicators (for same-digit merge)
-            for k in range(1, 10):
-                bv = m.NewBoolVar(f'is_val[{idx}=={k}]')
-                self.is_val[(p, k)] = bv
-                m.Add(self.val[p] == k).OnlyEnforceIf(bv)
-                m.Add(self.val[p] != k).OnlyEnforceIf(bv.Not())
-
-        # Root = minimum index among members
-        for p in get_all_pos(self.V, self.H):
-            m.Add(self.region[p] <= self.idx_of[p])
-
-        # Membership indicators in_region[p,k] <=> region[p] == k
-        for k in range(self.N):
-            for p in get_all_pos(self.V, self.H):
-                bmem = m.NewBoolVar(f'in_region[{self.idx_of[p]}=={k}]')
-                self.in_region[(p, k)] = bmem
-                m.Add(self.region[p] == k).OnlyEnforceIf(bmem)
-                m.Add(self.region[p] != k).OnlyEnforceIf(bmem.Not())
-
-        # same-digit neighbors must be in the same region
-        for u in get_all_pos(self.V, self.H):
-            for v in get_neighbors4(u, self.V, self.H):
-                if self.idx_of[v] < self.idx_of[u]:
-                    continue  # undirected pair once
-                # If val[u]==k and val[v]==k for any k in 1..9, then region[u]==region[v]
-                # Implement as: for each k, (is_val[u,k] & is_val[v,k]) ⇒ (region[u]==region[v])
-                for k in range(1, 10):
-                    m.Add(self.region[u] == self.region[v])\
-                        .OnlyEnforceIf([self.is_val[(u, k)], self.is_val[(v, k)]])
-
-        for u in get_all_pos(self.V, self.H):
-            for v in get_neighbors4(u, self.V, self.H):
-                b = self.model.NewBoolVar(f'same_region[{self.idx_of[u]},{self.idx_of[v]}]')
-                self.same_region[(u, v)] = b
-                self.model.Add(self.region[u] == self.region[v]).OnlyEnforceIf(b)
-                self.model.Add(self.region[u] != self.region[v]).OnlyEnforceIf(b.Not())
-                self.model.Add(self.val[u] == self.val[v]).OnlyEnforceIf(b)
-
-        # Distance variables dist[p,k] and masked AddMinEquality
-        for k in range(self.N):
-            root_pos = self.pos_of[k]
-
-            for p in get_all_pos(self.V, self.H):
-                dp = m.NewIntVar(0, self.INF, f'dist[{self.idx_of[p]},{k}]')
-                self.dist[(p, k)] = dp
-
-                # If p not in region k -> dist = INF
-                m.Add(dp == self.INF).OnlyEnforceIf(self.in_region[(p, k)].Not())
-
-                # Root distance: if k is active at its own position -> dist[root,k] = 0
-                if p == root_pos:
-                    m.Add(dp == 0).OnlyEnforceIf(self.is_root[root_pos])
-                    # If root_pos isn't the root for k, membership is 0 and above rule sets INF.
-
-            # For non-root members p of region k: dist[p,k] = 1 + min masked neighbor distances
-            for p in get_all_pos(self.V, self.H):
-                if p == root_pos:
-                    continue  # handled above
-
-                # Build masked neighbor candidates: INF if neighbor not in region k; else dist[n,k] + 1
-                cand_vars = []
-                for n in get_neighbors4(p, self.V, self.H):
-                    cn = m.NewIntVar(0, self.INF, f'canddist[{self.idx_of[p]},{k}->{self.idx_of[n]}]')
-                    cand_vars.append(cn)
-
-                    both_in_region_k = m.NewBoolVar(f'both_in_region_k[{self.idx_of[p]} in {k} and {self.idx_of[n]} in {k}]')
-                    and_constraint(m, both_in_region_k, [self.in_region[(p, k)], self.in_region[(n, k)]])
-
-                    # Reified equality:
-                    # in_region[n,k] => cn == dist[n,k] + 1
-                    m.Add(cn == self.dist[(n, k)] + 1).OnlyEnforceIf(both_in_region_k)
-                    # not in_region[n,k] => cn == INF
-                    m.Add(cn == self.INF).OnlyEnforceIf(both_in_region_k.Not())
-
-                # Only enforce the min equation when p is actually in region k (and not the root position).
-                # If p ∉ region k, dp is already INF via the earlier rule.
-                if cand_vars:
-                    m.AddMinEquality(self.dist[(p, k)], cand_vars)
-        for p in get_all_pos(self.V, self.H):
-            # every cell must have 1 dist != INF (lets just do at least 1 dist != INF)
-            not_infs = []
-            for k in range(self.N):
-                not_inf = m.NewBoolVar(f'not_inf[{self.idx_of[p]},{k}]')
-                m.Add(self.dist[(p, k)] != self.INF).OnlyEnforceIf(not_inf)
-                m.Add(self.dist[(p, k)] == self.INF).OnlyEnforceIf(not_inf.Not())
-                not_infs.append(not_inf)
-            m.AddBoolOr(not_infs)
-
-        # Region sizes
-        for k in range(self.N):
-            root_pos = self.pos_of[k]
-            members = [self.in_region[(p, k)] for p in get_all_pos(self.V, self.H)]
-            size_k = m.NewIntVar(0, self.N, f'size[{k}]')
-            m.Add(size_k == sum(members))
-            # Active root -> size equals its value
-            m.Add(size_k == self.val[root_pos]).OnlyEnforceIf(self.is_root[root_pos])
-            # Inactive root id -> size 0
-            m.Add(size_k == 0).OnlyEnforceIf(self.is_root[root_pos].Not())
-
-    def solve_and_print(self, verbose: bool = True):
-        def board_to_solution(board: "Board", solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
-            assignment: dict[Pos, int] = {}
-            for p in get_all_pos(board.V, board.H):
-                assignment[p] = solver.Value(board.val[p])
-            return SingleSolution(assignment=assignment)
-        def callback(single_res: SingleSolution):
-            print("Solution found")
-            res = np.full((self.V, self.H), ' ', dtype=object)
-            for pos in get_all_pos(self.V, self.H):
-                c = get_char(self.board, pos)
-                c = single_res.assignment[pos]
-                set_char(res, pos, c)
-            print(res)
-        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)
+from dataclasses import dataclass
+
+import numpy as np
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Pos, Shape, get_all_pos, get_char, set_char, polyominoes, in_bounds, get_next_pos, Direction
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, and_constraint
+
+
+@dataclass
+class ShapeOnBoard:
+    is_active: cp_model.IntVar
+    N: int
+    body: set[Pos]
+    disallow_same_shape: set[Pos]
+
+
+class Board:
+    def __init__(self, board: np.ndarray, digits = (1, 2, 3, 4, 5, 6, 7, 8, 9)):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        self.board = board
+        self.V, self.H = board.shape
+        assert all((c == ' ') or (str(c).isdecimal() and 0 <= int(c) <= 9) for c in np.nditer(board)), "board must contain space or digits 0..9"
+        self.digits = digits
+        self.polyominoes = {d: polyominoes(d) for d in self.digits}
+        # len_shapes = sum(len(shapes) for shapes in self.polyominoes.values())
+        # print(f'total number of shapes: {len_shapes}')
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.digit_to_shapes = {d: [] for d in self.digits}
+        self.body_loc_to_shape = {(d,p): [] for d in self.digits for p in get_all_pos(self.V, self.H)}
+        self.forced_pos: dict[Pos, int] = {}
+
+        self.create_vars()
+        self.constrain_numbers_on_board()
+        self.init_polyominoes_on_board()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            for d in self.digits:
+                self.model_vars[(d,pos)] = self.model.NewBoolVar(f'{d}:{pos}')
+
+    def constrain_numbers_on_board(self):
+        # force numbers already on the board
+        for pos in get_all_pos(self.V, self.H):
+            c = get_char(self.board, pos)
+            if c.isdecimal():
+                self.model.Add(self.model_vars[(int(c),pos)] == 1)
+                self.forced_pos[pos] = int(c)
+
+    def init_polyominoes_on_board(self):
+        # total_count = 0
+        for d in self.digits:  # all digits
+            digit_count = 0
+            for pos in get_all_pos(self.V, self.H):  # translate by shape
+                for shape in self.polyominoes[d]:  # all shapes of d digits
+                    body = {pos + p for p in shape}
+                    if any(not in_bounds(p, self.V, self.H) for p in body):
+                        continue
+                    if any(p in self.forced_pos and self.forced_pos[p] != d for p in body):  # part of this shape's body is already forced to a different digit, skip
+                        continue
+                    disallow_same_shape = set(get_next_pos(p, direction) for p in body for direction in Direction)
+                    disallow_same_shape = {p for p in disallow_same_shape if p not in body and in_bounds(p, self.V, self.H)}
+                    shape_on_board = ShapeOnBoard(
+                        is_active=self.model.NewBoolVar(f'd{d}:{digit_count}:{pos}:is_active'),
+                        N=d,
+                        body=body,
+                        disallow_same_shape=disallow_same_shape,
+                    )
+                    self.digit_to_shapes[d].append(shape_on_board)
+                    for p in body:
+                        self.body_loc_to_shape[(d,p)].append(shape_on_board)
+                    digit_count += 1
+        #             total_count += 1
+        #             if total_count % 1000 == 0:
+        #                 print(f'{total_count} shapes on board')
+        # print(f'total number of shapes on board: {total_count}')
+
+    def add_all_constraints(self):
+        for pos in get_all_pos(self.V, self.H):
+            # exactly one digit is active at every position
+            self.model.AddExactlyOne(self.model_vars[(d,pos)] for d in self.digits)
+            # exactly one shape is active at that position
+            self.model.AddExactlyOne(s.is_active for d in self.digits for s in self.body_loc_to_shape[(d,pos)])
+        # if a shape is active then all its body is active
+        
+        for s_list in self.body_loc_to_shape.values():
+            for s in s_list:
+                for p in s.body:
+                    self.model.Add(self.model_vars[(s.N,p)] == 1).OnlyEnforceIf(s.is_active)
+
+        # same shape cannot touch
+        for d, s_list in self.digit_to_shapes.items():
+            for s in s_list:
+                for disallow_pos in s.disallow_same_shape:
+                    self.model.Add(self.model_vars[(d,disallow_pos)] == 0).OnlyEnforceIf(s.is_active)
+
+    def solve_and_print(self, verbose: bool = True):
+        def board_to_solution(board: "Board", solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            assignment: dict[Pos, int] = {}
+            for pos in get_all_pos(self.V, self.H):
+                for d in self.digits:
+                    if solver.Value(self.model_vars[(d,pos)]) == 1:
+                        assignment[pos] = d
+                        break
+            return SingleSolution(assignment=assignment)
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            res = np.full((self.V, self.H), ' ', dtype=object)
+            for pos in get_all_pos(self.V, self.H):
+                c = get_char(self.board, pos)
+                c = single_res.assignment[pos]
+                set_char(res, pos, c)
+            print(res)
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/inertia/tsp.py

@@ -342,7 +342,9 @@ def solve_optimal_walk(
     for attempt in range(attempts):
         cluster_orders = shuffled_cluster_orders()
         for meta in meta_list:
+            # print('solve once')
             rep_idxs, _, _, _ = solve_once(cluster_orders, meta)
+            # print('solve once done')
             if rep_idxs is None:
                 continue
 
@@ -370,6 +372,7 @@ def solve_optimal_walk(
                         new_nodes, new_cost = reps_to_nodes_and_cost(new_reps)
                         if new_cost < cost:
                             reps = new_reps
+                            # print('2-opt improved cost from', cost, 'to', new_cost)
                             nodes_seq, cost = new_nodes, new_cost
                             improved = True
                             break
@@ -383,7 +386,7 @@ def solve_optimal_walk(
 
     if best_nodes is None:
         raise RuntimeError("No solution found.")
-
+    # print('final check')
     # Final checks and edge list
     edge_walk: List[Tuple[Pos, Pos]] = [(best_nodes[i], best_nodes[i+1]) for i in range(len(best_nodes)-1)]
     assert all(e in edges for e in edge_walk), "Output contains an edge not in the input directed edges."