{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from pathlib import Path\n",
    "import numpy as np\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to read hex color codes from a file and convert them to RGB tuples\n",
    "def load_palette(file_path):\n",
    "    palette = []\n",
    "    with open(file_path, 'r') as f:\n",
    "        for line in f:\n",
    "            hex_color = line.strip()\n",
    "            # Convert hex to RGB tuple\n",
    "            rgb_color = tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))\n",
    "            palette.append(rgb_color)\n",
    "    return palette\n",
    "\n",
    "pico8_palette = load_palette('../pico-8-secret-palette.hex')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(path: Path) -> Image.Image:\n",
    "    img = Image.open(path).convert('RGB')\n",
    "    w, h = img.size\n",
    "    if w > 128 or h > 128:\n",
    "        if w > h:\n",
    "            w, h = 128, int(round(h * (128 / w)))\n",
    "        else:\n",
    "            w, h = int(round(w * (128 / h))), 128\n",
    "        img = img.resize((w, h))\n",
    "    return img\n",
    "\n",
    "img = load_image('../peacock.jpg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def convert_to_indexed(img_np: np.array, palette: list) -> Image.Image:\n",
    "    pixels = img_np.reshape(-1, 3).astype(np.float32)\n",
    "    pal = np.array(palette, dtype=np.float32)\n",
    "    distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)\n",
    "    closest_idxs = np.argmin(distances, axis=1).astype(np.uint8).reshape(img_np.shape[1], -1)\n",
    "    result = Image.fromarray(closest_idxs, mode=\"P\")\n",
    "    result.putpalette(np.array(palette).reshape(-1).tolist())\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def squared_quantization_error(img_np: np.array, palette: np.array):\n",
    "    pixels = img_np.reshape(-1, 3).astype(np.float32)\n",
    "    pal = np.array(palette, dtype=np.float32)\n",
    "    sq_distances = np.sum((pixels[:, None, :] - pal[None, :, :]) ** 2, axis=2)\n",
    "    min_errs = np.min(sq_distances, axis=1)\n",
    "    return np.sum(min_errs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4252 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=P size=128x128>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def pick_palette_indexes(img, n, palette):\n",
    "    img_np = np.array(img)\n",
    "    result = list(range(n))\n",
    "    num_steps = 0\n",
    "    pal_idx = 0\n",
    "    while num_steps < n:\n",
    "        num_steps += 1\n",
    "        old_val = result[pal_idx]\n",
    "        result[pal_idx] = None\n",
    "        best_err = math.inf\n",
    "        best_val = None\n",
    "        for val in [i for i in range(len(palette)) if i not in result]:\n",
    "            result[pal_idx] = val\n",
    "            err = squared_quantization_error(img_np, [palette[i] for i in result])\n",
    "            if err < best_err:\n",
    "                best_err = err\n",
    "                best_val = val\n",
    "        result[pal_idx] = best_val\n",
    "        if best_val != old_val:\n",
    "            num_steps = 0\n",
    "        pal_idx = (pal_idx + 1) % len(result)\n",
    "    return result\n",
    "\n",
    "pal_idxs = pick_palette_indexes(img, 16, pico8_palette)\n",
    "\n",
    "\n",
    "pal = [pico8_palette[i] for i in pal_idxs]\n",
    "convert_to_indexed(np.array(img), pal)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAIAAAACABAMAAAAxEHz4AAAAMFBMVEUAAAAdK1N+JVMAh1GrUjZfV0/Cw8f/8ej/AE3/owD/7CcA5DYprf+Ddpz/d6j/zKruozp8AAAK7UlEQVR4AbWZMW/jOBbHHRyuNyG4S3CfgSAylYzFAnapxQixtvMEg8VoK4+7cbXIZwsEXOUtjNjuFDVzpe3idjyVrfv/H0lJlp1ksMA9TCSK5PvxvcdHUvJ0UpHcy3IwGAyTgX/Ml8PkdelYwMRrEJAkg6h6fl09Sf5PgFFlAcx5XTrplE5MKpOtC1UQfgRQktAGuCgsK8DwBUM607INQNcR7BgMcl68XlXwFe4OF8p9wwKoAEBNJ67fCLUXpZOWYoJ3AWro57XtExVHL1uQShh9FK3KqCa4YUc07JIwDzgPDgBFdrwEeMEEl0gOsDwDEEcBciiF9sUDZB6pbwdiQcRr8aGtK88eIC4sq24jligOMGL5bQB7OQ0WKQ2AK7YwzgKZxaaGlHlx/UdV6SJA9Jfs31CQR//cUqsfxQI7h6IwdE0jqz4YrOq+F0sEWH3J/GpArx+9CRDreVmKjh/FA/Lc17xw75wChr6bA0SXAHe3vhfuLUDbhYsAZV4BDF1bbcGq0dsWVXCLZlfdtsDVjxwADl4AXBHgetYAOwsu9V4FvFdG2sWGBqA5D6eA8akZ76+0bSehCciXOc1gtQfYII5PHYkRBPhmF0qn2sBRRbHV4EByld8iD8ZjaSGYMlgpY1CzFHc71Sm2nko3KrpNMQmNybQ2pgFY8dxUKviAzmJsBViX5SFfT1wgiHkykFCbzAQ3wh4mSZ4AtlKBed44EypAgd0ZJ8SEURC5M0FotJnPH81cQQ0KSEveV6pzXew3MAHICjAlAHJYWv3I0AFlggwmGE3FPB8LBS5cF2W5cQAbRXjghbYNIqiF4WMYzjVCFuob6nMyePmHNuj/Gf1qC3i+eCHgzsx1MDdZkF1nYWZ0T/Tt5edAmX+X5X7TBEy9Nu5wLjLmEf/mWaBgAMRGQQCrWHXULRRqAKJzClg+ieOMQpDnT8YEqm9HT6J8/UlhHouy3DITmEhL/J24kDMCUDZ9ztgqQBADgETW03LRBiDlWoChYf6YfhqjL5S1CYO+I+STrtJXBhYcxQKJ+akLh/ypH5tAB+vJVS+BuYCEag4AM3m80qjKGgDMexuQ5wtl1Dpd9JKkh9L4I1LKW5DDIUULSmcBAIgMnitJJ+uFUuP0k0pUEo/TdfoJXjhAOrkzmJopen9lEDFnBORFpV4in9OF7qUTGHz1z2SMnX8VaCVpTGtj0w0Mu9cA1OYVYJ9iyazEgVWi4ARWSAzACguJZ8gSgQ36DQDyjtwKgE4YbJ0vYHUSI2rpBDOp+mPEEL6mKwBME0B9EApHwCsf9NMJu91CPVmnkwUBqB0xWrnCBJ0DImdC+nu6wjBfkEJIIjiAECw6V8qGQAIZI7P+tC5gMUkQaxP2XxAyyCRWmDmDVYhyVwHQE2VYcCT73TnAmjD7Au01lFYcxxgU09UVc6mH3WCSL6N1GXftNB65p3WwkUkMBhGD4A0QAkKFO6JxhdUUmF9X+e4IwlcYZ27ReWsBXJXiAo2yw+OKSCax+pWlGLNmsDUVuYwxwSwEpj/FcnYArAlI9GwB9JRqC6VukAEpVr/GiBoL+LCM2Oeb6l6psKgtOIoP0Q5te8xAiuskyTGNAf4S3LD2AFD3+2IAtfKgINwPXAzElMHgFzSV+ymvkMMKDiD7vBDwriwONOEQdLudE4D4sKXenhcrC4zCZYelzEmUmUvL5a4s/8LyVAxiZQEBv6DlVOA7NjRsf6KNS788lAcYesDs4gHyVaYRhU2eiwEnhG+YbAjmXARbwBTtS7oAy+pEotJ2u93xfiqIf4DoBUHAeUC5KIvyUJS/IUcZgoYF2/Px2X6I6TqEFxAKVn4uP3J7ZLF2QZ7sZTZ7qJ8O7zEuCYikpgWQ/WRxLQFAeetjwAYrM4p/YM4w4gJAqWDDAzakPj60KMczgOg3CbFoW4gpxLr/YlWkdgs9A3x3gNoGmIAAcN7gzD0AD+VTovVN4SzApiold6kANeFOQxcAhEC/m832s8PTChVT0dhyV24CvAG4+2pmDUUjrYrZ7Mvv+/iuApy50AA8OMJetOFCYAh4WKcxUoqLGdIG1B6QJF3KErbzBYOAe1TGPUZD8pCrkDtSLZcAe0MAF2NgfprNuFUGgTPgDMBxa7HgPZzn6Jpb0uxhjJPS3K6/28bXLPB2HaBOQjdA/s/+WPUC6jv/vg5fcuHB65ewAMId3pifSpwwegx9235ctgDldzpQ+vAJBBZoSQIVIpXxtgP/fYCPSIOTIFbDNgoEwH8NE/QUKymf1o0MwdsA8SATLxq6pBy5Ib0J+A0AHYQIpD2Q6/H5kvYDgCcB4IVNa78LOAZOA4TgLQt4BmBTM9mjtudpbQEOtLcBWyxn5qDJQq3/VSujdMTB9hbguN0uBID9HXNZnAC2OM9eBWCnhsQCoAFKnQCOPJIRgnYM+HtMeRRdXjbQgxi8uuvetGkBDbCAXaOap3Kafq70t5wDAoJ5oG7QUVpwP25pAD1IOs0zYW9/U6sBC9FHGJGPN6jm2Z8/o7DxBhAAsVZLcx7J7Rm1uQ0B9MNHpT9s0AC7pZkFMUA+PNGZjfLajQaIPOX5xhtwjaXwDPWIjRFaea8BrrvUssUKarH/igR4dReAtLQAqLNAp9e4RbEDIIr6dpMvfZsdylngKy/cI6ePbRhvBw2A7fs2YOgA3A30uBGi1wH1r4kLC5gjj4NbeN42MqF0qi991wzDRq4YE3Ct51mmFQx4ATD03Z0W89MWrQeBwXenUpjEOoiNrvjsa8vQ/yD6M8bHjo6v37npSRKed4UL7cr6OQaAWTifZ+qDpFrdZktw9zUAtRH/OQD6+SJA9sQ2tXpmCHCg0ABze/wbgJgACCJgxnu7Eiu6K7zqAvQrQIGdYIN5qMQBEAQfxKg9y0M5U/FqgBiExf4zlzfnktuBcMioAaxaOqy9LbACICFcCDMLAKMWm1feAkGephrfZISAGOj7dFer2j3w84Y2D60LkQM0bMDHIRZxRg9CI/uh3bV2Vh1Xty1KDDygtiGK8I6NH08YAe7IlVqjYI+WywB+/qvreUZ1pAHe8BqKVZGEUxfEE7vpLZTEXwgAcE8/FwCGYoELgdwQB97vmgD8erc7V+fb+mBQASTKFWmDD+dQhsdFI6OKSwB+K3nAM79ZjjZFnrf4rMGZ3ACo+1cBEfOLPZivgvqm6UIoUcR+pHrFGwDRogmQHfp+JCDkNGTcj1pHs4dVLiBs9ogkgPrlAm7PBYC1qFTXppJX9PcaMABg56t5XyAE3Msg9KCr+kWz2ZVrQOQsqDrF9r2GiCxT3a4yedVWFxqAZ+uCb8O7VcCVRAuwI0JM8w3TdzsB7Hwt7werbwFwAUlxKRNeBnxrAmhAcJ3dN0ew5a1PJPz6cPrtumgAsuBRKfwW13rPJKJK5QsARiDLwhBpQMD8MfzzzAKuRpvKg0EriDaGshvgXX/+iM0hnE/bBHrgAQNkENtxPe74I0dAE5DKyAIAEAUzL1oAMaACwIRa8G4kADeHAghC96lWYcSAChDV6tvtfwRgPcDgAnA/4FXqOKvgQO0Co1DJAhaLAzaL8JIHkZ8Qj76P7MknAC4HyJKVFiCJiLWoPaDAfoHW0wNENhQqQapNOXYxCHkmBDqc04BQ9W23+jpEsQmoWjwAQRTA4zUAcGtZ9bCFFwEcEEHAUsBrek+r7FqjSgWjHwS4NwMAMvzXAtLYWIDqtQB8vOSCB8CDMMPW9giGlQuA/wHleVRDWkJkqQAAAABJRU5ErkJggg==",
      "text/plain": [
       "<PIL.Image.Image image mode=P size=128x128>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "convert_to_indexed(np.array(img), pico8_palette[:16])"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
