create_puzzles.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "#np.random.seed(42)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "task_name = \"vehicles\"\n",
    "\n",
    "# Add images here, i.e. pictures of vehicles with \"vehicles\": [\"car\", \"truck\", \"motorcycle\", \"jeep\", \"tractor\"]\n",
    "objects = {\"plants\": [\"apple\", \"pineapple\", \"banana\", \"avocado\", \"ginger\", \"kiwi\", \"pear\", \"watermelon\"]}\n",
    "all_objects = objects[task_name]\n",
    "\n",
    "background_colors = [\"darkkhaki\", \"navajowhite\", \"gold\"]\n",
    "lines = []\n",
    "file_num = 0\n",
    "for _ in range(5):\n",
    "    objects = list(np.random.choice(all_objects, np.random.randint(4,len(all_objects)+1)))\n",
    "    rows = 2*np.random.randint(2,5)\n",
    "    min_cols = int(4*len(objects)/(rows-2)+2)+1\n",
    "    cols = 4*np.random.randint(min_cols//2, min_cols//2+4)\n",
    "    assert(cols % 2 + rows % 2 == 0)\n",
    "    fig, ax = plt.subplots(figsize=(cols/2, rows/2))\n",
    "\n",
    "    correct = set(np.random.choice(objects, np.random.randint(1,len(objects)))) # No trick questions, there is always at least one object visible.\n",
    "    incorrect = set(objects) - correct\n",
    "\n",
    "    print(f\"Correct: {correct}\")\n",
    "\n",
    "    spaces = (cols/2 - 2) * (rows - 2) / 4\n",
    "    try:\n",
    "        assert spaces - len(objects) > 0, \"There is not enough space to fit all the objects\"\n",
    "    except Exception as e:\n",
    "        print(f\"Skipping due to error: {e}\")\n",
    "    optional_holes = int(max(0, min(spaces - len(objects), 4)))\n",
    "    if optional_holes > 0:\n",
    "        optional_holes = np.random.randint(0,optional_holes)\n",
    "\n",
    "    def rand_position():\n",
    "        return np.random.randint(0,1+cols/4-2), np.random.randint(0,1+rows/2-2)\n",
    "    \n",
    "    object_positions = {}\n",
    "    try: # Terrible and hacky, refactor immediately\n",
    "        for object in objects:\n",
    "            position = rand_position()\n",
    "            counter = 0\n",
    "            while position in object_positions.values():\n",
    "                position = rand_position()\n",
    "                counter += 1\n",
    "                if counter > 100:\n",
    "                    raise Exception\n",
    "            object_positions[object] = position\n",
    "    except:\n",
    "        continue\n",
    "\n",
    "    hole_positions = []\n",
    "    try: # Terrible and hacky, refactor immediately\n",
    "        for k in range(optional_holes):\n",
    "            position = rand_position()\n",
    "            counter = 0\n",
    "            while position in list(object_positions.values()) + hole_positions:\n",
    "                position = rand_position()\n",
    "                counter += 1\n",
    "                if counter > 100:\n",
    "                    raise Exception\n",
    "            hole_positions.append(position)\n",
    "    except:\n",
    "        continue\n",
    "    assert(len(object_positions.values()) == len(set(object_positions.values()))) # No duplicate positions\n",
    "\n",
    "    # Define the grid size and colors\n",
    "    grid_color = np.random.choice(background_colors, 1)[0]\n",
    "    border_color = 'black'\n",
    "\n",
    "    def add_small_patch(i,j):\n",
    "        rect = patches.Rectangle((i, j), 1, 1, linewidth=1, edgecolor=border_color, facecolor=grid_color)\n",
    "        ax.add_patch(rect)\n",
    "\n",
    "    \n",
    "    # Tile edge cells\n",
    "    for i in range(cols):\n",
    "        add_small_patch(i,0)\n",
    "        add_small_patch(i,rows-1)\n",
    "    for j in range(rows):\n",
    "        add_small_patch(0,j)\n",
    "        add_small_patch(cols-1,j)\n",
    "        add_small_patch(cols//2-1,j)\n",
    "        add_small_patch(cols//2,j)\n",
    "    \n",
    "    # Tile middle cells\n",
    "    for r in range((rows-2)//2):\n",
    "        for c in range((cols//2-2)//2):\n",
    "            col = int(1 + cols//2 + 2*c)\n",
    "            row = 1 + 2*r\n",
    "            \n",
    "            for i in [0,1]:\n",
    "                for j in [0,1]:\n",
    "                    add_small_patch(col+i,row+j)\n",
    "            \n",
    "            skip = False\n",
    "            for object, position in object_positions.items():\n",
    "                if position == (c,r) and object in correct:\n",
    "                    skip = True\n",
    "            if skip:\n",
    "                continue\n",
    "\n",
    "            if (c, r) in hole_positions:\n",
    "                continue\n",
    "\n",
    "            col = int(cols//2 - 3 - 2*c)\n",
    "            row = 1 + 2*r\n",
    "\n",
    "            for i in [0,1]:\n",
    "                for j in [0,1]:\n",
    "                    add_small_patch(col+i,row+j)      \n",
    "\n",
    "    thick_linewidth = 8\n",
    "    rect = patches.Rectangle((0, 0), cols, rows, linewidth=thick_linewidth, edgecolor=border_color, facecolor=\"#00000000\")\n",
    "    ax.add_patch(rect)\n",
    "    rect = patches.Rectangle((0, 0), cols/2, rows, linewidth=thick_linewidth / 2, edgecolor=border_color, facecolor=\"#00000000\")\n",
    "    ax.add_patch(rect)\n",
    "\n",
    "    for object, position in object_positions.items():\n",
    "        col, row = position\n",
    "        col = 1 +cols/2 + 2*col\n",
    "        row = 1 + 2*row\n",
    "        rect = patches.Rectangle((col, row), 2, 2, linewidth=1, edgecolor=border_color, facecolor=grid_color)\n",
    "        ax.add_patch(rect)\n",
    "\n",
    "        img = mpimg.imread(f'assets/{object}.png')\n",
    "        ax.imshow(img, extent=[col, col+2, row, row+2], aspect='auto', zorder=2)\n",
    "\n",
    "    ax.set_xlim(0, cols)\n",
    "    ax.set_ylim(0, rows)\n",
    "    ax.axis('off')\n",
    "\n",
    "    plt.gca().set_aspect('equal')\n",
    "    filename = f'puzzles/{task_name}-{file_num}.png'\n",
    "    plt.savefig(filename, bbox_inches=\"tight\")\n",
    "    puzzle = {\"id\": f\"{task_name}-{file_num}\", \"correct\": list(correct), \"incorrect\": list(incorrect), \"image\": filename}\n",
    "    lines.append(json.dumps(puzzle) + \"\\n\")\n",
    "    file_num += 1\n",
    "\n",
    "with open(f\"puzzles/{task_name}.jsonl\", \"w\") as fp:\n",
    "    fp.writelines(lines)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}