prepare_more_independent_data.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare `more_independent` data\n",
    "\n",
    "The data recorded while playing the game needs to be checked and cleaned before being used to train a model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from recording import FileBasedDatabase"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "recordings has 688 files.\n",
      "data\\more_independent has 0 files.\n",
      "data\\more_independent\\train has 15475 files.\n",
      "data\\more_independent\\test has 3868 files.\n"
     ]
    }
   ],
   "source": [
    "db_rec = FileBasedDatabase(\"recordings\")\n",
    "db = FileBasedDatabase(os.path.join(\"data\", \"more_independent\"))\n",
    "db_train = FileBasedDatabase(os.path.join(\"data\", \"more_independent\", \"train\"))\n",
    "db_test = FileBasedDatabase(os.path.join(\"data\", \"more_independent\", \"test\"))\n",
    "print(f\"{db_rec._path} has {len(db_rec)} files.\")\n",
    "print(f\"{db._path} has {len(db)} files.\")\n",
    "print(f\"{db_train._path} has {len(db_train)} files.\")\n",
    "print(f\"{db_test._path} has {len(db_test)} files.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Moved data from recordings to data\\more_independent\n"
     ]
    }
   ],
   "source": [
    "# Move files from recordings folder to our dataset\n",
    "for boards, events in db_rec:\n",
    "    db.insert(boards, events)\n",
    "\n",
    "print(f\"Moved data from {db_rec._path} to {db._path}\")\n",
    "\n",
    "db_rec.delete_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "boards: shape (2, 22, 10), dtype int32\n",
      "events: shape (1,), dtype int32\n"
     ]
    }
   ],
   "source": [
    "# Check we can read the files\n",
    "\n",
    "boards, events = next(iter(db))\n",
    "print(f\"boards: shape {boards.shape}, dtype {boards.dtype}\")\n",
    "print(f\"events: shape {events.shape}, dtype {events.dtype}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcoAAAG1CAYAAAB56LTIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8pXeV/AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAWMElEQVR4nO3de4xU9dnA8WegsCwX78gCAoKAWPDSVKBag4oaq1bjZavYKoJYUqtWTSxWTURt7B/VlLXYV6SNiLZaLiJpGxBti6aKgvVCta1RLFgKBrCCrAUXdc/7R1/ndV14WBTZXf18kk2Y35yZ+c3scL5z5szOKRVFUQQAsFVtmnsCANCSCSUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEkhbp7rvvjlKptM2fp556qrmnGIsWLYobbrghNmzYsMtv+4YbbmjweHTs2DF69+4dp556akybNi3q6up2+Zzgs+oLzT0ByNx0003Rt2/fRuP9+/dvhtk0tGjRorjxxhtjzJgxscceezTLHO64447o3Llz1NXVxapVq2LBggVx4YUXRk1NTfzud7+LXr16Ncu84LNEKGnRTjrppDj88MObexotVnV1deyzzz7l09dff3386le/itGjR8c3vvGN7W55b9q0KTp27PhpTxNaNW+90mq9++67sddee8XYsWMbnbdx48bo0KFDXHXVVeWxurq6mDhxYvTv3z8qKiqiV69eMWHChEZvU5ZKpbj00ktj7ty5MWTIkKioqIjBgwfHQw89VF7mhhtuiO9///sREdG3b9/yW6ArVqz4dO7sDvjWt74VF110USxevDgeeeSR8vgxxxwTQ4YMiWeeeSZGjBgRHTt2jGuvvTYiItauXRvjxo2Lbt26RYcOHeLQQw+N6dOnN7jeFStWRKlUiltvvTUmTZoUffr0icrKyjj66KPjxRdf3KX3EXYlW5S0aG+99Va88cYbDcZKpVLsvffe0a5duzjjjDNizpw5ceedd0b79u3Ly8ydOzfq6upi1KhRERFRX18fp512Wjz++OMxfvz4OOigg+KFF16ISZMmxcsvvxxz585tcBuPP/54zJkzJ7773e9Gly5d4qc//WmcddZZ8c9//jP23nvvOPPMM+Pll1+O+++/PyZNmlTequvates278umTZti06ZN273Pbdu2jT333LOpD9FWnX/++TF16tR4+OGH44QTTiiP//vf/46TTjopRo0aFeedd15069YtNm/eHMccc0wsW7YsLr300ujbt2/MmjUrxowZExs2bIjLL7+8wXXfc889UVtbG5dcckm88847cdttt8XIkSPjhRdeiG7dun2ieUOLVEALNG3atCIitvpTUVFRXm7BggVFRBS//e1vG1z+5JNPLvr161c+fe+99xZt2rQp/vSnPzVYbsqUKUVEFE888UR5LCKK9u3bF8uWLSuPLV26tIiIYvLkyeWxW265pYiIYvny5U26TxMnTtzmffrwT58+fZp8XevWrdvq+evXry8iojjjjDPKY0cffXQREcWUKVMaLFtTU1NERPHLX/6yPLZly5biiCOOKDp37lxs3LixKIqiWL58eRERRWVlZfGvf/2rvOzixYuLiCiuvPLKJj0O0NrYoqRF+9nPfhYDBw5sMNa2bdvyv0eOHBn77LNPzJgxI77+9a9HRMT69evjkUceafC266xZs+Kggw6KQYMGNdhCHTlyZERELFy4MI488sjy+PHHHx8HHHBA+fQhhxwSu+22W/zjH//42Pdl9OjRcdRRR213ucrKyo99Gx/o3LlzRETU1tY2GK+oqGj0VvW8efOiqqoqzj333PJYu3bt4nvf+16ce+658dhjj5Uf24iI008/PXr27Fk+PWzYsBg+fHjMmzcvfvKTn3ziuUNLI5S0aMOGDUs/zPOFL3whzjrrrLjvvvuirq4uKioqYs6cOfHuu+/GOeecU17ulVdeib///e/bfGt07dq1DU737t270TJ77rlnrF+//mPek4h+/fpFv379Pvbld8Tbb78dERFdunRpMN6zZ88Gb1FHRLz22msxYMCAaNOm4UcWDjrooPL5HzZgwIBGtzdw4MCYOXPmJ543tERCSas3atSouPPOO2P+/Plx+umnx8yZM2PQoEFx6KGHlpepr6+Pgw8+eJtbPB/9M4oPb7V+WFEUH3ueb7/9djlgmbZt26b7Opvigw/XfPTPaHbG1ip83gglrd6IESOie/fuMWPGjDjqqKPij3/8Y1x33XUNljnggANi6dKlcdxxx0WpVNopt7uj13PrrbfGjTfeuN3l+vTp84k/PXvvvfdGRMSJJ57YpNv7y1/+EvX19Q22Kl966aXy+R/2yiuvNLqOl19+Ofbff/9PMGNouYSSVq9NmzZRXV0dd911VwwbNizee++9Bm+7RkScffbZMW/evPj5z38e48ePb3De5s2bo76+Pjp16rRDt/vB8k39Zp5dtY/yvvvui1/84hdxxBFHxHHHHbfd5U8++eR4+OGHY8aMGeX9lO+9915Mnjw5OnfuHEcffXSD5efOnRurVq0q76dcsmRJLF68OK644opPNG9oqYSSFm3+/PnlLZsPO/LIIxvs7zvnnHNi8uTJMXHixDj44IPL+9c+cP7558fMmTPjO9/5TixcuDC++tWvxvvvvx8vvfRSzJw5MxYsWLDDX2zw5S9/OSIirrvuuhg1alS0a9cuTj311G0G99PYRzl79uzo3LlzbNmypfzNPE888UQceuihMWvWrCZdx/jx4+POO++MMWPGxDPPPBP7779/zJ49O5544omoqalptJ+zf//+cdRRR8XFF18cdXV1UVNTE3vvvXdMmDBhp943aCmEkhbt+uuv3+r4tGnTGkTnyCOPjF69esXKlSsbbU1G/Herc+7cuTFp0qS455574sEHH4yOHTtGv3794vLLL2/0ydqmGDp0aPzwhz+MKVOmxEMPPRT19fWxfPnyHd4y/SQuvvjiiIjo0KFD7LPPPnHYYYfFXXfdFd/85jejoqKiSddRWVkZjz76aPzgBz+I6dOnx8aNG+PAAw+MadOmxZgxYxotP3r06GjTpk3U1NTE2rVrY9iwYXH77bdH9+7dd+ZdgxajVHySTycAnxsrVqyIvn37xi233NLgT2/gs85X2AFAQigBICGUAJCwjxIAErYoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKLfihRdeiOrq6ujTp0906NAhevbsGSeccEJMnjy5wXI/+tGPYu7cuc0zyf9TV1cXV199dfTo0SMqKytj+PDh8cgjjzTrnODzqrWsO95+++2YOHFifO1rX4u99torSqVS3H333c02n5bON/N8xKJFi+LYY4+N3r17xwUXXBBVVVWxcuXKeOqpp+LVV1+NZcuWlZft3LlzVFdXN+sT7Nxzz43Zs2fHFVdcEQMGDIi77747nn766Vi4cGGTDhIM7Bytad3xwZFgevfuHf369YtHH310m4dVw/EoG7n55ptj9913j6effjr22GOPBuetXbu2eSa1DUuWLIlf//rXDQ57NHr06BgyZEhMmDAhFi1a1MwzhM+P1rTu6N69e7z++utRVVUVf/7zn2Po0KHNPaUWzVuvH/Hqq6/G4MGDGz3RIyL23Xff8r9LpVL85z//ienTp0epVIpSqdTg1diqVaviwgsvjG7dukVFRUUMHjw47rrrrgbX9+ijj0apVIoZM2bEtddeG1VVVdGpU6c47bTTYuXKldud6+zZs6Nt27Yxfvz48liHDh1i3Lhx8eSTTzbpOoCdozWtOyoqKqKqqupj39fPG1uUH9GnT5948skn48UXX4whQ4Zsc7l77703Lrroohg2bFg5VAcccEBERKxZsya+8pWvRKlUiksvvTS6du0a8+fPj3HjxsXGjRvjiiuuaHBdN998c5RKpbj66qtj7dq1UVNTE8cff3w8//zzUVlZuc05PPfcczFw4MDYbbfdGowPGzYsIiKef/756NWr18d5GIAd1JrWHeygggYefvjhom3btkXbtm2LI444opgwYUKxYMGCYsuWLY2W7dSpU3HBBRc0Gh83blzRvXv34o033mgwPmrUqGL33XcvNm3aVBRFUSxcuLCIiKJnz57Fxo0by8vNnDmziIjitttuS+c6ePDgYuTIkY3G//rXvxYRUUyZMqUpdxnYCVrTuuPDnn766SIiimnTpjX5Mp833nr9iBNOOCGefPLJOO2002Lp0qXx4x//OE488cTo2bNn/OY3v9nu5YuiiAceeCBOPfXUKIoi3njjjfLPiSeeGG+99VY8++yzDS4zevTo6NKlS/l0dXV1dO/ePebNm5fe1ubNm6OioqLReIcOHcrnA7tGa1p3sGOEciuGDh0ac+bMifXr18eSJUvimmuuidra2qiuro6//e1v6WXXrVsXGzZsiKlTp0bXrl0b/IwdOzYiGu/YHzBgQIPTpVIp+vfvHytWrEhvq7KyMurq6hqNv/POO+XzgV2ntaw72DH2USbat28fQ4cOjaFDh8bAgQNj7NixMWvWrJg4ceI2L1NfXx8REeedd15ccMEFW13mkEMO2Snz6969e6xatarR+Ouvvx4RET169NgptwPsmJa+7mDHCGUTHX744RHx/xGK+O+rt4/q2rVrdOnSJd5///04/vjjm3Tdr7zySoPTRVHEsmXLtvuf4rDDDouFCxfGxo0bG3ygZ/HixeXzgebVEtcd7BhvvX7EwoULo9jKdzB88J7/gQceWB7r1KlTbNiwocFybdu2jbPOOiseeOCBePHFFxtdz7p16xqN3XPPPVFbW1s+PXv27Hj99dfjpJNOSudaXV0d77//fkydOrU8VldXF9OmTYvhw4f7xCvsQq1p3cGO8c08HzFkyJDYtGlTnHHGGTFo0KDYsmVLLFq0KGbMmBG9evWK5557rvx3Uqeccko89thjcdNNN0WPHj2ib9++MXz48FizZk0MHz481q1bF9/+9rfji1/8Yrz55pvx7LPPxu9///t48803I+K/fwt17LHHxsEHHxylUinGjh0ba9asiZqamthvv/1i6dKl0bFjx3S+Z599djz44INx5ZVXRv/+/WP69OmxZMmS+MMf/hAjRoz4tB8u4P+0tnXH7bffHhs2bIjVq1fHHXfcEWeeeWZ86UtfioiIyy67LHbfffdP9fFqVZrr47Yt1fz584sLL7ywGDRoUNG5c+eiffv2Rf/+/YvLLrusWLNmTYNlX3rppWLEiBFFZWVlERENPu69Zs2a4pJLLil69epVtGvXrqiqqiqOO+64YurUqeVlPviI9/33319cc801xb777ltUVlYWp5xySvHaa681ab6bN28urrrqqqKqqqqoqKgohg4dWjz00EM75bEAmq61rTv69OlTRMRWf5YvX74zHpLPDFuUzeiDV4WzZs2K6urq5p4O0EpYd+xa9lECQEIoASAhlACQsI8SABK2KAEg0aRv5qmvr4/Vq1dHly5dtvqNEtCciqKI2tra6NGjR7Rp47VfS2LdQUvW1HVHk0K5evVq3/JCi7dy5crYb7/9mnsafIh1B63B9tYdTXr5/eHDuEBL5Xna8vid0Bps73napFB6y4TWwPO05fE7oTXY3vPUDh0ASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAokkHbqax/5n+cnNPIb57wcDmngLAZ54tSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkHDgZoBdpCUc8D3CQd93lC1KAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASDheJQfk+O5AXw+2KIEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJB24G2EUc8L11skUJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQaFIoi6L4tOcBn5jnacvjd0JrsL3naZNCWVtbu1MmA58mz9OWx++E1mB7z9NS0YSXfPX19bF69ero0qVLlEqlnTY52BmKooja2tro0aNHtGljb0JLYt1BS9bUdUeTQgkAn1defgNAQigBICGUAJAQSgBICCUAJIQSABJCCQCJ/wX/xVlbbMmyyAAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from tetris import CELL_COLORS\n",
    "from engines import EVENT_NAMES\n",
    "\n",
    "\n",
    "def render_boards(boards):\n",
    "    batch_size, height, width = boards.shape\n",
    "    imgs = np.zeros((batch_size, height, width, 3), dtype=np.int32)\n",
    "    for idx in range(batch_size):\n",
    "        for row in range(height):\n",
    "            for col in range(width):\n",
    "                imgs[idx, row, col, :] = CELL_COLORS[boards[idx, row, col]]\n",
    "    return imgs\n",
    "\n",
    "\n",
    "\n",
    "def show_example(boards, events):\n",
    "    imgs = render_boards(boards)\n",
    "    event_name = EVENT_NAMES[events.item()]\n",
    "\n",
    "    fig, axs = plt.subplots(1, len(boards))\n",
    "    fig.suptitle(f\"Event = {event_name}\")\n",
    "\n",
    "    for i, ax in enumerate(axs):\n",
    "        ax.xaxis.set_visible(False)\n",
    "        ax.yaxis.set_visible(False)\n",
    "        ax.set_title(f\"Step {i}\")\n",
    "        ax.imshow(imgs[i])\n",
    "\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "show_example(boards, events)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Showing example 13\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAcoAAAG1CAYAAAB56LTIAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8pXeV/AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAW8ElEQVR4nO3de4xU9d348c+wwrJcxBtyExAExIKXPhWo1qCixqrVeNkqtoogltSqVROLVRNRG588qaZgsY9IGxFttVxE0jYg2hZNFQXrhWpbo1iwCAawgosFF2XP748+zs8V+LAIsrv6eiWbOGfOnPnOOHzfe+bMzikVRVEEALBVLRp7AADQlAklACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhJIm6d57741SqbTNn2eeeaaxhxgLFiyIm266KdatW7fb7/umm26q93y0adMmevToEaeffnpMmTIlamtrd/uY4PNqj8YeAGRuueWW6NWr1xbL+/Tp0wijqW/BggVx8803x8iRI2OvvfZqlDHcdddd0a5du6itrY0VK1bEvHnz4uKLL44JEybE7373u+jevXujjAs+T4SSJu2UU06JI488srGH0WRVV1fHfvvtV7584403xq9+9asYMWJEfPOb39zunveGDRuiTZs2n/UwoVnz1ivN1gcffBD77LNPjBo1aovrampqonXr1nHNNdeUl9XW1sa4ceOiT58+UVlZGd27d4+xY8du8TZlqVSKyy+/PGbPnh0DBw6MysrKGDBgQDzyyCPldW666ab4wQ9+EBERvXr1Kr8FumzZss/mwe6Ab3/723HJJZfEwoUL47HHHisvP+6442LgwIHx3HPPxdChQ6NNmzZx/fXXR0TE6tWrY/To0dGpU6do3bp1HH744TF16tR62122bFmUSqW4/fbbY/z48dGzZ8+oqqqKY489Nl5++eXd+hhhd7JHSZP27rvvxttvv11vWalUin333TdatmwZZ511VsyaNSvuvvvuaNWqVXmd2bNnR21tbQwfPjwiIurq6uKMM86IJ598MsaMGROHHHJIvPTSSzF+/Ph49dVXY/bs2fXu48knn4xZs2bF9773vWjfvn389Kc/jXPOOSf++c9/xr777htnn312vPrqq/Hggw/G+PHjy3t1HTt23OZj2bBhQ2zYsGG7j7mioiL23nvvhj5FW3XhhRfG5MmT49FHH42TTjqpvPxf//pXnHLKKTF8+PC44IILolOnTrFx48Y47rjjYsmSJXH55ZdHr169YsaMGTFy5MhYt25dXHnllfW2fd9998X69evjsssui/fffz/uuOOOGDZsWLz00kvRqVOnnRo3NEkFNEFTpkwpImKrP5WVleX15s2bV0RE8dvf/rbe7U899dSid+/e5cv3339/0aJFi+JPf/pTvfUmTZpURETx1FNPlZdFRNGqVatiyZIl5WWLFy8uIqKYOHFiedltt91WRESxdOnSBj2mcePGbfMxffynZ8+eDd7WmjVrtnr92rVri4gozjrrrPKyY489toiIYtKkSfXWnTBhQhERxS9/+cvysk2bNhVHHXVU0a5du6KmpqYoiqJYunRpERFFVVVV8eabb5bXXbhwYRERxdVXX92g5wGaG3uUNGk/+9nPol+/fvWWVVRUlP972LBhsd9++8W0adPiG9/4RkRErF27Nh577LF6b7vOmDEjDjnkkOjfv3+9PdRhw4ZFRMT8+fPj6KOPLi8/8cQT46CDDipfPuyww2LPPfeMf/zjH5/6sYwYMSKOOeaY7a5XVVX1qe/jI+3atYuIiPXr19dbXllZucVb1XPmzInOnTvH+eefX17WsmXL+P73vx/nn39+PPHEE+XnNiLizDPPjG7dupUvDx48OIYMGRJz5syJn/zkJzs9dmhqhJImbfDgwemHefbYY48455xz4oEHHoja2tqorKyMWbNmxQcffBDnnXdeeb3XXnst/v73v2/zrdHVq1fXu9yjR48t1tl7771j7dq1n/KRRPTu3Tt69+79qW+/I957772IiGjfvn295d26dav3FnVExBtvvBF9+/aNFi3qf2ThkEMOKV//cX379t3i/vr16xfTp0/f6XFDUySUNHvDhw+Pu+++O+bOnRtnnnlmTJ8+Pfr37x+HH354eZ26uro49NBDt7nH88k/o/j4XuvHFUXxqcf53nvvlQOWqaioSI91NsRHH6755J/R7Iq9VfiiEUqavaFDh0aXLl1i2rRpccwxx8Qf//jHuOGGG+qtc9BBB8XixYvjhBNOiFKptEvud0e3c/vtt8fNN9+83fV69uy505+evf/++yMi4uSTT27Q/f3lL3+Jurq6enuVr7zySvn6j3vttde22Marr74aBx544E6MGJouoaTZa9GiRVRXV8c999wTgwcPjg8//LDe264REeeee27MmTMnfv7zn8eYMWPqXbdx48aoq6uLtm3b7tD9frR+Q7+ZZ3cdo3zggQfiF7/4RRx11FFxwgknbHf9U089NR599NGYNm1a+Tjlhx9+GBMnTox27drFscceW2/92bNnx4oVK8rHKRctWhQLFy6Mq666aqfGDU2VUNKkzZ07t7xn83FHH310veN95513XkycODHGjRsXhx56aPn42kcuvPDCmD59enz3u9+N+fPnx9e+9rXYvHlzvPLKKzF9+vSYN2/eDn+xwVe+8pWIiLjhhhti+PDh0bJlyzj99NO3GdzP4hjlzJkzo127drFp06byN/M89dRTcfjhh8eMGTMatI0xY8bE3XffHSNHjoznnnsuDjzwwJg5c2Y89dRTMWHChC2Oc/bp0yeOOeaYuPTSS6O2tjYmTJgQ++67b4wdO3aXPjZoKoSSJu3GG2/c6vIpU6bUi87RRx8d3bt3j+XLl2+xNxnxn73O2bNnx/jx4+O+++6Lhx9+ONq0aRO9e/eOK6+8cotP1jbEoEGD4kc/+lFMmjQpHnnkkairq4ulS5fu8J7pzrj00ksjIqJ169ax3377xRFHHBH33HNPfOtb34rKysoGbaOqqioef/zx+OEPfxhTp06NmpqaOPjgg2PKlCkxcuTILdYfMWJEtGjRIiZMmBCrV6+OwYMHx5133hldunTZlQ8NmoxSsTOfTgC+MJYtWxa9evWK2267rd6f3sDnna+wA4CEUAJAQigBIOEYJQAk7FECQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQbsVLL70U1dXV0bNnz2jdunV069YtTjrppJg4cWK99f77v/87Zs+e3TiD/D+1tbVx7bXXRteuXaOqqiqGDBkSjz32WKOOCb6omsvc8d5778W4cePi61//euyzzz5RKpXi3nvvbbTxNHW+mecTFixYEMcff3z06NEjLrrooujcuXMsX748nnnmmXj99ddjyZIl5XXbtWsX1dXVjfoCO//882PmzJlx1VVXRd++fePee++NZ599NubPn9+gkwQDu0Zzmjs+OhNMjx49onfv3vH4449v87RqOB/lFm699dbo0KFDPPvss7HXXnvVu2716tWNM6htWLRoUfz617+ud9qjESNGxMCBA2Ps2LGxYMGCRh4hfHE0p7mjS5cu8dZbb0Xnzp3jz3/+cwwaNKixh9Skeev1E15//fUYMGDAFi/0iIj999+//N+lUin+/e9/x9SpU6NUKkWpVKr329iKFSvi4osvjk6dOkVlZWUMGDAg7rnnnnrbe/zxx6NUKsW0adPi+uuvj86dO0fbtm3jjDPOiOXLl293rDNnzoyKiooYM2ZMeVnr1q1j9OjR8fTTTzdoG8Cu0ZzmjsrKyujcufOnfqxfNPYoP6Fnz57x9NNPx8svvxwDBw7c5nr3339/XHLJJTF48OByqA466KCIiFi1alV89atfjVKpFJdffnl07Ngx5s6dG6NHj46ampq46qqr6m3r1ltvjVKpFNdee22sXr06JkyYECeeeGK8+OKLUVVVtc0xvPDCC9GvX7/Yc8896y0fPHhwRES8+OKL0b1790/zNAA7qDnNHeyggnoeffTRoqKioqioqCiOOuqoYuzYscW8efOKTZs2bbFu27Zti4suumiL5aNHjy66dOlSvP322/WWDx8+vOjQoUOxYcOGoiiKYv78+UVEFN26dStqamrK602fPr2IiOKOO+5IxzpgwIBi2LBhWyz/61//WkREMWnSpIY8ZGAXaE5zx8c9++yzRUQUU6ZMafBtvmi89foJJ510Ujz99NNxxhlnxOLFi+PHP/5xnHzyydGtW7f4zW9+s93bF0URDz30UJx++ulRFEW8/fbb5Z+TTz453n333Xj++efr3WbEiBHRvn378uXq6uro0qVLzJkzJ72vjRs3RmVl5RbLW7duXb4e2D2a09zBjhHKrRg0aFDMmjUr1q5dG4sWLYrrrrsu1q9fH9XV1fG3v/0tve2aNWti3bp1MXny5OjYsWO9n1GjRkXElgf2+/btW+9yqVSKPn36xLJly9L7qqqqitra2i2Wv//+++Xrgd2nucwd7BjHKBOtWrWKQYMGxaBBg6Jfv34xatSomDFjRowbN26bt6mrq4uIiAsuuCAuuuiira5z2GGH7ZLxdenSJVasWLHF8rfeeisiIrp27bpL7gfYMU197mDHCGUDHXnkkRHx/yMU8Z/f3j6pY8eO0b59+9i8eXOceOKJDdr2a6+9Vu9yURSxZMmS7f6jOOKII2L+/PlRU1NT7wM9CxcuLF8PNK6mOHewY7z1+gnz58+PYivfwfDRe/4HH3xweVnbtm1j3bp19darqKiIc845Jx566KF4+eWXt9jOmjVrtlh23333xfr168uXZ86cGW+99Vaccsop6Virq6tj8+bNMXny5PKy2tramDJlSgwZMsQnXmE3ak5zBzvGN/N8wsCBA2PDhg1x1llnRf/+/WPTpk2xYMGCmDZtWnTv3j1eeOGF8t9JnXbaafHEE0/ELbfcEl27do1evXrFkCFDYtWqVTFkyJBYs2ZNfOc734kvfelL8c4778Tzzz8fv//97+Odd96JiP/8LdTxxx8fhx56aJRKpRg1alSsWrUqJkyYEAcccEAsXrw42rRpk4733HPPjYcffjiuvvrq6NOnT0ydOjUWLVoUf/jDH2Lo0KGf9dMF/J/mNnfceeedsW7duli5cmXcddddcfbZZ8eXv/zliIi44oorokOHDp/p89WsNNbHbZuquXPnFhdffHHRv3//ol27dkWrVq2KPn36FFdccUWxatWqeuu+8sorxdChQ4uqqqoiIup93HvVqlXFZZddVnTv3r1o2bJl0blz5+KEE04oJk+eXF7no494P/jgg8V1111X7L///kVVVVVx2mmnFW+88UaDxrtx48bimmuuKTp37lxUVlYWgwYNKh555JFd8lwADdfc5o6ePXsWEbHVn6VLl+6Kp+Rzwx5lI/rot8IZM2ZEdXV1Yw8HaCbMHbuXY5QAkBBKAEgIJQAkHKMEgIQ9SgBINOibeerq6mLlypXRvn37rX6jBDSmoihi/fr10bVr12jRwu9+TYm5g6asoXNHg0K5cuVK3/JCk7d8+fI44IADGnsYfIy5g+Zge3NHg379/vhpXKCp8jptevw/oTnY3uu0QaH0lgnNgddp0+P/Cc3B9l6nDugAQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAxB6NPQAa1w8v/3djDyEiIv7nzraNPQSArbJHCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkSkVRFNtbqaamJjp06LA7xkMz1NjntKzdVBPjJ3eJd999N/bcc89GHQv1mTs+vxr7333ErjuP7fbmDnuUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEg4cTNfG44cXPTY+5gW5rCiZ8betJ3e5QAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASCxR2MPAIAvnv+5s21jD6HB7FECQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASCxR2MPoDEUt/9XYw8hStc839hDAHZAU5g3IswdjcEeJQAkhBIAEkIJAAmhBICEUAJAQigBICGUAJAQSgBICCUAJIQSABJCCQAJoQSAhFACQEIoASAhlACQaHbno/yvHx65C7ZSt9NbuKvjr3fq9v87daeHEN+7qN/ObwS+IHZ+7mj8eSPC3NEY7FECQEIoASAhlACQEEoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBIDEDp24+fCrvxwVlRWf1ViAzylzB82ZPUoASAglACSEEgASQgkACaEEgIRQAkBCKAEgIZQAkBBKAEgIJQAkhBIAEkIJAAmhBICEUAJAQigBINGg81EWRREREZtrN3+mg9ldat6v2+ltbNz43i4YCbvSR69Tmo7P09xh3vj82t7cUSoaMLu8+eab0b179102KPgsLF++PA444IDGHgYfY+6gOdje3NGgUNbV1cXKlSujffv2USqVdukAYWcVRRHr16+Prl27RosWjiY0JeYOmrKGzh0NCiUAfFH59RsAEkIJAAmhBICEUAJAQigBICGUAJAQSgBI/D+tP60X4hYttgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Show a random example from the training set. Run this repeatedly to get an overview of the data.\n",
    "\n",
    "idx = random.randrange(len(db))\n",
    "print(f\"Showing example {idx}\")\n",
    "boards, events = db[idx]\n",
    "show_example(boards, events)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train/test split\n",
    "\n",
    "We split the data into training and test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files split into 551 training and 137 test examples.\n"
     ]
    }
   ],
   "source": [
    "test_split = 0.2\n",
    "test_count = int(test_split * len(db))\n",
    "train_count = len(db) - test_count\n",
    "\n",
    "idxs = list(range(len(db)))\n",
    "random.shuffle(idxs)\n",
    "\n",
    "for idx in range(train_count):\n",
    "    boards, events = db[idxs[idx]]\n",
    "    db_train.insert(boards, events)\n",
    "\n",
    "for idx in range(train_count, len(db)):\n",
    "    boards, events = db[idxs[idx]]\n",
    "    db_test.insert(boards, events)\n",
    "\n",
    "print(f\"Files split into {train_count} training and {test_count} test examples.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up old database\n",
    "db.delete_all()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Post-split checks"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's check the number of each type of block spawn in both the training and test folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "class RecordingDataset(Dataset):\n",
    "    def __init__(self, db: FileBasedDatabase):\n",
    "        self.db = db\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.db)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        boards, events = self.db[idx]\n",
    "        x = torch.tensor(boards[-2]) # Ignore all boards except the last two\n",
    "        e = torch.tensor(events[-1])\n",
    "        y = torch.tensor(boards[-1], dtype=torch.long)\n",
    "        return (x, e), y\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Describing training dataset...\n",
      "{'I': 206, 'O': 207, 'T': 219, 'Z': 235, 'S': 217, 'J': 244, 'L': 207}\n",
      "Dataset has 1535 block spawns.\n",
      "\n",
      "Describing test dataset...\n",
      "{'I': 44, 'O': 49, 'T': 57, 'Z': 51, 'S': 59, 'J': 59, 'L': 61}\n",
      "Dataset has 380 block spawns.\n"
     ]
    }
   ],
   "source": [
    "class BlockType:\n",
    "    I = \"I\"\n",
    "    O = \"O\"\n",
    "    T = \"T\"\n",
    "    Z = \"Z\"\n",
    "    S = \"S\"\n",
    "    J = \"J\"\n",
    "    L = \"L\"\n",
    "    All = [I, O, T, Z, S, J, L]\n",
    "    \n",
    "def get_block_spawn_type(example):\n",
    "    (x, e), y = example\n",
    "    if (x[0, :] > 0).any() | (y[0, :] == 0).all():\n",
    "        return None\n",
    "    if (y[1, :] == 0).all():\n",
    "        return BlockType.I\n",
    "    if y[0, 3] > 0:\n",
    "        if y[1, 3] > 0:\n",
    "            return BlockType.J\n",
    "        else:\n",
    "            if y[0, 5] > 0:\n",
    "                return BlockType.T\n",
    "            else:\n",
    "                return BlockType.Z\n",
    "    else:\n",
    "        if y[1, 3] > 0:\n",
    "            if y[0, 4] > 0:\n",
    "                return BlockType.S\n",
    "            else:\n",
    "                return BlockType.L\n",
    "        else:\n",
    "            return BlockType.O\n",
    "\n",
    "\n",
    "def count_block_spawns_by_type(dataset):\n",
    "    spawns_by_type = {block_type: 0 for block_type in BlockType.All}\n",
    "    for example in dataset:\n",
    "        spawn_type = get_block_spawn_type(example)\n",
    "        if spawn_type is not None:\n",
    "            spawns_by_type[spawn_type] += 1\n",
    "    return spawns_by_type\n",
    "\n",
    "\n",
    "def describe_block_spawns(dataset):\n",
    "    spawns_by_type = count_block_spawns_by_type(dataset)\n",
    "    print(spawns_by_type)\n",
    "    \n",
    "    num_block_spawns = sum(val for key, val in spawns_by_type.items() if key is not None)\n",
    "    frac_block_spawns = num_block_spawns / len(dataset)\n",
    "    print(f\"Dataset has {num_block_spawns} block spawns.\")\n",
    "\n",
    "\n",
    "train_dataset = RecordingDataset(db_train)\n",
    "test_dataset = RecordingDataset(db_test)\n",
    "\n",
    "print(\"Describing training dataset...\")\n",
    "describe_block_spawns(train_dataset)\n",
    "print()\n",
    "print(\"Describing test dataset...\")\n",
    "describe_block_spawns(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also check the number of each type of action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Describing training dataset...\n",
      "Dataset has 7677 events of type Drop\n",
      "Dataset has 2042 events of type Left\n",
      "Dataset has 3251 events of type Right\n",
      "Dataset has 1751 events of type Rotate\n",
      "Dataset has 1305 events of type Insta-drop\n",
      "\n",
      "Describing test dataset...\n",
      "Dataset has 1925 events of type Drop\n",
      "Dataset has 515 events of type Left\n",
      "Dataset has 793 events of type Right\n",
      "Dataset has 430 events of type Rotate\n",
      "Dataset has 342 events of type Insta-drop\n"
     ]
    }
   ],
   "source": [
    "def count_events(dataset):\n",
    "    counts = [0] * len(EVENT_NAMES)\n",
    "    for example in dataset:\n",
    "        (b, e), y = example\n",
    "        counts[e.item()] += 1\n",
    "    return counts\n",
    "\n",
    "\n",
    "def describe_events(dataset):\n",
    "    counts = count_events(dataset)\n",
    "    for i in range(len(EVENT_NAMES)):\n",
    "        print(f\"Dataset has {counts[i]} events of type {EVENT_NAMES[i]}\")\n",
    "    \n",
    "\n",
    "print(\"Describing training dataset...\")\n",
    "describe_events(train_dataset)\n",
    "print()\n",
    "print(\"Describing test dataset...\")\n",
    "describe_events(test_dataset)"
   ]
  },
  {
   "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.10.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}