Upload enhanced_heuristics.py

enhanced_heuristics.py

@@ -0,0 +1,841 @@
+"""
+Enhanced Heuristic Solvers for ARC-AGI
+Covers many more common patterns than the basic heuristics.
+These run instantly (no model) and serve as Track B of our pipeline.
+"""
+import copy
+import itertools
+from typing import List, Dict, Optional, Tuple
+from collections import Counter, defaultdict
+import numpy as np
+
+
+def grids_equal(g1, g2):
+    if g1 is None or g2 is None:
+        return False
+    if len(g1) != len(g2):
+        return False
+    for r1, r2 in zip(g1, g2):
+        if len(r1) != len(r2):
+            return False
+        if list(r1) != list(r2):
+            return False
+    return True
+
+
+class EnhancedHeuristicSolver:
+    """Comprehensive heuristic solver covering ~10-15% of ARC tasks."""
+    
+    def solve(self, task: Dict) -> Optional[List[List[int]]]:
+        """Try all heuristic solvers, return first match."""
+        solvers = [
+            self.try_identity,
+            self.try_color_map,
+            self.try_rotation,
+            self.try_flip,
+            self.try_transpose,
+            self.try_crop_nonzero,
+            self.try_crop_to_color,
+            self.try_scale,
+            self.try_tile,
+            self.try_gravity,
+            self.try_border,
+            self.try_fill_enclosed,
+            self.try_mirror_symmetric,
+            self.try_count_to_grid,
+            self.try_sort_rows,
+            self.try_sort_cols,
+            self.try_remove_color,
+            self.try_keep_largest_object,
+            self.try_overlay,
+            self.try_repeat_pattern,
+            self.try_extract_subgrid,
+            self.try_swap_colors,
+            self.try_mask_operation,
+            self.try_size_change_pattern,
+        ]
+        
+        for solver in solvers:
+            try:
+                result = solver(task)
+                if result is not None:
+                    # Validate result is reasonable
+                    if len(result) > 0 and all(len(r) > 0 for r in result):
+                        if all(all(isinstance(c, int) and 0 <= c <= 9 for c in row) for row in result):
+                            return result
+            except Exception:
+                continue
+        
+        return None
+    
+    # --- Geometric Transforms ---
+    
+    @staticmethod
+    def try_identity(task):
+        for p in task["train"]:
+            if p["input"] != p["output"]:
+                return None
+        return copy.deepcopy(task["test"][0]["input"])
+    
+    @staticmethod
+    def try_rotation(task):
+        for k in [1, 2, 3]:
+            if all(np.rot90(np.array(p["input"]), k=-k).tolist() == p["output"] for p in task["train"]):
+                return np.rot90(np.array(task["test"][0]["input"]), k=-k).tolist()
+        return None
+    
+    @staticmethod
+    def try_flip(task):
+        for fn in [np.fliplr, np.flipud]:
+            if all(fn(np.array(p["input"])).tolist() == p["output"] for p in task["train"]):
+                return fn(np.array(task["test"][0]["input"])).tolist()
+        return None
+    
+    @staticmethod
+    def try_transpose(task):
+        if all(np.array(p["input"]).T.tolist() == p["output"] for p in task["train"]):
+            return np.array(task["test"][0]["input"]).T.tolist()
+        return None
+    
+    @staticmethod
+    def try_mirror_symmetric(task):
+        """Check if output is input with added mirror symmetry."""
+        # Try completing horizontal/vertical symmetry
+        for axis in ['h', 'v']:
+            ok = True
+            for p in task["train"]:
+                arr = np.array(p["input"])
+                out = np.array(p["output"])
+                if arr.shape != out.shape:
+                    ok = False
+                    break
+                if axis == 'h':
+                    mirrored = np.fliplr(arr)
+                else:
+                    mirrored = np.flipud(arr)
+                # Check if output = combine input and its mirror (non-zero takes priority)
+                expected = arr.copy()
+                mask = arr == 0
+                expected[mask] = mirrored[mask]
+                if not np.array_equal(expected, out):
+                    ok = False
+                    break
+            if ok:
+                arr = np.array(task["test"][0]["input"])
+                if axis == 'h':
+                    mirrored = np.fliplr(arr)
+                else:
+                    mirrored = np.flipud(arr)
+                result = arr.copy()
+                mask = arr == 0
+                result[mask] = mirrored[mask]
+                return result.tolist()
+        return None
+    
+    # --- Color Operations ---
+    
+    @staticmethod
+    def try_color_map(task):
+        """Global color replacement: each color maps to another."""
+        inp0, out0 = task["train"][0]["input"], task["train"][0]["output"]
+        if len(inp0) != len(out0) or len(inp0[0]) != len(out0[0]):
+            return None
+        cmap = {}
+        for r in range(len(inp0)):
+            for c in range(len(inp0[0])):
+                k, v = inp0[r][c], out0[r][c]
+                if k in cmap and cmap[k] != v:
+                    return None
+                cmap[k] = v
+        for p in task["train"][1:]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+            for r in range(len(p["input"])):
+                for c in range(len(p["input"][0])):
+                    if cmap.get(p["input"][r][c]) != p["output"][r][c]:
+                        return None
+        return [[cmap.get(c, c) for c in row] for row in task["test"][0]["input"]]
+    
+    @staticmethod
+    def try_swap_colors(task):
+        """Swap two specific colors."""
+        for p in task["train"]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+        
+        # Find changed cells
+        swaps = set()
+        for p in task["train"]:
+            for r in range(len(p["input"])):
+                for c in range(len(p["input"][0])):
+                    if p["input"][r][c] != p["output"][r][c]:
+                        swaps.add((p["input"][r][c], p["output"][r][c]))
+        
+        if len(swaps) == 2:
+            s = list(swaps)
+            if s[0] == (s[1][1], s[1][0]):
+                a, b = s[0]
+                # Verify
+                ok = True
+                for p in task["train"]:
+                    for r in range(len(p["input"])):
+                        for c in range(len(p["input"][0])):
+                            ic = p["input"][r][c]
+                            oc = p["output"][r][c]
+                            if ic == a:
+                                if oc != b: ok = False
+                            elif ic == b:
+                                if oc != a: ok = False
+                            else:
+                                if oc != ic: ok = False
+                            if not ok: break
+                        if not ok: break
+                    if not ok: break
+                if ok:
+                    return [[b if c==a else (a if c==b else c) for c in row] 
+                            for row in task["test"][0]["input"]]
+        return None
+    
+    @staticmethod
+    def try_remove_color(task):
+        """Remove a specific color (replace with background)."""
+        for bg in range(10):
+            for remove_color in range(10):
+                if remove_color == bg:
+                    continue
+                ok = True
+                for p in task["train"]:
+                    if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                        ok = False
+                        break
+                    for r in range(len(p["input"])):
+                        for c in range(len(p["input"][0])):
+                            ic = p["input"][r][c]
+                            oc = p["output"][r][c]
+                            if ic == remove_color:
+                                if oc != bg: ok = False
+                            else:
+                                if oc != ic: ok = False
+                            if not ok: break
+                        if not ok: break
+                    if not ok: break
+                if ok:
+                    return [[bg if c==remove_color else c for c in row] 
+                            for row in task["test"][0]["input"]]
+        return None
+    
+    # --- Cropping ---
+    
+    @staticmethod
+    def try_crop_nonzero(task):
+        """Crop to bounding box of non-zero cells."""
+        for p in task["train"]:
+            arr = np.array(p["input"])
+            nz = np.argwhere(arr != 0)
+            if len(nz) == 0:
+                return None
+            r1, c1 = nz.min(0)
+            r2, c2 = nz.max(0)
+            if arr[r1:r2+1, c1:c2+1].tolist() != p["output"]:
+                return None
+        arr = np.array(task["test"][0]["input"])
+        nz = np.argwhere(arr != 0)
+        if len(nz) == 0:
+            return None
+        r1, c1 = nz.min(0)
+        r2, c2 = nz.max(0)
+        return arr[r1:r2+1, c1:c2+1].tolist()
+    
+    @staticmethod
+    def try_crop_to_color(task):
+        """Crop to bounding box of specific non-background color."""
+        for bg_color in range(10):
+            ok = True
+            for p in task["train"]:
+                arr = np.array(p["input"])
+                nz = np.argwhere(arr != bg_color)
+                if len(nz) == 0:
+                    ok = False
+                    break
+                r1, c1 = nz.min(0)
+                r2, c2 = nz.max(0)
+                if arr[r1:r2+1, c1:c2+1].tolist() != p["output"]:
+                    ok = False
+                    break
+            if ok:
+                arr = np.array(task["test"][0]["input"])
+                nz = np.argwhere(arr != bg_color)
+                if len(nz) == 0:
+                    return None
+                r1, c1 = nz.min(0)
+                r2, c2 = nz.max(0)
+                return arr[r1:r2+1, c1:c2+1].tolist()
+        return None
+    
+    @staticmethod
+    def try_extract_subgrid(task):
+        """Extract a specific rectangular region."""
+        # Check if all outputs are same size
+        out_shapes = set()
+        for p in task["train"]:
+            out_shapes.add((len(p["output"]), len(p["output"][0])))
+        if len(out_shapes) != 1:
+            return None
+        
+        oh, ow = out_shapes.pop()
+        
+        # Try finding the subgrid at every possible position
+        # Check if there's a consistent position rule
+        for p in task["train"]:
+            inp = np.array(p["input"])
+            out = np.array(p["output"])
+            ih, iw = inp.shape
+            
+            if oh > ih or ow > iw:
+                return None
+            
+            found = False
+            for r in range(ih - oh + 1):
+                for c in range(iw - ow + 1):
+                    if np.array_equal(inp[r:r+oh, c:c+ow], out):
+                        found = True
+                        break
+                if found:
+                    break
+            if not found:
+                return None
+        
+        return None  # Too ambiguous without more logic
+    
+    # --- Scaling ---
+    
+    @staticmethod
+    def try_scale(task):
+        for factor in [2, 3, 4, 5]:
+            ok = True
+            for p in task["train"]:
+                inp, out = np.array(p["input"]), np.array(p["output"])
+                if out.shape[0] != inp.shape[0]*factor or out.shape[1] != inp.shape[1]*factor:
+                    ok = False
+                    break
+                expected = np.repeat(np.repeat(inp, factor, axis=0), factor, axis=1)
+                if not np.array_equal(expected, out):
+                    ok = False
+                    break
+            if ok:
+                inp = np.array(task["test"][0]["input"])
+                return np.repeat(np.repeat(inp, factor, axis=0), factor, axis=1).tolist()
+        return None
+    
+    @staticmethod
+    def try_tile(task):
+        """Check if output is input tiled NxM times."""
+        for nr in range(1, 6):
+            for nc in range(1, 6):
+                if nr == 1 and nc == 1:
+                    continue
+                ok = True
+                for p in task["train"]:
+                    inp, out = np.array(p["input"]), np.array(p["output"])
+                    ih, iw = inp.shape
+                    if out.shape != (ih*nr, iw*nc):
+                        ok = False
+                        break
+                    expected = np.tile(inp, (nr, nc))
+                    if not np.array_equal(expected, out):
+                        ok = False
+                        break
+                if ok:
+                    inp = np.array(task["test"][0]["input"])
+                    return np.tile(inp, (nr, nc)).tolist()
+        return None
+    
+    # --- Gravity / Movement ---
+    
+    @staticmethod
+    def try_gravity(task):
+        """Check if colored cells fall down/up/left/right."""
+        for direction in ['down', 'up', 'left', 'right']:
+            ok = True
+            for p in task["train"]:
+                arr = np.array(p["input"])
+                out = np.array(p["output"])
+                if arr.shape != out.shape:
+                    ok = False
+                    break
+                
+                # Get background color (most common)
+                bg = Counter(arr.flatten().tolist()).most_common(1)[0][0]
+                
+                result = np.full_like(arr, bg)
+                h, w = arr.shape
+                
+                if direction == 'down':
+                    for c in range(w):
+                        non_bg = [arr[r, c] for r in range(h) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[h - len(non_bg) + i, c] = val
+                elif direction == 'up':
+                    for c in range(w):
+                        non_bg = [arr[r, c] for r in range(h) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[i, c] = val
+                elif direction == 'right':
+                    for r in range(h):
+                        non_bg = [arr[r, c] for c in range(w) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[r, w - len(non_bg) + i] = val
+                elif direction == 'left':
+                    for r in range(h):
+                        non_bg = [arr[r, c] for c in range(w) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[r, i] = val
+                
+                if not np.array_equal(result, out):
+                    ok = False
+                    break
+            
+            if ok:
+                arr = np.array(task["test"][0]["input"])
+                bg = Counter(arr.flatten().tolist()).most_common(1)[0][0]
+                result = np.full_like(arr, bg)
+                h, w = arr.shape
+                
+                if direction == 'down':
+                    for c in range(w):
+                        non_bg = [arr[r, c] for r in range(h) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[h - len(non_bg) + i, c] = val
+                elif direction == 'up':
+                    for c in range(w):
+                        non_bg = [arr[r, c] for r in range(h) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[i, c] = val
+                elif direction == 'right':
+                    for r in range(h):
+                        non_bg = [arr[r, c] for c in range(w) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[r, w - len(non_bg) + i] = val
+                elif direction == 'left':
+                    for r in range(h):
+                        non_bg = [arr[r, c] for c in range(w) if arr[r, c] != bg]
+                        for i, val in enumerate(non_bg):
+                            result[r, i] = val
+                
+                return result.tolist()
+        return None
+    
+    # --- Border / Frame ---
+    
+    @staticmethod
+    def try_border(task):
+        """Check if output adds a border around non-zero region."""
+        for p in task["train"]:
+            inp, out = np.array(p["input"]), np.array(p["output"])
+            if inp.shape != out.shape:
+                return None
+        
+        # Check if the change is adding a border color around objects
+        # Try: output = input with border cells colored
+        for border_color in range(10):
+            ok = True
+            for p in task["train"]:
+                inp = np.array(p["input"])
+                out = np.array(p["output"])
+                h, w = inp.shape
+                
+                expected = inp.copy()
+                bg = Counter(inp.flatten().tolist()).most_common(1)[0][0]
+                
+                for r in range(h):
+                    for c in range(w):
+                        if inp[r, c] == bg:
+                            # Check if adjacent to non-bg
+                            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                                nr, nc = r+dr, c+dc
+                                if 0 <= nr < h and 0 <= nc < w and inp[nr, nc] != bg:
+                                    expected[r, c] = border_color
+                                    break
+                
+                if not np.array_equal(expected, out):
+                    ok = False
+                    break
+            
+            if ok:
+                inp = np.array(task["test"][0]["input"])
+                h, w = inp.shape
+                bg = Counter(inp.flatten().tolist()).most_common(1)[0][0]
+                result = inp.copy()
+                for r in range(h):
+                    for c in range(w):
+                        if inp[r, c] == bg:
+                            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                                nr, nc = r+dr, c+dc
+                                if 0 <= nr < h and 0 <= nc < w and inp[nr, nc] != bg:
+                                    result[r, c] = border_color
+                                    break
+                return result.tolist()
+        return None
+    
+    # --- Fill Operations ---
+    
+    @staticmethod
+    def try_fill_enclosed(task):
+        """Fill enclosed regions with a specific color."""
+        for p in task["train"]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+        
+        # Try flood fill from edges
+        for fill_color in range(10):
+            ok = True
+            for p in task["train"]:
+                inp = np.array(p["input"])
+                out = np.array(p["output"])
+                h, w = inp.shape
+                
+                bg = Counter(inp.flatten().tolist()).most_common(1)[0][0]
+                
+                # Find enclosed bg regions (not connected to border)
+                visited = np.zeros_like(inp, dtype=bool)
+                
+                # BFS from border
+                from collections import deque
+                queue = deque()
+                for r in range(h):
+                    for c in [0, w-1]:
+                        if inp[r, c] == bg and not visited[r, c]:
+                            queue.append((r, c))
+                            visited[r, c] = True
+                for c in range(w):
+                    for r in [0, h-1]:
+                        if inp[r, c] == bg and not visited[r, c]:
+                            queue.append((r, c))
+                            visited[r, c] = True
+                
+                while queue:
+                    r, c = queue.popleft()
+                    for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                        nr, nc = r+dr, c+dc
+                        if 0 <= nr < h and 0 <= nc < w and not visited[nr, nc] and inp[nr, nc] == bg:
+                            visited[nr, nc] = True
+                            queue.append((nr, nc))
+                
+                # Fill enclosed
+                expected = inp.copy()
+                for r in range(h):
+                    for c in range(w):
+                        if inp[r, c] == bg and not visited[r, c]:
+                            expected[r, c] = fill_color
+                
+                if not np.array_equal(expected, out):
+                    ok = False
+                    break
+            
+            if ok:
+                from collections import deque
+                inp = np.array(task["test"][0]["input"])
+                h, w = inp.shape
+                bg = Counter(inp.flatten().tolist()).most_common(1)[0][0]
+                visited = np.zeros_like(inp, dtype=bool)
+                queue = deque()
+                for r in range(h):
+                    for c in [0, w-1]:
+                        if inp[r, c] == bg and not visited[r, c]:
+                            queue.append((r, c))
+                            visited[r, c] = True
+                for c in range(w):
+                    for r in [0, h-1]:
+                        if inp[r, c] == bg and not visited[r, c]:
+                            queue.append((r, c))
+                            visited[r, c] = True
+                while queue:
+                    r, c = queue.popleft()
+                    for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                        nr, nc = r+dr, c+dc
+                        if 0 <= nr < h and 0 <= nc < w and not visited[nr, nc] and inp[nr, nc] == bg:
+                            visited[nr, nc] = True
+                            queue.append((nr, nc))
+                result = inp.copy()
+                for r in range(h):
+                    for c in range(w):
+                        if inp[r, c] == bg and not visited[r, c]:
+                            result[r, c] = fill_color
+                return result.tolist()
+        return None
+    
+    # --- Counting / Size ---
+    
+    @staticmethod
+    def try_count_to_grid(task):
+        """Output is a small grid whose size encodes a count."""
+        # Check if output dimensions relate to counting something in input
+        return None  # Complex, skip for now
+    
+    # --- Sorting ---
+    
+    @staticmethod
+    def try_sort_rows(task):
+        """Check if rows are sorted by some criterion."""
+        for p in task["train"]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+        
+        # Try sorting rows by color values
+        for reverse in [False, True]:
+            ok = True
+            for p in task["train"]:
+                sorted_rows = sorted(p["input"], reverse=reverse)
+                if sorted_rows != p["output"]:
+                    ok = False
+                    break
+            if ok:
+                return sorted(task["test"][0]["input"], reverse=reverse)
+        return None
+    
+    @staticmethod
+    def try_sort_cols(task):
+        """Check if columns are sorted."""
+        for p in task["train"]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+        
+        for reverse in [False, True]:
+            ok = True
+            for p in task["train"]:
+                arr = np.array(p["input"])
+                sorted_arr = np.sort(arr, axis=0)
+                if reverse:
+                    sorted_arr = sorted_arr[::-1]
+                if not np.array_equal(sorted_arr, np.array(p["output"])):
+                    ok = False
+                    break
+            if ok:
+                arr = np.array(task["test"][0]["input"])
+                sorted_arr = np.sort(arr, axis=0)
+                if reverse:
+                    sorted_arr = sorted_arr[::-1]
+                return sorted_arr.tolist()
+        return None
+    
+    # --- Object Operations ---
+    
+    @staticmethod
+    def try_keep_largest_object(task):
+        """Keep only the largest connected component."""
+        from collections import deque
+        
+        for p in task["train"]:
+            if len(p["input"]) != len(p["output"]) or len(p["input"][0]) != len(p["output"][0]):
+                return None
+        
+        def get_objects(grid, bg=0):
+            arr = np.array(grid)
+            h, w = arr.shape
+            visited = np.zeros_like(arr, dtype=bool)
+            objects = []
+            
+            for r in range(h):
+                for c in range(w):
+                    if not visited[r, c] and arr[r, c] != bg:
+                        cells = []
+                        color = arr[r, c]
+                        queue = deque([(r, c)])
+                        visited[r, c] = True
+                        while queue:
+                            cr, cc = queue.popleft()
+                            cells.append((cr, cc))
+                            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+                                nr, nc = cr+dr, cc+dc
+                                if 0 <= nr < h and 0 <= nc < w and not visited[nr, nc] and arr[nr, nc] != bg:
+                                    visited[nr, nc] = True
+                                    queue.append((nr, nc))
+                        objects.append({"cells": cells, "color": color})
+            
+            return objects
+        
+        for bg in [0]:
+            ok = True
+            for p in task["train"]:
+                objects = get_objects(p["input"], bg)
+                if not objects:
+                    ok = False
+                    break
+                largest = max(objects, key=lambda o: len(o["cells"]))
+                
+                expected = np.full_like(np.array(p["input"]), bg)
+                inp_arr = np.array(p["input"])
+                for r, c in largest["cells"]:
+                    expected[r, c] = inp_arr[r, c]
+                
+                if not np.array_equal(expected, np.array(p["output"])):
+                    ok = False
+                    break
+            
+            if ok:
+                inp = task["test"][0]["input"]
+                objects = get_objects(inp, bg)
+                if not objects:
+                    return None
+                largest = max(objects, key=lambda o: len(o["cells"]))
+                result = np.full_like(np.array(inp), bg)
+                inp_arr = np.array(inp)
+                for r, c in largest["cells"]:
+                    result[r, c] = inp_arr[r, c]
+                return result.tolist()
+        
+        return None
+    
+    # --- Overlay / Combine ---
+    
+    @staticmethod
+    def try_overlay(task):
+        """Check if output is OR/AND overlay of two halves of input."""
+        for split_type in ['horizontal', 'vertical']:
+            for op in ['or', 'and', 'xor']:
+                ok = True
+                for p in task["train"]:
+                    inp = np.array(p["input"])
+                    out = np.array(p["output"])
+                    h, w = inp.shape
+                    
+                    if split_type == 'horizontal' and h % 2 == 0:
+                        top = inp[:h//2]
+                        bot = inp[h//2:]
+                        if out.shape != top.shape:
+                            ok = False
+                            break
+                    elif split_type == 'vertical' and w % 2 == 0:
+                        top = inp[:, :w//2]
+                        bot = inp[:, w//2:]
+                        if out.shape != top.shape:
+                            ok = False
+                            break
+                    else:
+                        ok = False
+                        break
+                    
+                    bg = 0
+                    if op == 'or':
+                        expected = np.where(top != bg, top, bot)
+                    elif op == 'and':
+                        expected = np.where((top != bg) & (bot != bg), top, bg)
+                    elif op == 'xor':
+                        expected = np.where((top != bg) ^ (bot != bg), 
+                                          np.where(top != bg, top, bot), bg)
+                    
+                    if not np.array_equal(expected, out):
+                        ok = False
+                        break
+                
+                if ok:
+                    inp = np.array(task["test"][0]["input"])
+                    h, w = inp.shape
+                    
+                    if split_type == 'horizontal':
+                        top = inp[:h//2]
+                        bot = inp[h//2:]
+                    else:
+                        top = inp[:, :w//2]
+                        bot = inp[:, w//2:]
+                    
+                    bg = 0
+                    if op == 'or':
+                        result = np.where(top != bg, top, bot)
+                    elif op == 'and':
+                        result = np.where((top != bg) & (bot != bg), top, bg)
+                    elif op == 'xor':
+                        result = np.where((top != bg) ^ (bot != bg),
+                                        np.where(top != bg, top, bot), bg)
+                    
+                    return result.tolist()
+        return None
+    
+    # --- Pattern Repetition ---
+    
+    @staticmethod
+    def try_repeat_pattern(task):
+        """Check if output repeats a pattern found in input."""
+        return None  # Complex pattern detection, skip
+    
+    @staticmethod
+    def try_mask_operation(task):
+        """Check if one color acts as a mask for another."""
+        return None  # Complex, skip
+    
+    @staticmethod
+    def try_size_change_pattern(task):
+        """Detect systematic size change patterns."""
+        # Check if output is always a fixed size
+        out_sizes = set()
+        for p in task["train"]:
+            out_sizes.add((len(p["output"]), len(p["output"][0])))
+        
+        if len(out_sizes) == 1:
+            oh, ow = out_sizes.pop()
+            # Check if it's always 1x1 (counting/classification)
+            if oh == 1 and ow == 1:
+                # Try: output color = number of unique non-bg colors
+                for bg in [0]:
+                    ok = True
+                    for p in task["train"]:
+                        colors = set(c for row in p["input"] for c in row) - {bg}
+                        if p["output"] != [[len(colors)]]:
+                            ok = False
+                            break
+                    if ok:
+                        bg = 0
+                        colors = set(c for row in task["test"][0]["input"] for c in row) - {bg}
+                        return [[len(colors)]]
+        
+        return None
+
+
+# ============================================================
+# Test
+# ============================================================
+
+if __name__ == "__main__":
+    from datasets import load_dataset
+    
+    print("Loading ARC-AGI-2...")
+    ds = load_dataset("arc-agi-community/arc-agi-2", split="train")
+    tasks = []
+    for row in ds:
+        tasks.append({"train": row["fewshots"], "test": row["question"]})
+    
+    print("Loading ARC-AGI-1...")
+    ds_v1 = load_dataset("lordspline/arc-agi", split="training")
+    tasks_v1 = []
+    for row in ds_v1:
+        tasks_v1.append({"train": row["train"], "test": row["test"]})
+    
+    solver = EnhancedHeuristicSolver()
+    
+    # Eval ARC-AGI-2
+    correct_v2 = 0
+    total_v2 = 0
+    for task in tasks:
+        gt = task["test"][0].get("output")
+        if gt is None:
+            continue
+        total_v2 += 1
+        pred = solver.solve(task)
+        if pred is not None and grids_equal(pred, gt):
+            correct_v2 += 1
+    
+    print(f"ARC-AGI-2: {correct_v2}/{total_v2} = {correct_v2/total_v2*100:.1f}%")
+    
+    # Eval ARC-AGI-1
+    correct_v1 = 0
+    total_v1 = 0
+    for task in tasks_v1:
+        gt = task["test"][0].get("output")
+        if gt is None:
+            continue
+        total_v1 += 1
+        pred = solver.solve(task)
+        if pred is not None and grids_equal(pred, gt):
+            correct_v1 += 1
+    
+    print(f"ARC-AGI-1: {correct_v1}/{total_v1} = {correct_v1/total_v1*100:.1f}%")