{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "from collections import Counter\n",
    "import heapq\n",
    "from pathlib import Path\n",
    "import statistics\n",
    "from bitarray import bitarray\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def transform_context(img_np: np.ndarray):\n",
    "    h, w = img_np.shape\n",
    "    front = [0] * (h + 1 + w + 1)\n",
    "    frequency_ctx = {None: Counter()}\n",
    "    order_ctx = {None: list(range(16))}\n",
    "\n",
    "    result = np.zeros((h, w), dtype=np.uint8)\n",
    "    for y in range(h):\n",
    "        for x in range(w):\n",
    "            # Get current pixel value and context\n",
    "            corner_idx = h - y + x\n",
    "            context = tuple(front[corner_idx - 1:corner_idx + 3])\n",
    "            assert len(context) == 4, f\"bad context {context} ({x=},{y=},{corner_idx=})\"\n",
    "            pixel = int(img_np[y, x])\n",
    "            front[corner_idx] = pixel\n",
    "\n",
    "            # Get appropriate frequency Counter and order list\n",
    "            if context in frequency_ctx:\n",
    "                frequency = frequency_ctx[context]\n",
    "                order = order_ctx[context]\n",
    "            else:\n",
    "                frequency = frequency_ctx[None] # separate context for misses\n",
    "                order = order_ctx[None]\n",
    "\n",
    "            # Eager M2F transform\n",
    "            idx = order.index(pixel)\n",
    "            result[y, x] = idx\n",
    "            frequency[pixel] += 1\n",
    "            # if idx > 0:\n",
    "            #     print(f\"encoding {pixel=}, surprised at {x=},{y=}: {context=},{frequency=},{order=}\")\n",
    "            while idx > 0 and frequency[order[idx - 1]] <= frequency[pixel]:\n",
    "                order[idx], order[idx - 1] = order[idx - 1], order[idx]\n",
    "                idx -= 1\n",
    "\n",
    "            # Populate context state based on a copy of the default\n",
    "            if context not in frequency_ctx:\n",
    "                frequency_ctx[context] = frequency.copy()\n",
    "                order_ctx[context] = order.copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       "       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=uint8)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def checkerboard(h, w):\n",
    "    result = np.zeros((h, w), dtype=np.uint8)\n",
    "    for y in range(h):\n",
    "        for x in range(w):\n",
    "            result[y,x] = (x ^ y) & 1\n",
    "    return result\n",
    "\n",
    "transform_context(checkerboard(7,11))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def huffman(counts: dict):\n",
    "    h = [(v, (k,)) for k, v in counts.items()]\n",
    "    prefixes = {k: \"\" for k in counts.keys()}\n",
    "    heapq.heapify(h)\n",
    "    while(len(h) > 1):\n",
    "        (c1, v1), (c2, v2) = heapq.heappop(h), heapq.heappop(h)\n",
    "        heapq.heappush(h, (c1 + c2, v1 + v2))\n",
    "        for v in v1:\n",
    "            prefixes[v] = \"0\" + prefixes[v]\n",
    "        for v in v2:\n",
    "            prefixes[v] = \"1\" + prefixes[v]\n",
    "    keys = sorted(prefixes.keys())\n",
    "    return {k: prefixes[k] for k in keys}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average length of 0th code = 1.00 (quartiles [1.0, 1.0, 1.0])\n",
      "Average length of 1th code = 2.67 (quartiles [2.0, 3.0, 3.0])\n",
      "Average length of 2th code = 3.52 (quartiles [3.0, 4.0, 4.0])\n",
      "Average length of 3th code = 4.10 (quartiles [4.0, 4.0, 4.0])\n",
      "Average length of 4th code = 4.67 (quartiles [4.0, 5.0, 5.0])\n",
      "Average length of 5th code = 5.00 (quartiles [4.5, 5.0, 5.5])\n",
      "Average length of 6th code = 5.45 (quartiles [5.0, 5.0, 6.0])\n",
      "Average length of 7th code = 5.95 (quartiles [5.0, 6.0, 6.0])\n",
      "Average length of 8th code = 6.40 (quartiles [5.25, 6.0, 7.0])\n",
      "Average length of 9th code = 6.90 (quartiles [6.0, 6.5, 7.0])\n",
      "Average length of 10th code = 6.95 (quartiles [6.0, 7.0, 8.0])\n",
      "Average length of 11th code = 7.53 (quartiles [7.0, 7.0, 8.0])\n",
      "Average length of 12th code = 8.22 (quartiles [7.0, 8.0, 9.0])\n",
      "Average length of 13th code = 8.62 (quartiles [8.0, 9.0, 9.0])\n",
      "Average length of 14th code = 8.69 (quartiles [8.0, 8.0, 10.0])\n",
      "Average length of 15th code = 8.75 (quartiles [8.0, 8.0, 9.75])\n"
     ]
    }
   ],
   "source": [
    "img_code_lengths = []\n",
    "for p in Path('../corpus/quantized/').glob(\"*.*\"):\n",
    "    img = Image.open(p)\n",
    "    #np.array(img)\n",
    "    img_np = np.array(img)\n",
    "    transformed = transform_context(img_np)\n",
    "    c = Counter(transformed.reshape(-1))\n",
    "    code_lengths = [len(x) for x in huffman(c).values()]\n",
    "    img_code_lengths.append(code_lengths)\n",
    "\n",
    "for i in range(16):\n",
    "    x = []\n",
    "    for code_lengths in img_code_lengths:\n",
    "        if i < len(code_lengths):\n",
    "            x.append(code_lengths[i])\n",
    "    mean = statistics.mean(x)\n",
    "    quarts = statistics.quantiles(x)\n",
    "    print(f\"Average length of {i}th code = {mean:1.2f} (quartiles {quarts})\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'e', 'e', 'e'], [0, 1, 2])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def transform_rle(inputs):\n",
    "    outputs = []\n",
    "    for val in inputs:\n",
    "        last_val = outputs[-1][0] if len(outputs) > 0 else None\n",
    "        if last_val == val:\n",
    "            outputs[-1][1] += 1\n",
    "        else:\n",
    "            outputs.append([val, 1])\n",
    "    return [tuple(x) for x in outputs]\n",
    "\n",
    "def transform_trig_rle(inputs, trigger_length):\n",
    "    runs = transform_rle(inputs)\n",
    "    values = []\n",
    "    additional_run_lengths = []\n",
    "    for value, run_length in runs:\n",
    "        num_literals = min(run_length, trigger_length)\n",
    "        values.extend(value for _ in range(num_literals))\n",
    "        if run_length >= trigger_length:\n",
    "            additional_run_lengths.append(run_length - trigger_length)\n",
    "    return values, additional_run_lengths\n",
    "\n",
    "transform_trig_rle(\"abbcccddddeeeee\", 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Encoding value 0 as 0 (length 1)\n",
      "Encoding value 1 as 111 (length 3)\n",
      "Encoding value 2 as 100 (length 3)\n",
      "Encoding value 3 as 1100 (length 4)\n",
      "Encoding value 4 as 1010 (length 4)\n",
      "Encoding value 5 as 11010 (length 5)\n",
      "Encoding value 6 as 10110 (length 5)\n",
      "Encoding value 7 as 110111 (length 6)\n",
      "Encoding value 8 as 101111 (length 6)\n",
      "Encoding value 9 as 1101101 (length 7)\n",
      "Encoding value 10 as 1011101 (length 7)\n",
      "Encoding value 11 as 1011100 (length 7)\n",
      "Encoding value 12 as 11011000 (length 8)\n",
      "Encoding value 13 as 110110011 (length 9)\n",
      "Encoding value 14 as 1101100101 (length 10)\n",
      "Encoding value 15 as 1101100100 (length 10)\n"
     ]
    }
   ],
   "source": [
    "def analyze_huffman(counts, name):\n",
    "    huffman_map = huffman(counts)\n",
    "    for k, v in huffman_map.items():\n",
    "        print(f\"Encoding {name} {k} as {v} (length {len(v)})\")\n",
    "\n",
    "value_counter = Counter()\n",
    "length_counter = Counter()\n",
    "for p in Path('../corpus/quantized/').glob(\"*.*\"):\n",
    "    img = Image.open(p)\n",
    "    img_np = np.array(img)\n",
    "    context_coded = transform_context(img_np)\n",
    "    values, lengths = transform_trig_rle(context_coded.reshape(-1), 3)\n",
    "    value_counter.update(values)\n",
    "    length_counter.update(lengths)\n",
    "\n",
    "analyze_huffman(value_counter, 'value')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What did I learn about the distribution of values (post trig-RLE transform)?\n",
    "\n",
    "The main takeaway is that the codes look very geometric. For example, running trigger length 2 on the corpus and Huffman coding the values results in codes of length 2, 2, 3, 3, 4, 4, and so on. Note that the number of codes of each length is pretty much constant, much like a Golomb/Rice code. And those are optimal/near-optimal for geometric distributions.\n",
    "\n",
    "Running trigger length 3 and above seems to give a slightly different distribution: the Huffman codes have lengths *1*, 3, 3, 4, 4, etc. This is basically just a single prefix bit with a possible Rice code attached: the zero value is encoded as \"0\", while subsequent values are encoded as \"1\" + Rice code.\n",
    "\n",
    "Keep in mind, this is specific to the RLE _values_. The RLE _lengths_ are a different story."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(1, 5988), (2, 4715), (3, 4074), (4, 2536), (5, 1396), (6, 815), (7, 542), (8, 16), (9, 11), (10, 10), (11, 8), (12, 2)]\n",
      "Encoding bucket 1 as 11 (length 2)\n",
      "Encoding bucket 2 as 01 (length 2)\n",
      "Encoding bucket 3 as 00 (length 2)\n",
      "Encoding bucket 4 as 100 (length 3)\n",
      "Encoding bucket 5 as 1010 (length 4)\n",
      "Encoding bucket 6 as 10111 (length 5)\n",
      "Encoding bucket 7 as 101101 (length 6)\n",
      "Encoding bucket 8 as 10110011 (length 8)\n",
      "Encoding bucket 9 as 10110010 (length 8)\n",
      "Encoding bucket 10 as 10110000 (length 8)\n",
      "Encoding bucket 11 as 101100011 (length 9)\n",
      "Encoding bucket 12 as 101100010 (length 9)\n"
     ]
    }
   ],
   "source": [
    "# Try buckets based on number of bits\n",
    "bucket_frequencies = Counter()\n",
    "for length, count in length_counter.items():\n",
    "    bucket = (length + 1).bit_length()\n",
    "    bucket_frequencies[bucket] += count\n",
    "print([(i, bucket_frequencies[i]) for i in range(1, max(bucket_frequencies.keys()) + 1)])\n",
    "analyze_huffman(bucket_frequencies, 'bucket')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we use bucket algorithm `length.bit_length()` and tweak the codes slightly to make them monotonically increasing in size, we get an encoding scheme that looks like this:\n",
    "\n",
    "```\n",
    "00  (run 0)\n",
    "01  (run 1)\n",
    "100x    (runs 2-3) (implied leading 1 bit)\n",
    "101xx   (runs 4-7)\n",
    "110xxx  (runs 8-15)\n",
    "1110xxxx    (runs 16-31)\n",
    "11110xxxxx  (runs 32-63)\n",
    "111110xxxxxx    (runs 64-128)\n",
    "```\n",
    "\n",
    "If we change the bucket algorithm to `(length + 1).bit_length` so there's always a leading 1 bit, we don't have to reorder any codes, and we get something similar:\n",
    "\n",
    "```\n",
    "00  (run 0)\n",
    "01x (runs 1-2)\n",
    "10xx    (runs 3-6)\n",
    "110xxx  (runs 7-14)\n",
    "1110xxxx    (runs 15-30)\n",
    "11110xxxxx  (runs 31-62)\n",
    "111110xxxxx (runs 63-126)\n",
    "```\n",
    "\n",
    "The code lengths aren't exactly the same, but this is reminiscient of Exp-Golomb coding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../corpus/quantized/floodedcaves_0.png compressed to 32.4% (2658 bytes)\n",
      "../corpus/quantized/age_of_ants-title.png compressed to 77.1% (6318 bytes)\n",
      "../corpus/quantized/build a jetpack_0.png compressed to 20.4% (1668 bytes)\n",
      "../corpus/quantized/shinkansen 1 3_1.png compressed to 21.3% (1744 bytes)\n",
      "../corpus/quantized/pico8_build_a_jetpack-1.png compressed to 25% (2050 bytes)\n",
      "../corpus/quantized/iso_ray-09_000.png compressed to 54.1% (4435 bytes)\n",
      "../corpus/quantized/age of ants_2.png compressed to 41.4% (3389 bytes)\n",
      "../corpus/quantized/pico8_mot_raymarch3-0.png compressed to 62.7% (5136 bytes)\n",
      "../corpus/quantized/age of ants_5.png compressed to 43.6% (3569 bytes)\n",
      "../corpus/quantized/donsol for pico-8  v1 8 4 _2.png compressed to 15.2% (1243 bytes)\n",
      "../corpus/quantized/pico8_witchcrafttd-6.png compressed to 25.1% (2057 bytes)\n",
      "../corpus/quantized/pico8_bunnysurvivor-9.png compressed to 27.2% (2228 bytes)\n",
      "../corpus/quantized/pico8_rotslimepires_1_1-0.png compressed to 57% (4670 bytes)\n",
      "../corpus/quantized/pico8_donsol8_v1-14.png compressed to 18.3% (1499 bytes)\n",
      "../corpus/quantized/pico8_ppwr-5.png compressed to 41.3% (3381 bytes)\n",
      "../corpus/quantized/pico8_px9-9.png compressed to 23.3% (1905 bytes)\n",
      "../corpus/quantized/build a jetpack_5.png compressed to 20.3% (1665 bytes)\n",
      "../corpus/quantized/pico20068.png compressed to 34.1% (2790 bytes)\n",
      "../corpus/quantized/storming the grandmothership_1.png compressed to 37.5% (3070 bytes)\n",
      "../corpus/quantized/build a jetpack_9.png compressed to 20.4% (1675 bytes)\n",
      "../corpus/quantized/hersheys_train_line_0.png compressed to 52% (4259 bytes)\n"
     ]
    }
   ],
   "source": [
    "def append_exp_golomb(dest: bitarray, val: int):\n",
    "    val += 1\n",
    "    dest.extend(1 for _ in range(val.bit_length() - 1))\n",
    "    dest.append(0)\n",
    "    for i in range(val.bit_length() - 2, -1, -1):\n",
    "        dest.append((val >> i) & 1)\n",
    "\n",
    "def append_rice1(dest: bitarray, val: int, frequentZero=False):\n",
    "    if frequentZero:\n",
    "        if val == 0:\n",
    "            dest.append(0)\n",
    "            return\n",
    "        else:\n",
    "            dest.append(1)\n",
    "            val -= 1\n",
    "    dest.extend(1 for _ in range(val >> 1))\n",
    "    dest.append(0)\n",
    "    dest.append(val & 1)\n",
    "\n",
    "def encode_mixed_trig_rle(values, xlengths, trigger_length):\n",
    "    result = bitarray()\n",
    "    runs = transform_rle(values)\n",
    "    xlengths_iter = iter(xlengths)\n",
    "    for value, length in runs:\n",
    "        for _ in range(length):\n",
    "            append_rice1(result, value, True)\n",
    "        assert length <= trigger_length\n",
    "        if length == trigger_length:\n",
    "            append_exp_golomb(result, next(xlengths_iter))\n",
    "    assert next(xlengths_iter, None) == None\n",
    "    return result\n",
    "\n",
    "def encode_image(img, trigger_length=3):\n",
    "    img_np = np.array(img)\n",
    "    ctx_transformed = transform_context(img_np)\n",
    "    values, xlengths = transform_trig_rle(ctx_transformed.reshape(-1), trigger_length)\n",
    "    return encode_mixed_trig_rle(values, xlengths, trigger_length)\n",
    "\n",
    "for p in Path('../corpus/quantized/').glob('*.*'):\n",
    "    img = Image.open(p)\n",
    "    raw_size_bytes = math.ceil(img.width * img.height * 4 / 8)\n",
    "    encoded = encode_image(img, trigger_length=2)\n",
    "    encoded_size_bytes = math.ceil(len(encoded) / 8)\n",
    "    encoded_percent = encoded_size_bytes * 100 / raw_size_bytes\n",
    "    print(f\"{p} compressed to {encoded_percent:.3g}% ({encoded_size_bytes} bytes)\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These percents are pretty much universally worse than PX9. The only image that benefited was the pattern-dithered `iso_ray-09_000`. Possible reasons for the poor performance:\n",
    "- Too much context causes dilution?\n",
    "- Maybe frequency-based M2F is suboptimal somehow?"
   ]
  }
 ],
 "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
}