bank-code-arc-solution.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import os\n",
    "\n",
    "client = OpenAI(\n",
    "    base_url=\"https://openrouter.ai/api/v1\",\n",
    "    api_key=os.getenv(\"OPENROUTER_API_KEY\"),\n",
    ")\n",
    "\n",
    "def ask(messages):\n",
    "    output =  client.chat.completions.create(\n",
    "        model=\"anthropic/claude-3.5-sonnet\",\n",
    "        messages=messages,\n",
    "        temperature=0.5,\n",
    "        max_tokens=2048,\n",
    "        top_p=0.95,\n",
    "        stop=[\"</COLLECT>\", \"</EXPLORE>\", \"</COLLECT_ADD>\", \"</EXPLORE_ADD>\", \"</VALIDATE_SOLUTION>\", \"</SOLUTION>\"],\n",
    "    )\n",
    "\n",
    "    return output.choices[0].message.content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "import traceback\n",
    "import builtins\n",
    "\n",
    "timestamp = datetime.now().strftime(\"%Y%m%d-%H%M%S\")\n",
    "log_filename = f'attempts/attempt-{timestamp}.txt'\n",
    "\n",
    "# Redirect print statements to writing to a file\n",
    "def print(*args, **kwargs):\n",
    "    builtins.print(*args, **kwargs)\n",
    "    with open(log_filename, 'a') as f:\n",
    "        f.write(' '.join(map(str, args)) + '\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('puzzles/arc-agi_test_challenges.json', 'r') as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "# loading example puzzle\n",
    "data = data[\"1190e5a7\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train 0\n",
      "Input:\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠🟠\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "🟢🟠🟢🟢🟢🟢🟢🟢🟢🟢🟠🟢🟢🟠🟢\n",
      "Output:\n",
      "🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢\n",
      "\n",
      "Train 1\n",
      "Input:\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔹🔹🔹🔹🔹🔹🔹🔹🔹🔹🔹\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "🔹🔹🔹🔹🔹🔹🔹🔹🔹🔹🔹\n",
      "🔵🔵🔵🔵🔹🔵🔵🔵🔵🔵🔵\n",
      "Output:\n",
      "🔵🔵\n",
      "🔵🔵\n",
      "🔵🔵\n",
      "\n",
      "Train 2\n",
      "Input:\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "🟢🟢🟢🟢🟢🟢🔵🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🟢🔵🟢🔵🟢🔵🟢\n",
      "Output:\n",
      "🟢🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢🟢\n",
      "🟢🟢🟢🟢🟢\n",
      "\n",
      "Test 0\n",
      "Input:\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪⚪\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵🔵⚪🔵🔵🔵🔵⚪🔵\n",
      "Output:\n",
      "<find transformation to predict the test output>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "emoji_map = {\n",
    "    0: '⚫', 1: '🔵', 2: '🔴', 3: '🟢', 4: '🟡',\n",
    "    5: '⚪', 6: '🟣', 7: '🟠', 8: '🔹', 9: '🟥'\n",
    "}\n",
    "\n",
    "def grid_to_emoji(grid):\n",
    "    return '\\n'.join([''.join([emoji_map[cell] for cell in row]) for row in grid])\n",
    "    # return \"\\n\".join([\"\".join([str(ele) for ele in row]) for row in grid])\n",
    "\n",
    "add_to_prompt = \"\"\n",
    "\n",
    "for i in range(len(data[\"train\"])):\n",
    "    add_to_prompt += f\"Train {i}\\n\"\n",
    "    add_to_prompt += \"Input:\\n\"\n",
    "    add_to_prompt += grid_to_emoji(data[\"train\"][i][\"input\"]) + \"\\n\"\n",
    "    add_to_prompt += \"Output:\\n\"\n",
    "    add_to_prompt += grid_to_emoji(data[\"train\"][i][\"output\"]) + \"\\n\\n\"\n",
    "\n",
    "for i in range(len(data[\"test\"])):\n",
    "    add_to_prompt += f\"Test {i}\\n\"\n",
    "    add_to_prompt += \"Input:\\n\"\n",
    "    add_to_prompt += grid_to_emoji(data[\"test\"][i][\"input\"]) + \"\\n\"\n",
    "    add_to_prompt += \"Output:\\n\"\n",
    "    add_to_prompt += \"<find transformation to predict the test output>\\n\\n\"\n",
    "    \n",
    "print(add_to_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('prompts/bank_prompt.txt', 'r') as f:\n",
    "    initial_prompt =  f.read()\n",
    "\n",
    "initial_prompt = initial_prompt.split(\"<PUZZLE_DATA>\")[0]+\"<PUZZLE_DATA>\\n\"+add_to_prompt+\"\\n</PUZZLE_DATA>\"+initial_prompt.split(\"</PUZZLE_DATA>\")[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_initial_prompt(initial_prompt, explore_history, collect_history):\n",
    "    collect = \"\"\n",
    "    if(len(collect_history) > 0):\n",
    "        for idx, e in enumerate(collect_history):\n",
    "            collect += f\"<collect_history_item>{e}</collect_history_item>\\n\"\n",
    "    else:\n",
    "        collect = \"-- no collect history --\\n\"\n",
    "    \n",
    "    collect = \"<collect_history>\\n\"+collect+\"</collect_history>\"\n",
    "    initial_prompt = initial_prompt.split(\"<collect_history>\")[0]+collect+initial_prompt.split(\"</collect_history>\")[1]\n",
    "\n",
    "    explore = \"\"\n",
    "    if(len(explore_history) > 0):\n",
    "        for idx, e in enumerate(explore_history):\n",
    "            explore += f\"<explore_history_item>{e}</explore_history_item>\\n\"\n",
    "    else:\n",
    "        explore = \"-- no explore history --\\n\"\n",
    "\n",
    "    explore = \"<explore_history>\\n\"+explore+\"</explore_history>\"\n",
    "    initial_prompt = initial_prompt.split(\"<explore_history>\")[0]+explore+initial_prompt.split(\"</explore_history>\")[1]\n",
    "\n",
    "    return initial_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Let me start by collecting some basic information about the grid patterns. I'll first examine the structure and color patterns.\n",
      "\n",
      "<COLLECT>\n",
      "Let's analyze the dimensions and colors in each input/output pair\n",
      "<DSL>\n",
      "def run_dsl():\n",
      "    results = []\n",
      "    for idx in range(len(data[\"train\"])):\n",
      "        input_grid = data[\"train\"][idx][\"input\"]\n",
      "        output_grid = data[\"train\"][idx][\"output\"]\n",
      "        \n",
      "        results.append({\n",
      "            f\"Train {idx}\": {\n",
      "                \"input_dims\": f\"{len(input_grid)}x{len(input_grid[0])}\",\n",
      "                \"output_dims\": f\"{len(output_grid)}x{len(output_grid[0])}\",\n",
      "                \"input_colors\": sorted(list(set(sum(input_grid, [])))),\n",
      "                \"output_colors\": sorted(list(set(sum(output_grid, []))))\n",
      "            }\n",
      "        })\n",
      "    return results\n",
      "</DSL>\n",
      "\n",
      "</COLLECT>\n",
      "\n",
      "\n",
      "Output: [{'Train 0': {'input_dims': '15x15', 'output_dims': '2x4', 'input_colors': [3, 7], 'output_colors': [3]}}, {'Train 1': {'input_dims': '11x11', 'output_dims': '3x2', 'input_colors': [1, 8], 'output_colors': [1]}}, {'Train 2': {'input_dims': '27x27', 'output_dims': '6x5', 'input_colors': [1, 3], 'output_colors': [3]}}]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "From this initial collection, I notice some interesting patterns that I should add to our knowledge:\n",
      "\n",
      "<COLLECT_ADD>\n",
      "1. Each input grid contains exactly 2 colors, while output grids contain just 1 color\n",
      "2. The output color is always one of the colors from the input (specifically the non-line color)\n",
      "3. Input grids are always square (15x15, 11x11, 27x27)\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "Let me analyze this puzzle systematically. Looking at the grids visually, I notice that:\n",
      "1. There are horizontal lines of one color (orange/light blue/white) that divide the grid into sections\n",
      "2. The output seems to be a smaller grid made up of the dominant color (non-line color)\n",
      "3. The output dimensions might be related to the number of sections created by the horizontal lines\n",
      "\n",
      "Let me explore this hypothesis:\n",
      "\n",
      "<EXPLORE>\n",
      "Let me count the number of sections created by horizontal lines and compare with output dimensions\n",
      "<DSL>\n",
      "def run_dsl():\n",
      "    results = []\n",
      "    for idx in range(3):  # Check all training examples\n",
      "        input_grid = data[\"train\"][idx][\"input\"]\n",
      "        output_grid = data[\"train\"][idx][\"output\"]\n",
      "        \n",
      "        # Find the horizontal line color (by checking middle row which always has lines)\n",
      "        middle_row = input_grid[len(input_grid)//2]\n",
      "        line_color = middle_row[0]  # The line color is always first in these rows\n",
      "        \n",
      "        # Count sections by counting groups of non-line rows\n",
      "        section_count = 1\n",
      "        prev_was_line = False\n",
      "        for row in input_grid:\n",
      "            if all(cell == line_color for cell in row):\n",
      "                if not prev_was_line:\n",
      "                    section_count += 1\n",
      "                prev_was_line = True\n",
      "            else:\n",
      "                prev_was_line = False\n",
      "        \n",
      "        # Get output dimensions\n",
      "        output_height = len(output_grid)\n",
      "        output_width = len(output_grid[0])\n",
      "        \n",
      "        results.append({\n",
      "            'train_idx': idx,\n",
      "            'sections': section_count,\n",
      "            'output_dims': (output_height, output_width),\n",
      "            'visualization': grid_to_emoji(output_grid)\n",
      "        })\n",
      "    \n",
      "    return results\n",
      "</DSL>\n",
      "\n",
      "\n",
      "\n",
      "Output: [{'train_idx': 0, 'sections': 1, 'output_dims': (2, 4), 'visualization': '🟢🟢🟢🟢\\n🟢🟢🟢🟢'}, {'train_idx': 1, 'sections': 1, 'output_dims': (3, 2), 'visualization': '🔵🔵\\n🔵🔵\\n🔵🔵'}, {'train_idx': 2, 'sections': 1, 'output_dims': (6, 5), 'visualization': '🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢'}]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "I notice that my initial hypothesis about sections wasn't quite right - each example shows section_count as 1 because I need to fix my counting logic. Let me try a different exploration - I think the key might be counting the number of horizontal lines to determine output dimensions.\n",
      "\n",
      "<EXPLORE>\n",
      "Let me count the actual horizontal lines and see if they relate to output dimensions\n",
      "<DSL>\n",
      "def run_dsl():\n",
      "    results = []\n",
      "    for idx in range(3):\n",
      "        input_grid = data[\"train\"][idx][\"input\"]\n",
      "        output_grid = data[\"train\"][idx][\"output\"]\n",
      "        \n",
      "        # Find the line color (by checking the row we know is a line)\n",
      "        for row in input_grid:\n",
      "            if len(set(row)) == 1:  # If row has all same color, that's our line color\n",
      "                line_color = row[0]\n",
      "                break\n",
      "        \n",
      "        # Count horizontal lines\n",
      "        line_count = 0\n",
      "        for row in input_grid:\n",
      "            if all(cell == line_color for cell in row):\n",
      "                line_count += 1\n",
      "                \n",
      "        # Get the non-line color (the color that appears in non-line rows)\n",
      "        other_color = None\n",
      "        for row in input_grid:\n",
      "            if not all(cell == line_color for cell in row):\n",
      "                other_color = next(cell for cell in row if cell != line_color)\n",
      "                break\n",
      "        \n",
      "        results.append({\n",
      "            'train_idx': idx,\n",
      "            'horizontal_lines': line_count,\n",
      "            'line_color': line_color,\n",
      "            'other_color': other_color,\n",
      "            'output_dims': (len(output_grid), len(output_grid[0])),\n",
      "            'visualization': grid_to_emoji(output_grid)\n",
      "        })\n",
      "    \n",
      "    return results\n",
      "</DSL>\n",
      "\n",
      "\n",
      "\n",
      "Output: [{'train_idx': 0, 'horizontal_lines': 1, 'line_color': 7, 'other_color': 3, 'output_dims': (2, 4), 'visualization': '🟢🟢🟢🟢\\n🟢🟢🟢🟢'}, {'train_idx': 1, 'horizontal_lines': 2, 'line_color': 8, 'other_color': 1, 'output_dims': (3, 2), 'visualization': '🔵🔵\\n🔵🔵\\n🔵🔵'}, {'train_idx': 2, 'horizontal_lines': 5, 'line_color': 1, 'other_color': 3, 'output_dims': (6, 5), 'visualization': '🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢'}]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "I notice a potential pattern! Let me analyze what I see:\n",
      "\n",
      "Train 0: 1 horizontal line -> output is 2x4\n",
      "Train 1: 2 horizontal lines -> output is 3x2\n",
      "Train 2: 5 horizontal lines -> output is 6x5\n",
      "\n",
      "I think I see the pattern:\n",
      "- The number of horizontal lines (n) determines the output dimensions\n",
      "- Output height = n+1\n",
      "- Output width seems to be related to the number of vertical sections between vertical lines in the input\n",
      "\n",
      "Let me validate this hypothesis:\n",
      "\n",
      "<VALIDATE_SOLUTION>\n",
      "Testing if output dimensions are determined by horizontal lines and vertical sections\n",
      "<DSL>\n",
      "def run_dsl():\n",
      "    def get_vertical_sections(grid, line_color):\n",
      "        # Look at first non-line row to count vertical sections\n",
      "        for row in grid:\n",
      "            if not all(cell == line_color for cell in row):\n",
      "                test_row = row\n",
      "                break\n",
      "        \n",
      "        section_count = 1\n",
      "        for i in range(1, len(test_row)):\n",
      "            if test_row[i] == line_color and test_row[i-1] != line_color:\n",
      "                section_count += 1\n",
      "        return section_count\n",
      "\n",
      "    results = []\n",
      "    for idx in range(3):\n",
      "        input_grid = data[\"train\"][idx][\"input\"]\n",
      "        output_grid = data[\"train\"][idx][\"output\"]\n",
      "        \n",
      "        # Find line color\n",
      "        for row in input_grid:\n",
      "            if len(set(row)) == 1:\n",
      "                line_color = row[0]\n",
      "                break\n",
      "        \n",
      "        # Count horizontal lines\n",
      "        h_lines = sum(1 for row in input_grid if all(cell == line_color for cell in row))\n",
      "        \n",
      "        # Get vertical sections\n",
      "        v_sections = get_vertical_sections(input_grid, line_color)\n",
      "        \n",
      "        # Get actual output dimensions\n",
      "        actual_height = len(output_grid)\n",
      "        actual_width = len(output_grid[0])\n",
      "        \n",
      "        # Predicted dimensions\n",
      "        predicted_height = h_lines + 1\n",
      "        predicted_width = v_sections\n",
      "        \n",
      "        results.append({\n",
      "            'train_idx': idx,\n",
      "            'horizontal_lines': h_lines,\n",
      "            'vertical_sections': v_sections,\n",
      "            'predicted_dims': (predicted_height, predicted_width),\n",
      "            'actual_dims': (actual_height, actual_width),\n",
      "            'matches': (predicted_height == actual_height and predicted_width == actual_width),\n",
      "            'output': grid_to_emoji(output_grid)\n",
      "        })\n",
      "    \n",
      "    return results\n",
      "</DSL>\n",
      "\n",
      "\n",
      "\n",
      "Output: [{'train_idx': 0, 'horizontal_lines': 1, 'vertical_sections': 4, 'predicted_dims': (2, 4), 'actual_dims': (2, 4), 'matches': True, 'output': '🟢🟢🟢🟢\\n🟢🟢🟢🟢'}, {'train_idx': 1, 'horizontal_lines': 2, 'vertical_sections': 2, 'predicted_dims': (3, 2), 'actual_dims': (3, 2), 'matches': True, 'output': '🔵🔵\\n🔵🔵\\n🔵🔵'}, {'train_idx': 2, 'horizontal_lines': 5, 'vertical_sections': 5, 'predicted_dims': (6, 5), 'actual_dims': (6, 5), 'matches': True, 'output': '🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢\\n🟢🟢🟢🟢🟢'}]\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "----------------------------------------------------------------------------------------------------\n",
      "\n",
      "\n",
      "\n",
      "The validation shows my hypothesis is correct! Let me explain the complete pattern:\n",
      "\n",
      "1. The output dimensions are determined by:\n",
      "   - Height = number of horizontal lines + 1\n",
      "   - Width = number of vertical sections (counting regions between vertical lines)\n",
      "2. The output color is the non-line color from the input grid\n",
      "\n",
      "Let me write a complete solution:\n",
      "\n",
      "<SOLUTION>\n",
      "<DSL>\n",
      "def run_dsl():\n",
      "    def get_output_dimensions(input_grid, line_color):\n",
      "        # Count horizontal lines\n",
      "        h_lines = sum(1 for row in input_grid if all(cell == line_color for cell in row))\n",
      "        \n",
      "        # Count vertical sections by looking at first non-line row\n",
      "        for row in input_grid:\n",
      "            if not all(cell == line_color for cell in row):\n",
      "                test_row = row\n",
      "                break\n",
      "        \n",
      "        section_count = 1\n",
      "        for i in range(1, len(test_row)):\n",
      "            if test_row[i] == line_color and test_row[i-1] != line_color:\n",
      "                section_count += 1\n",
      "                \n",
      "        return (h_lines + 1, section_count)\n",
      "    \n",
      "    def get_grid_colors(input_grid):\n",
      "        # Get the two colors in the grid\n",
      "        colors = set()\n",
      "        for row in input_grid:\n",
      "            colors.update(row)\n",
      "        \n",
      "        # Find line color (the color that makes up complete rows)\n",
      "        for row in input_grid:\n",
      "            if len(set(row)) == 1:\n",
      "                line_color = row[0]\n",
      "                other_color = (colors - {line_color}).pop()\n",
      "                return line_color, other_color\n",
      "    \n",
      "    outputs = []\n",
      "    \n",
      "    # Process each input grid\n",
      "    for idx in range(len(data[\"test\" if \"test\" in data else \"train\"])):\n",
      "        input_grid = data[\"test\"][idx][\"input\"] if \"test\" in data else data[\"train\"][idx][\"input\"]\n",
      "        \n",
      "        # Get colors and dimensions\n",
      "        line_color, fill_color = get_grid_colors(input_grid)\n",
      "        height, width = get_output_dimensions(input_grid, line_color)\n",
      "        \n",
      "        # Create output grid filled with the non-line color\n",
      "        output_grid = [[fill_color for _ in range(width)] for _ in range(height)]\n",
      "        outputs.append(output_grid)\n",
      "    \n",
      "    return outputs\n",
      "</DSL>\n",
      "\n",
      "SOLUTION FOUND:\n",
      "[[[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1]]]\n"
     ]
    }
   ],
   "source": [
    "explore_history = []\n",
    "collect_history = []\n",
    "\n",
    "initial_prompt = compute_initial_prompt(initial_prompt, explore_history, collect_history)\n",
    "messages = [\n",
    "    {\n",
    "        \"role\": \"user\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\",\n",
    "                \"text\": initial_prompt\n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "]\n",
    "\n",
    "solution_found = False  \n",
    "\n",
    "for i in range(100):\n",
    "    response = ask(messages)\n",
    "    print(response)\n",
    "\n",
    "    def run_dsl():\n",
    "        return \"< run_dsl function not implemented >\"\n",
    "    \n",
    "    if(\"<EXPLORE>\" in response):\n",
    "        response += \"</EXPLORE>\"\n",
    "        string_output_from_claude = response.split(\"<EXPLORE>\")[1].split(\"</EXPLORE>\")[0] # full output where you generated DSL code\n",
    "        code = string_output_from_claude.split(\"<DSL>\")[1].split(\"</DSL>\")[0].strip().strip(\"\\n\") # extract just the code\n",
    "        try:\n",
    "            exec(code) # run eval on your code to use the run_dsl function that you wrote\n",
    "            output = run_dsl() # this function must have been written by you, and created when i run eval\n",
    "        except Exception as e:\n",
    "            output = f\"Error: {str(e)}\\n\\nFull traceback:\\n{traceback.format_exc()}\"\n",
    "        print(f\"\\n\\nOutput: {output}\\n\\n\")\n",
    "        messages.extend([\n",
    "            {\n",
    "                \"role\": \"assistant\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": response\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": \"<DSL_OUTPUT>\\n\"+str(output)+\"\\n</DSL_OUTPUT>\\n\\n Usually, you would now perform a explore add step. But, if you have an idea of the solution, you should perform a validate solution step.\"\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ])\n",
    "\n",
    "    if(\"<VALIDATE_SOLUTION>\" in response):\n",
    "        response += \"</VALIDATE_SOLUTION>\"\n",
    "        string_output_from_claude = response.split(\"<VALIDATE_SOLUTION>\")[1].split(\"</VALIDATE_SOLUTION>\")[0] # full output where you generated DSL code\n",
    "        code = string_output_from_claude.split(\"<DSL>\")[1].split(\"</DSL>\")[0].strip().strip(\"\\n\") # extract just the code\n",
    "        try:\n",
    "            exec(code) # run eval on your code to use the run_dsl function that you wrote\n",
    "            output = run_dsl() # this function must have been written by you, and created when i run eval\n",
    "        except Exception as e:\n",
    "            output = f\"Error: {str(e)}\\n\\nFull traceback:\\n{traceback.format_exc()}\"\n",
    "        print(f\"\\n\\nOutput: {output}\\n\\n\")\n",
    "        messages.extend([\n",
    "            {\n",
    "                \"role\": \"assistant\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": response\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": \"<DSL_OUTPUT>\\n\"+str(output)+\"\\n</DSL_OUTPUT>\\nUsually, you would attempt to validate a solution based on the information and patterns that you recognised. If you are not sure, you would perform an explore add step.\"\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ])\n",
    "\n",
    "    if(\"<SOLUTION>\" in response):\n",
    "        response += \"</SOLUTION>\"\n",
    "        string_output_from_claude = response.split(\"<SOLUTION>\")[1].split(\"</SOLUTION>\")[0] # full output where you generated DSL code\n",
    "        code = string_output_from_claude.split(\"<DSL>\")[1].split(\"</DSL>\")[0].strip().strip(\"\\n\") # extract just the code\n",
    "        try:\n",
    "            exec(code) # run eval on your code to use the run_dsl function that you wrote\n",
    "            output = run_dsl() # this function must have been written by you, and created when i run eval\n",
    "        except Exception as e:\n",
    "            output = f\"Error: {str(e)}\\n\\nFull traceback:\\n{traceback.format_exc()}\"\n",
    "            print(f\"\\n\\nOutput: {output}\\n\\n\")\n",
    "            messages.extend([\n",
    "                {\n",
    "                    \"role\": \"assistant\",\n",
    "                    \"content\": [\n",
    "                        {\n",
    "                            \"type\": \"text\",\n",
    "                            \"text\": response\n",
    "                        }\n",
    "                    ]\n",
    "                },\n",
    "                {\n",
    "                    \"role\": \"user\",\n",
    "                    \"content\": [\n",
    "                        {\n",
    "                            \"type\": \"text\",\n",
    "                            \"text\": \"<DSL_OUTPUT>\\n\"+str(output)+\"\\n</DSL_OUTPUT>\"\n",
    "                        },\n",
    "                    ],\n",
    "                }\n",
    "            ])\n",
    "            continue\n",
    "\n",
    "        print(\"SOLUTION FOUND:\")\n",
    "        print(output)\n",
    "        solution_found = True\n",
    "        break\n",
    "\n",
    "    elif(\"<COLLECT>\" in response):\n",
    "        response += \"</COLLECT>\"\n",
    "        print(\"</COLLECT>\")\n",
    "        string_output_from_claude = response.split(\"<COLLECT>\")[1].split(\"</COLLECT>\")[0]\n",
    "        code = string_output_from_claude.split(\"<DSL>\")[1].split(\"</DSL>\")[0].strip().strip(\"\\n\")\n",
    "        try:\n",
    "            exec(code) # run eval on your code to use the run_dsl function that you wrote\n",
    "            output = run_dsl() # this function must have been written by you, and created when i run eval\n",
    "        except Exception as e:\n",
    "            output = f\"Error: {str(e)}\\n\\nFull traceback:\\n{traceback.format_exc()}\"\n",
    "        print(f\"\\n\\nOutput: {output}\\n\\n\")\n",
    "        messages.extend([\n",
    "            {\n",
    "                \"role\": \"assistant\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": response\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": \"<DSL_OUTPUT>\\n\"+str(output)+\"\\n</DSL_OUTPUT>\\n\\n Usually, you would attempt to validate a solution based on the information and patterns that you recognised. If you are not sure, you would perform an collect add step.\"\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ])\n",
    "\n",
    "    elif(\"<EXPLORE_ADD>\" in response):\n",
    "        response += \"</EXPLORE_ADD>\"\n",
    "        print(\"</EXPLORE_ADD>\")\n",
    "        string_output_from_claude = response.split(\"<EXPLORE_ADD>\")[1].split(\"</EXPLORE_ADD>\")[0]\n",
    "        explore_history.append(string_output_from_claude)\n",
    "        initial_prompt = compute_initial_prompt(initial_prompt, explore_history, collect_history)\n",
    "        messages = [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": initial_prompt\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ]\n",
    "\n",
    "    elif(\"<COLLECT_ADD>\" in response):\n",
    "        response += \"</COLLECT_ADD>\"\n",
    "        string_output_from_claude = response.split(\"<COLLECT_ADD>\")[1].split(\"</COLLECT_ADD>\")[0]\n",
    "        collect_history.append(string_output_from_claude)\n",
    "        initial_prompt = compute_initial_prompt(initial_prompt, explore_history, collect_history)\n",
    "        messages = [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": initial_prompt\n",
    "                    },\n",
    "                ],\n",
    "            }\n",
    "        ]\n",
    "    \n",
    "    print(\"\\n\\n\")\n",
    "    print(\"-\"*100)\n",
    "    print(\"\\n\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution found!\n",
      "[[[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1]]]\n",
      "🔵🔵🔵\n",
      "🔵🔵🔵\n",
      "🔵🔵🔵\n",
      "🔵🔵🔵\n",
      "🔵🔵🔵\n"
     ]
    }
   ],
   "source": [
    "if(solution_found):\n",
    "    print(\"Solution found!\")\n",
    "    print(output)\n",
    "    print(grid_to_emoji(output[0]))\n",
    "else:\n",
    "    print(\"Solution not found.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "explore_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['\\n1. Each input grid contains exactly 2 colors, while output grids contain just 1 color\\n2. The output color is always one of the colors from the input (specifically the non-line color)\\n3. Input grids are always square (15x15, 11x11, 27x27)\\n']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "collect_history"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}