iterative-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=4096,\n",
    "        top_p=0.95,\n",
    "        stop=[\"</DSL>\", \"</SOLUTION>\", \"</VALIDATE_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 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[\"1f85a75f\"]"
   ]
  },
  {
   "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",
      "⚪⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚪⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫🔵\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚪⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫🔵⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚪⚫⚫⚫⚫🔵⚫⚪⚫⚫⚫⚫⚫⚫🔵\n",
      "⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚪⚫⚫⚫⚫⚫\n",
      "🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫🔵⚫⚫⚫⚫🔵\n",
      "⚫⚫⚫⚫⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫🔵⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔵⚫🔵⚫⚫⚪\n",
      "⚫⚫⚫⚫⚫⚫⚫🔵🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫⚪⚫⚫⚫⚪⚫⚫⚫⚪⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚪⚫⚫⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "Output:\n",
      "🟢🟢⚫\n",
      "🟢🟢🟢\n",
      "🟢⚫🟢\n",
      "🟢🟢🟢\n",
      "⚫🟢🟢\n",
      "\n",
      "Train 1\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",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔴🔴⚫⚫🔴⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🔴⚫⚫\n",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "Output:\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",
      "⚫⚫⚫⚫⚫⚫🔹⚫⚫⚫⚫⚫🟢⚫🔵⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🟢\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/iterative_2_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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "----------------------------------------------------------------------------------------------------\n",
      "Iteration 1\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Response received:\n",
      "Let me analyze the puzzle step by step.\n",
      "\n",
      "1) First, let me look at the training examples:\n",
      "- Train 0: Has a 3x5 green (🟢) pattern in the middle of a grid with blue (🔵) and white (⚪) background\n",
      "- Train 1: Has a 3x3 yellow (🟡) pattern in the middle of a grid with red (🔴) background\n",
      "\n",
      "2) Looking at Test 0:\n",
      "- Has a similar grid with blue (🔵), green (🟢), light blue (🔹) background\n",
      "- And notably has a red (🔴) pattern in the middle forming what appears to be a 3x3 shape\n",
      "\n",
      "3) Key observations:\n",
      "- Each input has a distinct pattern in the middle\n",
      "- The output seems to be that middle pattern extracted and preserved\n",
      "- The output dimensions match the pattern size found in the input\n",
      "\n",
      "Let me validate this hypothesis:\n",
      "\n",
      "<DSL>\n",
      "def dsl_code():\n",
      "    def find_pattern(grid):\n",
      "        # Get dimensions\n",
      "        h = height(grid)\n",
      "        w = width(grid)\n",
      "        \n",
      "        # Find all objects\n",
      "        objs = objects(grid, True, False, True)\n",
      "        \n",
      "        # Look for a compact object near center\n",
      "        center_row = h // 2\n",
      "        center_col = w // 2\n",
      "        \n",
      "        # Find object closest to center that forms a rectangle\n",
      "        min_dist = float('inf')\n",
      "        target_obj = None\n",
      "        \n",
      "        for obj in objs:\n",
      "            # Get object bounds\n",
      "            obj_indices = toindices(obj)\n",
      "            if size(obj_indices) < 4:  # Skip tiny objects\n",
      "                continue\n",
      "                \n",
      "            # Check if forms rectangle\n",
      "            min_row = minimum(apply(first, obj_indices))\n",
      "            max_row = maximum(apply(first, obj_indices))\n",
      "            min_col = minimum(apply(last, obj_indices))\n",
      "            max_col = maximum(apply(last, obj_indices))\n",
      "            \n",
      "            obj_center_row = (min_row + max_row) // 2\n",
      "            obj_center_col = (min_col + max_col) // 2\n",
      "            \n",
      "            dist = abs(obj_center_row - center_row) + abs(obj_center_col - center_col)\n",
      "            \n",
      "            if dist < min_dist:\n",
      "                min_dist = dist\n",
      "                target_obj = obj\n",
      "        \n",
      "        if target_obj is None:\n",
      "            return None\n",
      "            \n",
      "        # Get bounds of target object\n",
      "        indices = toindices(target_obj)\n",
      "        min_row = minimum(apply(first, indices))\n",
      "        max_row = maximum(apply(first, indices))\n",
      "        min_col = minimum(apply(last, indices))\n",
      "        max_col = maximum(apply(last, indices))\n",
      "        \n",
      "        # Extract subgrid\n",
      "        return crop(grid, asIterable(min_row, min_col), \n",
      "                   asIterable(max_row - min_row + 1, max_col - min_col + 1))\n",
      "    \n",
      "    # Test on training examples\n",
      "    train_results = []\n",
      "    for i in range(2):\n",
      "        input_grid = data[\"train\"][i][\"input\"]\n",
      "        expected = data[\"train\"][i][\"output\"]\n",
      "        result = find_pattern(input_grid)\n",
      "        train_results.append((result, expected))\n",
      "    \n",
      "    return train_results\n",
      "</DSL>\n",
      "Code output:\n",
      "[(([3, 3, 0], [3, 3, 3], [3, 0, 3], [3, 3, 3], [0, 3, 3]), [[3, 3, 0], [3, 3, 3], [3, 0, 3], [3, 3, 3], [0, 3, 3]]), (([0, 4, 0], [4, 4, 4], [0, 4, 4]), [[0, 4, 0], [4, 4, 4], [0, 4, 4]])]\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Iteration 2\n",
      "----------------------------------------------------------------------------------------------------\n",
      "Response received:\n",
      "Let me analyze this step by step:\n",
      "\n",
      "1) I notice that in both training examples:\n",
      "- Train 0 has a 3x5 green (🟢) pattern in the middle\n",
      "- Train 1 has a 3x3 yellow (🟡) pattern in the middle\n",
      "\n",
      "2) Looking at Test 0:\n",
      "- Has a similar grid layout\n",
      "- Contains a 3x3 red (🔴) pattern in the middle\n",
      "\n",
      "3) The pattern seems to be:\n",
      "- Find a solid colored shape in the middle of the grid\n",
      "- Extract that shape exactly as it appears\n",
      "- Preserve the dimensions and colors of the found pattern\n",
      "\n",
      "Let me write a solution:\n",
      "\n",
      "<SOLUTION>\n",
      "def dsl_code():\n",
      "    def find_pattern(grid):\n",
      "        # Get grid dimensions\n",
      "        h = height(grid)\n",
      "        w = width(grid)\n",
      "        \n",
      "        # Find center region\n",
      "        center_row = h // 2\n",
      "        center_col = w // 2\n",
      "        \n",
      "        # Get all objects\n",
      "        objs = objects(grid, True, False, True)\n",
      "        \n",
      "        # Find the object closest to center\n",
      "        min_dist = float('inf')\n",
      "        target_obj = None\n",
      "        \n",
      "        for obj in objs:\n",
      "            # Get object bounds\n",
      "            indices = toindices(obj)\n",
      "            if size(indices) < 4:  # Skip tiny objects\n",
      "                continue\n",
      "                \n",
      "            # Calculate center of object\n",
      "            min_row = minimum(apply(first, indices))\n",
      "            max_row = maximum(apply(first, indices))\n",
      "            min_col = minimum(apply(last, indices))\n",
      "            max_col = maximum(apply(last, indices))\n",
      "            \n",
      "            obj_center_row = (min_row + max_row) // 2\n",
      "            obj_center_col = (min_col + max_col) // 2\n",
      "            \n",
      "            # Calculate distance to grid center\n",
      "            dist = abs(obj_center_row - center_row) + abs(obj_center_col - center_col)\n",
      "            \n",
      "            if dist < min_dist:\n",
      "                min_dist = dist\n",
      "                target_obj = obj\n",
      "        \n",
      "        if target_obj is None:\n",
      "            return None\n",
      "            \n",
      "        # Extract the pattern\n",
      "        indices = toindices(target_obj)\n",
      "        min_row = minimum(apply(first, indices))\n",
      "        max_row = maximum(apply(first, indices))\n",
      "        min_col = minimum(apply(last, indices))\n",
      "        max_col = maximum(apply(last, indices))\n",
      "        \n",
      "        # Return the cropped pattern\n",
      "        return crop(grid, asIterable(min_row, min_col), \n",
      "                   asIterable(max_row - min_row + 1, max_col - min_col + 1))\n",
      "    \n",
      "    # Process each test input\n",
      "    results = []\n",
      "    for i in range(1):  # Only 1 test case\n",
      "        test_grid = data[\"test\"][i][\"input\"]\n",
      "        result = find_pattern(test_grid)\n",
      "        results.append(result)\n",
      "    \n",
      "    return results[0]  # Return the result for test case 0\n",
      "</SOLUTION>\n",
      "Solution found:\n",
      "\n",
      "def dsl_code():\n",
      "    def find_pattern(grid):\n",
      "        # Get grid dimensions\n",
      "        h = height(grid)\n",
      "        w = width(grid)\n",
      "        \n",
      "        # Find center region\n",
      "        center_row = h // 2\n",
      "        center_col = w // 2\n",
      "        \n",
      "        # Get all objects\n",
      "        objs = objects(grid, True, False, True)\n",
      "        \n",
      "        # Find the object closest to center\n",
      "        min_dist = float('inf')\n",
      "        target_obj = None\n",
      "        \n",
      "        for obj in objs:\n",
      "            # Get object bounds\n",
      "            indices = toindices(obj)\n",
      "            if size(indices) < 4:  # Skip tiny objects\n",
      "                continue\n",
      "                \n",
      "            # Calculate center of object\n",
      "            min_row = minimum(apply(first, indices))\n",
      "            max_row = maximum(apply(first, indices))\n",
      "            min_col = minimum(apply(last, indices))\n",
      "            max_col = maximum(apply(last, indices))\n",
      "            \n",
      "            obj_center_row = (min_row + max_row) // 2\n",
      "            obj_center_col = (min_col + max_col) // 2\n",
      "            \n",
      "            # Calculate distance to grid center\n",
      "            dist = abs(obj_center_row - center_row) + abs(obj_center_col - center_col)\n",
      "            \n",
      "            if dist < min_dist:\n",
      "                min_dist = dist\n",
      "                target_obj = obj\n",
      "        \n",
      "        if target_obj is None:\n",
      "            return None\n",
      "            \n",
      "        # Extract the pattern\n",
      "        indices = toindices(target_obj)\n",
      "        min_row = minimum(apply(first, indices))\n",
      "        max_row = maximum(apply(first, indices))\n",
      "        min_col = minimum(apply(last, indices))\n",
      "        max_col = maximum(apply(last, indices))\n",
      "        \n",
      "        # Return the cropped pattern\n",
      "        return crop(grid, asIterable(min_row, min_col), \n",
      "                   asIterable(max_row - min_row + 1, max_col - min_col + 1))\n",
      "    \n",
      "    # Process each test input\n",
      "    results = []\n",
      "    for i in range(1):  # Only 1 test case\n",
      "        test_grid = data[\"test\"][i][\"input\"]\n",
      "        result = find_pattern(test_grid)\n",
      "        results.append(result)\n",
      "    \n",
      "    return results[0]  # Return the result for test case 0\n",
      "\n",
      "Solution output:\n",
      "([0, 2, 2, 2], [2, 2, 0, 2], [2, 2, 0, 2], [0, 2, 2, 2])\n"
     ]
    }
   ],
   "source": [
    "import traceback\n",
    "from dsl import *\n",
    "\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",
    "ITERATIONS = 100\n",
    "for i in range(ITERATIONS):\n",
    "    print(f\"{'-'*100}\\nIteration {i+1}\\n{'-'*100}\")\n",
    "    response = ask(messages)\n",
    "\n",
    "    if(\"<SOLUTION>\" in response):\n",
    "        response += \"</SOLUTION>\"\n",
    "    elif(\"<DSL>\" in response):\n",
    "        response += \"</DSL>\"\n",
    "    elif(\"<VALIDATE_SOLUTION>\" in response):\n",
    "        response += \"</VALIDATE_SOLUTION>\"\n",
    "\n",
    "    print(\"Response received:\")\n",
    "    print(response)\n",
    "    messages.append({\n",
    "        \"role\": \"system\",\n",
    "        \"content\": [\n",
    "            {\n",
    "                \"type\": \"text\",\n",
    "                \"text\": response\n",
    "            }\n",
    "        ]\n",
    "    })\n",
    "\n",
    "    if \"<SOLUTION>\" in response:\n",
    "        solution = response.split(\"<SOLUTION>\")[1].split(\"</SOLUTION>\")[0]\n",
    "        print(\"Solution found:\")\n",
    "        print(solution)\n",
    "        def dsl_code():\n",
    "            return \"-- dsl_code() function not defined() function not defined --\"\n",
    "        try:\n",
    "            exec(solution)\n",
    "            solution_output = dsl_code()\n",
    "            solution_found = True\n",
    "        except Exception as e:\n",
    "            solution_output = \"ERROR: \" + str(e) + \"\\n\" + \"ERROR TRACE: \\n\" + traceback.format_exc()\n",
    "            messages.append({\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\": solution_output\n",
    "                    }\n",
    "                ]\n",
    "            })\n",
    "            print(\"Error in solution execution:\")\n",
    "            print(solution_output)\n",
    "            continue\n",
    "        print(\"Solution output:\")\n",
    "        print(solution_output)\n",
    "        break\n",
    "\n",
    "    if \"<DSL>\" in response or \"<VALIDATE_SOLUTION>\" in response:\n",
    "        response = response.replace(\"VALIDATE_SOLUTION\" , \"DSL\")\n",
    "        code = response.split(\"<DSL>\")[1].split(\"</DSL>\")[0]\n",
    "        \n",
    "        def dsl_code():\n",
    "            return \"-- dsl_code() function not defined --\"\n",
    "        \n",
    "        try:\n",
    "            exec(code)\n",
    "            code_output = dsl_code()\n",
    "            if(code_output == \"-- dsl_code() function not defined --\"):\n",
    "                code_output = \"You have not defined the dsl_code() function.\"\n",
    "            if(code_output == \"\" or code_output == None):\n",
    "                code_output = \"No output from your dsl_code() function.\"\n",
    "            if(len(code_output) > 3000):\n",
    "                code_output = code_output[:3000] + \"... (truncated output)\"\n",
    "            if(i == ITERATIONS-1):\n",
    "                code_output += \"\\n\\n(YOU HAVE RUN OUT OF TIME. THIS IS YOUR FINAL THROUGHT PROCESS PLEASE SUBMIT A SOLUTION, based on your thought process. The code must work first try on all train and test data without any modifications.)\"\n",
    "        except Exception as e:\n",
    "            code_output = \"ERROR: \" + str(e) + \"\\n\" + \"ERROR TRACE: \\n\" + traceback.format_exc()\n",
    "        \n",
    "        print(\"Code output:\")\n",
    "        print(code_output)\n",
    "\n",
    "        messages.append({\n",
    "            \"role\": \"user\",\n",
    "            \"content\": [\n",
    "                {\n",
    "                    \"type\": \"text\",\n",
    "                    \"text\": f\"<CODE_OUTPUT>\\n{code_output}\\n</CODE_OUTPUT>\"\n",
    "                }\n",
    "            ]\n",
    "        })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Solution found in 2 iterations.\n",
      "⚫🔴🔴🔴\n",
      "🔴🔴⚫🔴\n",
      "🔴🔴⚫🔴\n",
      "⚫🔴🔴🔴\n"
     ]
    }
   ],
   "source": [
    "if(solution_found):\n",
    "    print(f\"Solution found in {i+1} iterations.\")\n",
    "    print(grid_to_emoji(solution_output))\n",
    "else:\n",
    "    print(\"Solution not found.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⚫⚫🔵⚫🟢⚫🔹⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🟢⚫🔹⚫⚫⚫⚫⚫⚫⚫⚫\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",
      "⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫⚫🟢⚫⚫⚫⚫🟢⚫⚫⚫⚫⚫⚫⚫⚫\n",
      "⚫⚫🔹⚫⚫⚫⚫🔹⚫⚫⚫⚫⚫⚫⚫⚫⚫🟢⚫⚫⚫🟢⚫⚫⚫⚫⚫⚫⚫🟢\n"
     ]
    }
   ],
   "source": [
    "print(\"INPUT:\")\n",
    "print(grid_to_emoji(data[\"test\"][0][\"input\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}