Update app.py

app.py

@@ -1,9 +1,6 @@
 """
-ARC-AGI-3 Feature Visualizer
+ARC-AGI-3 Puzzle Interface + Feature Visualizer
 Hugging Face Space: beanapologist/arc-agi
-
-Interactive demo of the Re/Im symmetry-aware feature extractor.
-Draw a grid, see what the CNN sees before any training happens.
 """
 
 import gradio as gr
@@ -11,186 +8,134 @@ import numpy as np
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
-import matplotlib.patches as mpatches
 from matplotlib.colors import ListedColormap
 import io
 from PIL import Image
 
-# ── Inline feature extractor (no external deps) ───────────────────────────────
+# ── ARC color palette ─────────────────────────────────────────────────────────
+
+ARC_HEX = [
+    '#000000','#1a6faf','#e03a3a','#3aa63a','#f5c400',
+    '#c060c0','#d07030','#aaaaaa','#60b8d0','#874010'
+]
+ARC_CMAP = ListedColormap(ARC_HEX)
+COLOR_NAMES = ['black','blue','red','green','yellow',
+               'purple','orange','gray','azure','maroon']
+
+# ── Feature extractor ─────────────────────────────────────────────────────────
 
 def _connected_components(mask):
     labels = np.zeros_like(mask, dtype=np.int32)
-    current_label = 0
+    cur = 0
     H, W = mask.shape
     for r in range(H):
         for c in range(W):
-            if mask[r, c] and labels[r, c] == 0:
-                current_label += 1
-                queue = [(r, c)]
-                labels[r, c] = current_label
-                while queue:
-                    y, x = queue.pop()
-                    for dy, dx in [(-1,0),(1,0),(0,-1),(0,1)]:
-                        ny, nx = y+dy, x+dx
+            if mask[r,c] and labels[r,c]==0:
+                cur += 1
+                q = [(r,c)]; labels[r,c] = cur
+                while q:
+                    y,x = q.pop()
+                    for dy,dx in [(-1,0),(1,0),(0,-1),(0,1)]:
+                        ny,nx = y+dy,x+dx
                         if 0<=ny<H and 0<=nx<W and mask[ny,nx] and labels[ny,nx]==0:
-                            labels[ny,nx] = current_label
-                            queue.append((ny,nx))
+                            labels[ny,nx]=cur; q.append((ny,nx))
     return labels
 
-def _distance_transform(mask, max_dist=8):
-    H, W = mask.shape
-    dist = np.full((H, W), max_dist, dtype=np.float32)
-    queue = []
-    for r in range(H):
-        for c in range(W):
-            if mask[r, c]:
-                dist[r, c] = 0
-                queue.append((r, c))
-    head = 0
-    while head < len(queue):
-        y, x = queue[head]; head += 1
-        for dy, dx in [(-1,0),(1,0),(0,-1),(0,1)]:
-            ny, nx = y+dy, x+dx
-            if 0<=ny<H and 0<=nx<W and dist[ny,nx] > dist[y,x]+1:
-                dist[ny,nx] = dist[y,x]+1
-                queue.append((ny,nx))
-    return np.clip(dist / max_dist, 0, 1)
-
-def _sobel(grid_f):
-    p = np.pad(grid_f, 1, mode='edge')
-    gx = (-p[:-2,:-2] - 2*p[1:-1,:-2] - p[2:,:-2]
-          + p[:-2,2:]  + 2*p[1:-1,2:]  + p[2:,2:]) / 8.0
-    gy = (-p[:-2,:-2] - 2*p[:-2,1:-1] - p[:-2,2:]
-          + p[2:,:-2]  + 2*p[2:,1:-1]  + p[2:,2:]) / 8.0
+def _sobel(f):
+    p = np.pad(f, 1, mode='edge')
+    gx = (-p[:-2,:-2]-2*p[1:-1,:-2]-p[2:,:-2]+p[:-2,2:]+2*p[1:-1,2:]+p[2:,2:])/8
+    gy = (-p[:-2,:-2]-2*p[:-2,1:-1]-p[:-2,2:]+p[2:,:-2]+2*p[2:,1:-1]+p[2:,2:])/8
     return gx, gy
 
-def _reflection_symmetry_map(grid, axis):
-    H, W = grid.shape
-    score_map = np.zeros((H, W), dtype=np.float32)
+def _sym_map(grid, axis):
+    H,W = grid.shape
+    s = np.zeros((H,W), np.float32)
     if axis == 'h':
         for x in range(W):
-            reach = min(x, W-1-x)
-            if reach == 0: score_map[:,x] = 1.0; continue
-            left  = grid[:, x-reach:x]
-            right = grid[:, x+1:x+reach+1][:, ::-1]
-            score_map[:,x] = (left == right).mean()
+            r = min(x, W-1-x)
+            if r==0: s[:,x]=1.; continue
+            s[:,x] = (grid[:,x-r:x]==grid[:,x+1:x+r+1][:,::-1]).mean()
     else:
         for y in range(H):
-            reach = min(y, H-1-y)
-            if reach == 0: score_map[y,:] = 1.0; continue
-            top    = grid[y-reach:y, :]
-            bottom = grid[y+1:y+reach+1,:][::-1,:]
-            score_map[y,:] = (top == bottom).mean()
-    return score_map
-
-def _winding_proxy(gx, gy):
-    p_gx = np.pad(gx, 1, mode='edge')
-    p_gy = np.pad(gy, 1, mode='edge')
-    dgydx = (p_gy[1:-1,2:] - p_gy[1:-1,:-2]) / 2.0
-    dgxdy = (p_gx[2:,1:-1] - p_gx[:-2,1:-1]) / 2.0
-    curl = dgydx - dgxdy
-    mx = np.abs(curl).max()
-    if mx > 0: curl = curl / mx
-    return curl.astype(np.float32)
-
-def extract_features_visual(grid_2d):
-    """Extract and return named Im-side feature maps for visualization."""
-    H, W = grid_2d.shape
-    color_float = grid_2d.astype(np.float32) / 9.0
-    gx, gy = _sobel(color_float)
-
-    h_sym  = _reflection_symmetry_map(grid_2d, 'h')
-    v_sym  = _reflection_symmetry_map(grid_2d, 'v')
-    curl   = _winding_proxy(gx, gy)
-
-    padded = np.pad(grid_2d, 1, mode='edge')
-    boundary = (
-        (padded[1:-1,1:-1] != padded[:-2,1:-1]) |
-        (padded[1:-1,1:-1] != padded[2:,1:-1])  |
-        (padded[1:-1,1:-1] != padded[1:-1,:-2]) |
-        (padded[1:-1,1:-1] != padded[1:-1,2:])
-    ).astype(np.float32)
-
-    edge_mag = np.sqrt(gx**2 + gy**2)
-    edge_mag = edge_mag / (edge_mag.max() + 1e-8)
-
-    # Connected component map
-    global_labels = np.zeros((H, W), dtype=np.int32)
-    current = 0
+            r = min(y, H-1-y)
+            if r==0: s[y,:]=1.; continue
+            s[y,:] = (grid[y-r:y,:]==grid[y+1:y+r+1,:][::-1,:]).mean()
+    return s
+
+def _boundary(grid):
+    p = np.pad(grid, 1, mode='edge')
+    return ((p[1:-1,1:-1]!=p[:-2,1:-1])|(p[1:-1,1:-1]!=p[2:,1:-1])|
+            (p[1:-1,1:-1]!=p[1:-1,:-2])|(p[1:-1,1:-1]!=p[1:-1,2:])).astype(np.float32)
+
+def compute_scores(grid):
+    h_sym = _sym_map(grid,'h').max()
+    v_sym = _sym_map(grid,'v').max()
+    b_cnt = int(_boundary(grid).sum())
+    gx,gy = _sobel(grid.astype(np.float32)/9)
+    edge  = float(np.sqrt(gx**2+gy**2).mean())
+    n_comp = 0
     for c in range(10):
-        mask = (grid_2d == c)
-        if not mask.any(): continue
-        labels = _connected_components(mask)
-        for comp_id in range(1, labels.max()+1):
-            current += 1
-            global_labels[labels == comp_id] = current
-    norm_labels = global_labels.astype(np.float32)
-    if norm_labels.max() > 0: norm_labels /= norm_labels.max()
-
-    return dict(
-        h_symmetry=h_sym,
-        v_symmetry=v_sym,
-        boundary=boundary,
-        edge_magnitude=edge_mag,
-        winding_curl=curl,
-        component_map=norm_labels,
-    )
-
-# ── ARC color palette ─────────────────────────────────────────────────────────
-
-ARC_COLORS = [
-    '#000000', '#1a6faf', '#e03a3a', '#3aa63a',
-    '#f5c400', '#c060c0', '#d07030', '#aaaaaa',
-    '#60b8d0', '#874010'
-]
-ARC_CMAP = ListedColormap(ARC_COLORS)
-
-# ── Preset puzzles ────────────────────────────────────────────────────────────
-
-PRESETS = {
-    "Horizontal mirror": """0 1 2 0 0 0 2 1 0
-0 1 2 0 0 0 2 1 0
-0 1 2 0 0 0 2 1 0
-0 0 0 0 0 0 0 0 0
-0 0 0 3 3 3 0 0 0
-0 0 0 3 0 3 0 0 0
-0 0 0 3 3 3 0 0 0""",
-
-    "Colored objects": """0 0 0 0 0 0 0 0 0
-0 1 1 0 0 2 2 0 0
-0 1 1 0 0 2 2 0 0
-0 0 0 0 0 0 0 0 0
-0 3 3 0 0 4 4 0 0
-0 3 3 0 0 4 4 0 0
-0 0 0 0 0 0 0 0 0""",
-
-    "Border pattern": """5 5 5 5 5 5 5
-5 0 0 0 0 0 5
-5 0 3 3 3 0 5
-5 0 3 0 3 0 5
-5 0 3 3 3 0 5
-5 0 0 0 0 0 5
-5 5 5 5 5 5 5""",
-
-    "Diagonal": """1 0 0 0 0 0 0
-0 1 0 0 0 0 0
-0 0 1 0 0 0 0
-0 0 0 2 0 0 0
-0 0 0 0 2 0 0
-0 0 0 0 0 2 0
-0 0 0 0 0 0 3""",
-
-    "Scattered dots": """0 0 0 0 0 0 0 0
-0 1 0 0 0 2 0 0
-0 0 0 0 0 0 0 0
-0 0 0 3 0 0 0 0
-0 0 0 0 0 4 0 0
-0 5 0 0 0 0 0 0
-0 0 0 0 6 0 0 0
-0 0 0 0 0 0 0 0""",
-}
+        mask = (grid==c)
+        if mask.any(): n_comp += _connected_components(mask).max()
+    return dict(h_sym=float(h_sym), v_sym=float(v_sym),
+                boundary=b_cnt, edge=round(edge,4), components=int(n_comp))
+
+# ── Grid rendering ────────────────────────────────────────────────────────────
+
+def render_grid(grid, title='', highlight_wrong=None):
+    H,W = grid.shape
+    cell = max(32, min(60, 280//max(H,W)))
+    fig,ax = plt.subplots(figsize=((W*cell+4)/72, (H*cell+24)/72), dpi=72)
+    fig.patch.set_facecolor('#1e1e2e')
+    ax.set_facecolor('#1e1e2e')
+    ax.imshow(grid, cmap=ARC_CMAP, vmin=0, vmax=9,
+              interpolation='nearest', aspect='equal')
+    for x in range(W+1): ax.axvline(x-0.5, color='#555', lw=0.5)
+    for y in range(H+1): ax.axhline(y-0.5, color='#555', lw=0.5)
+    for r in range(H):
+        for c in range(W):
+            v = int(grid[r,c])
+            col = 'white' if v in [0,1,2,3,5,6,9] else 'black'
+            ax.text(c, r, str(v), ha='center', va='center',
+                    fontsize=max(7, cell//5), color=col,
+                    fontweight='bold', fontfamily='monospace')
+            if highlight_wrong is not None and highlight_wrong[r,c]:
+                ax.add_patch(plt.Rectangle(
+                    (c-0.5,r-0.5),1,1, fill=False, edgecolor='#ff4444', lw=2))
+    ax.set_xlim(-0.5, W-0.5); ax.set_ylim(H-0.5, -0.5); ax.axis('off')
+    if title:
+        ax.set_title(title, color='#cdd6f4', fontsize=9, pad=3)
+    plt.tight_layout(pad=0.3)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=72, bbox_inches='tight', facecolor='#1e1e2e')
+    buf.seek(0); img = Image.open(buf).copy(); plt.close()
+    return img
+
+def render_feature_maps(grid):
+    gx,gy = _sobel(grid.astype(np.float32)/9)
+    maps = [
+        ('H-symmetry',  _sym_map(grid,'h'),           'YlOrRd'),
+        ('V-symmetry',  _sym_map(grid,'v'),           'YlOrRd'),
+        ('Boundary',    _boundary(grid),               'plasma'),
+        ('Edge mag',    np.sqrt(gx**2+gy**2),         'hot'),
+    ]
+    fig,axes = plt.subplots(1,4,figsize=(12,2.8))
+    fig.patch.set_facecolor('#1e1e2e')
+    for ax,(title,data,cmap) in zip(axes,maps):
+        ax.set_facecolor('#0d0d1a')
+        vmax = 1 if data.max()<=1 else data.max()
+        im = ax.imshow(data, cmap=cmap, vmin=0, vmax=vmax,
+                       interpolation='nearest', aspect='equal')
+        ax.set_title(title, color='white', fontsize=9, pad=3)
+        ax.axis('off')
+        plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+    plt.tight_layout(pad=1)
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', dpi=100, bbox_inches='tight', facecolor='#1e1e2e')
+    buf.seek(0); img = Image.open(buf).copy(); plt.close()
+    return img
 
-# ── Parse grid text ───────────────────────────────────────────────────────────
+# ── Parse ─────────────────────────────────────────────────────────────────────
 
 def parse_grid(text):
     try:
@@ -198,169 +143,191 @@ def parse_grid(text):
         for line in text.strip().split('\n'):
             line = line.strip()
             if not line: continue
-            vals = [int(x) for x in line.split()]
-            rows.append(vals)
-        if not rows: return None, "No rows found"
+            rows.append([int(x) for x in line.split()])
+        if not rows: return None, "Empty"
         W = len(rows[0])
-        if any(len(r) != W for r in rows):
-            return None, "Rows have different lengths"
-        grid = np.array(rows, dtype=np.int64)
-        if grid.min() < 0 or grid.max() > 9:
-            return None, "Values must be 0–9"
-        return grid, None
+        if any(len(r)!=W for r in rows): return None, "Rows have different lengths"
+        g = np.array(rows, dtype=np.int64)
+        if g.min()<0 or g.max()>9: return None, "Values must be 0–9"
+        return g, None
     except Exception as e:
         return None, str(e)
 
-# ── Main analysis function ────────────────────────────────────────────────────
+def blank_like(grid_txt):
+    g, err = parse_grid(grid_txt)
+    if err or g is None: return ""
+    return "\n".join(" ".join(["0"]*g.shape[1]) for _ in range(g.shape[0]))
+
+# ── Puzzles ───────────────────────────────────────────────────────────────────
+
+PUZZLES = {
+    "Mirror complete": {
+        "desc": "Complete the right half to mirror the left half.",
+        "train_in":  "0 1 2 3 0 0 0 0\n0 4 0 3 0 0 0 0\n0 1 2 3 0 0 0 0",
+        "train_out": "0 1 2 3 3 2 1 0\n0 4 0 3 3 0 4 0\n0 1 2 3 3 2 1 0",
+        "test_in":   "0 2 1 0 0 0 0 0\n0 2 3 0 0 0 0 0\n0 0 1 0 0 0 0 0",
+        "answer":    "0 2 1 0 0 1 2 0\n0 2 3 0 0 3 2 0\n0 0 1 0 0 1 0 0",
+    },
+    "Color shift +1": {
+        "desc": "Each nonzero color increases by 1 (9 stays 9).",
+        "train_in":  "0 1 0\n2 0 3\n0 4 0",
+        "train_out": "0 2 0\n3 0 4\n0 5 0",
+        "test_in":   "0 3 1\n0 0 2\n4 0 0",
+        "answer":    "0 4 2\n0 0 3\n5 0 0",
+    },
+    "Border only": {
+        "desc": "Keep only the outer border, fill interior with 0.",
+        "train_in":  "2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2\n2 2 2 2 2",
+        "train_out": "2 2 2 2 2\n2 0 0 0 2\n2 0 0 0 2\n2 0 0 0 2\n2 2 2 2 2",
+        "test_in":   "3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3\n3 3 3 3 3",
+        "answer":    "3 3 3 3 3\n3 0 0 0 3\n3 0 0 0 3\n3 0 0 0 3\n3 3 3 3 3",
+    },
+    "Gravity down": {
+        "desc": "Each colored pixel falls to the bottom of its column.",
+        "train_in":  "1 0 2\n0 0 0\n0 3 0",
+        "train_out": "0 0 0\n0 0 2\n1 3 0",
+        "test_in":   "4 0 0\n0 5 0\n0 0 6",
+        "answer":    "0 0 0\n0 0 0\n4 5 6",
+    },
+    "Count → block": {
+        "desc": "Count the 1s; fill an N×N block of 2s in the bottom-left.",
+        "train_in":  "0 1 0 1 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0",
+        "train_out": "0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n2 2 0 0 0\n2 2 0 0 0",
+        "test_in":   "0 1 0 1 1\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0\n0 0 0 0 0",
+        "answer":    "0 0 0 0 0\n0 0 0 0 0\n2 2 2 0 0\n2 2 2 0 0\n2 2 2 0 0",
+    },
+}
 
-def analyze_grid(grid_text):
-    grid, err = parse_grid(grid_text)
-    if err:
-        return None, f"Parse error: {err}", ""
-
-    features = extract_features_visual(grid)
-    H, W = grid.shape
-
-    fig, axes = plt.subplots(2, 4, figsize=(14, 7))
-    fig.patch.set_facecolor('#1a1a2e')
-
-    titles = [
-        ("Input grid", grid, ARC_CMAP, 0, 9, False),
-        ("H-symmetry\n(mirror axes →)", features['h_symmetry'], 'YlOrRd', 0, 1, False),
-        ("V-symmetry\n(mirror axes ↕)", features['v_symmetry'], 'YlOrRd', 0, 1, False),
-        ("Boundary contour\n(Cauchy edge)", features['boundary'], 'plasma', 0, 1, False),
-        ("Edge magnitude\n(Sobel)", features['edge_magnitude'], 'hot', 0, 1, False),
-        ("Winding / curl\n(Im rotation)", features['winding_curl'], 'RdBu', -1, 1, True),
-        ("Component map\n(object IDs)", features['component_map'], 'tab20', 0, 1, False),
-    ]
+# ── Handlers ──────────────────────────────────────────────────────────────────
 
-    flat_axes = axes.flatten()
-    for i, ax in enumerate(flat_axes):
-        ax.set_facecolor('#0d0d1a')
-        if i < len(titles):
-            title, data, cmap, vmin, vmax, center = titles[i]
-            if cmap == ARC_CMAP:
-                im = ax.imshow(data, cmap=cmap, vmin=vmin, vmax=vmax,
-                               interpolation='nearest', aspect='equal')
-                for r in range(H):
-                    for c in range(W):
-                        v = int(data[r, c])
-                        color = 'white' if v in [0,6,7,8] else 'black'
-                        ax.text(c, r, str(v), ha='center', va='center',
-                                fontsize=max(6, 10 - max(H, W)//2),
-                                color=color, fontweight='bold')
-            else:
-                kwargs = dict(cmap=cmap, vmin=vmin, vmax=vmax,
-                              interpolation='nearest', aspect='equal')
-                im = ax.imshow(data, **kwargs)
-                plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
-
-            ax.set_title(title, color='white', fontsize=9, pad=4)
-            ax.tick_params(colors='#666', labelsize=7)
-            for spine in ax.spines.values():
-                spine.set_edgecolor('#333')
-        else:
-            ax.axis('off')
-
-    flat_axes[-1].axis('off')
-    flat_axes[-1].set_facecolor('#0d0d1a')
-
-    plt.tight_layout(pad=1.5)
+def on_load_puzzle(name):
+    p = PUZZLES[name]
+    ti_img, to_img, test_img = [
+        render_grid(parse_grid(t)[0], title=lbl)
+        for t, lbl in [(p["train_in"],"Training input"),
+                       (p["train_out"],"Training output"),
+                       (p["test_in"],"Test input — solve this")]
+    ]
+    blank = blank_like(p["answer"])
+    return (p["train_in"], p["train_out"], p["test_in"],
+            p["desc"], p["answer"],
+            blank,
+            ti_img, to_img, test_img,
+            None, None, None,
+            "Edit the answer grid above and click Submit.")
+
+def on_submit(answer_txt, answer_hidden, test_in_txt):
+    answer_g, err = parse_grid(answer_txt)
+    if err:
+        return None, None, None, f"**Parse error:** {err}"
 
-    buf = io.BytesIO()
-    plt.savefig(buf, format='png', dpi=120, bbox_inches='tight',
-                facecolor='#1a1a2e')
-    buf.seek(0)
-    img = Image.open(buf)
-    plt.close()
+    target_g, err2 = parse_grid(answer_hidden)
+    if err2:
+        return None, None, None, f"**Target parse error:** {err2}"
 
-    # Stats summary
-    h_max = features['h_symmetry'].max()
-    v_max = features['v_symmetry'].max()
-    boundary_count = int((features['boundary'] > 0.5).sum())
-    curl_max = float(np.abs(features['winding_curl']).max())
-    n_components = int(features['component_map'].max() * 20) if features['component_map'].max() > 0 else 0
+    if answer_g.shape != target_g.shape:
+        return None, None, None, (
+            f"**Wrong shape:** your answer is {answer_g.shape}, "
+            f"target is {target_g.shape}. Check row count and width.")
 
-    dominant = "H-symmetry" if h_max >= v_max else "V-symmetry"
-    sym_score = max(h_max, v_max)
+    correct = (answer_g == target_g)
+    pct = correct.mean() * 100
+    solved = bool(correct.all())
+    wrong = ~correct
 
-    summary = f"""**Grid:** {H}×{W}  |  **Colors used:** {len(np.unique(grid))-1} (excl. background)
+    answer_img = render_grid(answer_g, title="Your answer", highlight_wrong=wrong)
+    target_img = render_grid(target_g, title="Target output")
+    feat_img   = render_feature_maps(answer_g)
 
-**Im-side (orientation / structure):**
-- Strongest symmetry axis: {dominant} = {sym_score:.2f} {'★ strong' if sym_score > 0.8 else ''}
-- Boundary pixels (Cauchy edge): {boundary_count}
-- Winding / curl max: {curl_max:.3f}
-- Connected components: ~{n_components}
+    sc  = compute_scores(answer_g)
+    tsc = compute_scores(target_g)
 
-**What the CNN sees before training:**
-The 56-channel extractor separates *how big* (Re: color counts, distances) from *where pointed* (Im: symmetry axes, boundaries, rotation). High symmetry scores tell the CNN which actions are likely to preserve structure; high boundary counts indicate complex objects worth clicking on precisely."""
+    emoji = "✅" if solved else "❌"
+    score_md = f"""### {emoji} {"Solved!" if solved else "Not yet — keep trying"}
 
-    return img, summary, ""
+**Pixel accuracy: {pct:.1f}%** &nbsp;·&nbsp; {int(correct.sum())} / {answer_g.size} cells correct
 
+| Im-side feature | Your answer | Target |
+|---|---|---|
+| H-symmetry | {sc['h_sym']:.2f} | {tsc['h_sym']:.2f} |
+| V-symmetry | {sc['v_sym']:.2f} | {tsc['v_sym']:.2f} |
+| Boundary pixels | {sc['boundary']} | {tsc['boundary']} |
+| Components | {sc['components']} | {tsc['components']} |
+| Edge magnitude | {sc['edge']:.4f} | {tsc['edge']:.4f} |
 
-def load_preset(preset_name):
-    return PRESETS.get(preset_name, "")
+*These Im-side signals are exactly what the CNN reads from your grid before taking any action.*
+"""
+    return answer_img, target_img, feat_img, score_md
 
+def on_reset(answer_hidden):
+    return blank_like(answer_hidden)
 
-# ── Gradio UI ─────────────────────────────────────────────────────────────────
+# ── UI ────────────────────────────────────────────────────────────────────────
 
-CSS = """
-.gradio-container { max-width: 1100px !important; }
-#grid-input textarea { font-family: monospace; font-size: 14px; }
-"""
+with gr.Blocks(title="ARC-AGI-3 Puzzle Interface") as demo:
 
-with gr.Blocks(css=CSS, title="ARC-AGI-3 Feature Visualizer") as demo:
+    # State
+    s_train_in  = gr.State("")
+    s_train_out = gr.State("")
+    s_test_in   = gr.State("")
+    s_answer    = gr.State("")
 
     gr.Markdown("""
-# ARC-AGI-3 Feature Visualizer
-### Re/Im duality — what the CNN sees before any training
-
-Enter a grid of integers (0–9, space-separated rows) or pick a preset.
-The extractor computes **56 feature channels** split into:
-- **Re side** (local / multiplicative): one-hot colors, object sizes, distance maps
-- **Im side** (global / angular): symmetry axes, boundary contours, winding structure
+# ARC-AGI-3 Puzzle Interface
+Study the training example, figure out the rule, then complete the test input.
+Your pixel score updates on every submission — just like the real competition.
 """)
 
     with gr.Row():
-        with gr.Column(scale=1):
-            preset_dd = gr.Dropdown(
-                choices=list(PRESETS.keys()),
-                label="Load a preset puzzle",
-                value=None,
-            )
-            grid_input = gr.Textbox(
-                label="Grid (space-separated integers 0–9, one row per line)",
-                placeholder="0 1 2\n3 0 1\n2 3 0",
-                lines=10,
-                elem_id="grid-input",
-                value=PRESETS["Horizontal mirror"],
-            )
-            analyze_btn = gr.Button("Analyze", variant="primary")
-
-        with gr.Column(scale=2):
-            output_img = gr.Image(label="Feature maps", type="pil")
-            output_md  = gr.Markdown()
-            error_box  = gr.Textbox(label="", visible=False)
+        puzzle_dd = gr.Dropdown(choices=list(PUZZLES.keys()),
+                                value=list(PUZZLES.keys())[0],
+                                label="Select puzzle", scale=2)
+        desc_md = gr.Markdown(scale=3)
 
-    preset_dd.change(load_preset, inputs=preset_dd, outputs=grid_input)
-    analyze_btn.click(analyze_grid,
-                      inputs=grid_input,
-                      outputs=[output_img, output_md, error_box])
+    gr.Markdown("### Training example")
+    with gr.Row():
+        train_in_img  = gr.Image(label="Input",  type="pil", interactive=False, height=180)
+        train_out_img = gr.Image(label="Output", type="pil", interactive=False, height=180)
+        test_in_img   = gr.Image(label="Test input — what should the output be?",
+                                  type="pil", interactive=False, height=180)
 
-    gr.Markdown("""
----
-**How this works in ARC-AGI-3:** The agent receives a 64×64 grid frame from the game engine.
-Instead of feeding raw one-hot colors to the CNN, we pre-compute these 56 channels —
-giving the model geometric priors about symmetry and structure *before* it takes a single action.
+    gr.Markdown("---\n### Your answer")
+    gr.Markdown("Enter space-separated integers 0–9, one row per line. "
+                "**Colors:** " + " · ".join(f"`{i}`={n}" for i,n in enumerate(COLOR_NAMES)))
 
-Source: [Kaggle competition](https://www.kaggle.com/competitions/arc-prize-2026-arc-agi-3) |
-Agent based on [StochasticGoose](https://github.com/DriesSmit/ARC3-solution) by Dries Smit & Jack Cole (Tufa Labs)
-""")
+    with gr.Row():
+        with gr.Column(scale=1):
+            answer_box = gr.Textbox(label="Answer grid", lines=8,
+                                    placeholder="0 0 0\n0 0 0")
+            with gr.Row():
+                submit_btn = gr.Button("Submit", variant="primary")
+                reset_btn  = gr.Button("Reset to zeros")
 
-    demo.load(
-        fn=lambda: analyze_grid(PRESETS["Horizontal mirror"]),
-        outputs=[output_img, output_md, error_box],
-    )
+        with gr.Column(scale=2):
+            with gr.Row():
+                answer_img = gr.Image(label="Your answer", type="pil",
+                                      interactive=False, height=200)
+                target_img = gr.Image(label="Target",      type="pil",
+                                      interactive=False, height=200)
+            score_md = gr.Markdown("*Submit your answer to see your score.*")
+
+    gr.Markdown("---\n### Re/Im feature maps — what the CNN reads from your answer")
+    feat_img = gr.Image(label="Feature maps", type="pil", interactive=False)
+
+    # Wire up
+    all_outputs = [s_train_in, s_train_out, s_test_in, desc_md, s_answer,
+                   answer_box,
+                   train_in_img, train_out_img, test_in_img,
+                   answer_img, target_img, feat_img, score_md]
+
+    puzzle_dd.change(on_load_puzzle, inputs=puzzle_dd, outputs=all_outputs)
+    submit_btn.click(on_submit,
+                     inputs=[answer_box, s_answer, s_test_in],
+                     outputs=[answer_img, target_img, feat_img, score_md])
+    reset_btn.click(on_reset, inputs=s_answer, outputs=answer_box)
+    demo.load(on_load_puzzle, inputs=puzzle_dd, outputs=all_outputs)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(css="""
+        #answer-box textarea { font-family: monospace; font-size: 13px; }
+    """)