{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bitarray import bitarray\n",
    "from itertools import cycle\n",
    "from util import walk_corpus\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "123\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "bitarray('1000')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "BitSequence = list[int] | bitarray\n",
    "\n",
    "def bans_encode_bignum(bits: BitSequence, zero_probabilities: int | list[int], n=8):\n",
    "    if type(zero_probabilities) is int:\n",
    "        zero_probabilities = [zero_probabilities] * len(bits)\n",
    "    else:\n",
    "        assert len(bits) == len(zero_probabilities)\n",
    "    probability_total = 1 << n\n",
    "    assert all(0 < p < probability_total for p in zero_probabilities)\n",
    "\n",
    "    result = probability_total - 1\n",
    "    for bit, p0 in zip(reversed(bits), reversed(zero_probabilities)):\n",
    "        p = p0 if bit == 0 else probability_total - p0\n",
    "        quotient, remainder = divmod(result, p)\n",
    "        symbol_idx = remainder + (0 if bit == 0 else p0)\n",
    "        result = (quotient << n) + symbol_idx\n",
    "    return result\n",
    "\n",
    "def bans_decode_bignum(value: int, zero_probabilities: int | list[int], n=8):\n",
    "    total_probability = 1 << n\n",
    "    if type(zero_probabilities) is int:\n",
    "        assert 0 < zero_probabilities < total_probability\n",
    "        zero_probabilities = cycle([zero_probabilities])\n",
    "    else:\n",
    "        assert all(0 < p < total_probability for p in zero_probabilities)\n",
    "\n",
    "    result = bitarray()\n",
    "    mask = total_probability - 1\n",
    "    while value > mask:\n",
    "        symbol_idx = value & mask\n",
    "        p0 = next(zero_probabilities)\n",
    "        bit = 0 if symbol_idx < p0 else 1\n",
    "        result.append(bit)\n",
    "        p = p0 if bit == 0 else total_probability - p0\n",
    "        remainder = symbol_idx - (0 if bit == 0 else p0)\n",
    "        value = (value >> n) * p + remainder\n",
    "    return result\n",
    "\n",
    "t = bans_encode_bignum([1, 0, 0, 0], 11, n=4)\n",
    "print(t)\n",
    "bans_decode_bignum(t, 11, n=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bytearray(b'\\x99\\x15\\xfdu\\xc0\\xa1\\xf4!\\xccj\\x80') 11\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "bitarray('100100100100100100100100100100100100100100100100100100100100100100100100100100100100100100')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def bans_encode(bits: BitSequence, zero_probabilities: int | list[int], n=8):\n",
    "    if type(zero_probabilities) is int:\n",
    "        zero_probabilities = [zero_probabilities] * len(bits)\n",
    "    else:\n",
    "        assert len(bits) == len(zero_probabilities)\n",
    "    total_probability = 1 << n\n",
    "    assert all(0 < p < total_probability for p in zero_probabilities)\n",
    "\n",
    "    result = bytearray()\n",
    "    value = total_probability - 1\n",
    "    def flush_byte():\n",
    "        nonlocal result\n",
    "        nonlocal value\n",
    "        result.append(value & 0xff)\n",
    "        value >>= 8\n",
    "    def encode(x, p0, bit):\n",
    "        p = p0 if bit == 0 else total_probability - p0\n",
    "        quotient, remainder = divmod(x, p)\n",
    "        symbol_idx = remainder + (0 if bit == 0 else p0)\n",
    "        return (quotient << n) | symbol_idx\n",
    "\n",
    "    for bit, p0 in zip(reversed(bits), reversed(zero_probabilities)):\n",
    "        next_value = encode(value, p0, bit)\n",
    "        if next_value.bit_length() > n + 8:\n",
    "            flush_byte()\n",
    "            next_value = encode(value, p0, bit)\n",
    "        value = next_value\n",
    "    while value > 0:\n",
    "        flush_byte()\n",
    "    result.reverse()\n",
    "    return result\n",
    "\n",
    "\n",
    "def bans_decode(data: bytes | bytearray, zero_probabilities: int | list[int], n=8):\n",
    "    total_probability = 1 << n\n",
    "    if type(zero_probabilities) is int:\n",
    "        assert 0 < zero_probabilities < total_probability\n",
    "        zero_probabilities = cycle([zero_probabilities])\n",
    "    else:\n",
    "        assert all(0 < p < total_probability for p in zero_probabilities)\n",
    "\n",
    "    result = bitarray()\n",
    "    value = 0\n",
    "    data = iter(data)\n",
    "    def read_byte():\n",
    "        nonlocal value, data\n",
    "        try:\n",
    "            value = (value << 8) | next(data)\n",
    "            return True\n",
    "        except StopIteration:\n",
    "            return False\n",
    "\n",
    "    mask = total_probability - 1\n",
    "    while value.bit_length() <= n:\n",
    "        if not read_byte():\n",
    "            break\n",
    "    while True:\n",
    "        if value.bit_length() <= n:\n",
    "            if not read_byte():\n",
    "                break\n",
    "\n",
    "        symbol_idx = value & mask\n",
    "        p0 = next(zero_probabilities)\n",
    "        bit = 0 if symbol_idx < p0 else 1\n",
    "        result.append(bit)\n",
    "        p = p0 if bit == 0 else total_probability - p0\n",
    "        remainder = symbol_idx - (0 if bit == 0 else p0)\n",
    "        value = (value >> n) * p + remainder\n",
    "    return result\n",
    "\n",
    "t = bans_encode([1, 0, 0] * 30, 21, n=5)\n",
    "print(t, len(t))\n",
    "bans_decode(t, 21, n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'-0.91, -0.90, -0.89, -0.88, -0.86, -0.83, -0.80, -0.75, -0.67, -0.50, 0.00, 0.50, 0.67, 0.75, 0.80, 0.83, 0.86, 0.88, 0.89, 0.90, 0.91'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def sigmoid(x):\n",
    "    return x / (1.0 + abs(x))\n",
    "\n",
    "\", \".join(f\"{sigmoid(x):.2f}\" for x in range(-10, 11))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'-40, -36, -32, -28, -24, -20, -16, -12, -8, -4, 0, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def isigmoid(x):\n",
    "    # Trying to make something that returns an i8. Might not be 100% true.\n",
    "    # Precision of input is on my mind\n",
    "    return ((x << 3) // (1 + abs(x >> 4)))\n",
    "\n",
    "\", \".join(f\"{isigmoid(x)}\" for x in range(-10, 11))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What does the context look like for binary-encoded data?\n",
    "\n",
    "We can divide it up into discrete pieces, each of which are counted separately, and are combined via the sigmoid function:\n",
    "\n",
    "- Partial encoding of the current pixel: for example, if the pixel is \"1011\" we will pass through contexts \"\", \"1\", \"10\", and \"101\".\n",
    "- Full encoding of NW neighbor (omitted if it doesn't exist)\n",
    "- Full encoding of N neighbor (ditto for all of these)\n",
    "- Full encoding of NE neighbor\n",
    "- Full encoding of W neighbor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p=PosixPath('../corpus/quantized/floodedcaves_0.png'), bytes=1743 (percent=21.3%, bits per pixel=0.85107421875)\n",
      "p=PosixPath('../corpus/quantized/age_of_ants-title.png'), bytes=3222 (percent=39.3%, bits per pixel=1.5732421875)\n",
      "p=PosixPath('../corpus/quantized/build a jetpack_0.png'), bytes=1223 (percent=14.9%, bits per pixel=0.59716796875)\n",
      "p=PosixPath('../corpus/quantized/shinkansen 1 3_1.png'), bytes=1178 (percent=14.4%, bits per pixel=0.5751953125)\n",
      "p=PosixPath('../corpus/quantized/pico8_build_a_jetpack-1.png'), bytes=1116 (percent=13.6%, bits per pixel=0.544921875)\n",
      "p=PosixPath('../corpus/quantized/iso_ray-09_000.png'), bytes=3614 (percent=44.1%, bits per pixel=1.7646484375)\n",
      "p=PosixPath('../corpus/quantized/age of ants_2.png'), bytes=2190 (percent=26.7%, bits per pixel=1.0693359375)\n",
      "p=PosixPath('../corpus/quantized/pico8_mot_raymarch3-0.png'), bytes=4233 (percent=51.7%, bits per pixel=2.06689453125)\n",
      "p=PosixPath('../corpus/quantized/age of ants_5.png'), bytes=1988 (percent=24.3%, bits per pixel=0.970703125)\n",
      "p=PosixPath('../corpus/quantized/donsol for pico-8  v1 8 4 _2.png'), bytes=1006 (percent=12.3%, bits per pixel=0.4912109375)\n",
      "p=PosixPath('../corpus/quantized/pico8_witchcrafttd-6.png'), bytes=1272 (percent=15.5%, bits per pixel=0.62109375)\n",
      "p=PosixPath('../corpus/quantized/pico8_bunnysurvivor-9.png'), bytes=1189 (percent=14.5%, bits per pixel=0.58056640625)\n",
      "p=PosixPath('../corpus/quantized/pico8_rotslimepires_1_1-0.png'), bytes=3144 (percent=38.4%, bits per pixel=1.53515625)\n",
      "p=PosixPath('../corpus/quantized/pico8_donsol8_v1-14.png'), bytes=1236 (percent=15.1%, bits per pixel=0.603515625)\n",
      "p=PosixPath('../corpus/quantized/pico8_ppwr-5.png'), bytes=1742 (percent=21.3%, bits per pixel=0.8505859375)\n",
      "p=PosixPath('../corpus/quantized/pico8_px9-9.png'), bytes=1064 (percent=13.0%, bits per pixel=0.51953125)\n",
      "p=PosixPath('../corpus/quantized/build a jetpack_5.png'), bytes=1117 (percent=13.6%, bits per pixel=0.54541015625)\n",
      "p=PosixPath('../corpus/quantized/pico20068.png'), bytes=1464 (percent=17.9%, bits per pixel=0.71484375)\n",
      "p=PosixPath('../corpus/quantized/storming the grandmothership_1.png'), bytes=2231 (percent=27.2%, bits per pixel=1.08935546875)\n",
      "p=PosixPath('../corpus/quantized/build a jetpack_9.png'), bytes=1315 (percent=16.1%, bits per pixel=0.64208984375)\n",
      "p=PosixPath('../corpus/quantized/hersheys_train_line_0.png'), bytes=2613 (percent=31.9%, bits per pixel=1.27587890625)\n",
      "Total bytes: 39900\n"
     ]
    }
   ],
   "source": [
    "class Predictor:\n",
    "    def __init__(self):\n",
    "        self.counts = [Counter() for _ in range(4)]\n",
    "        gain = 1.27\n",
    "        weights = [30, 100] * 2\n",
    "        self.weights = [x/sum(weights) * gain for x in weights]\n",
    "        self.decay = 0.9\n",
    "\n",
    "    def _keys(self, neighbors, partial_pixel):\n",
    "        assert len(neighbors) == 4\n",
    "        return tuple(neighbor | (partial_pixel << 4) for neighbor in neighbors)\n",
    "\n",
    "    def predict(self, contexts):\n",
    "        \"\"\"Returns probability that the next bit is zero, out of 256\"\"\"\n",
    "        total = 0.0\n",
    "        for counter, key, weight in zip(self.counts, contexts, self.weights):\n",
    "            total += weight * counter[key]\n",
    "        prob_float = (sigmoid(total) + 1.0) / 2.0\n",
    "        prob_int = round(prob_float * 256.0)\n",
    "        return min(max(1, prob_int), 255)\n",
    "\n",
    "    def update(self, contexts, bit):\n",
    "        \"\"\"Updates the model with the actual bit\"\"\"\n",
    "        # Note for the future: code might be simpler if we used P(bit=1) everywhere\n",
    "        delta = -int(bit) * 2 + 1\n",
    "        for counter, key in zip(self.counts, contexts):\n",
    "            counter[key] = counter[key] * self.decay + delta\n",
    "\n",
    "\n",
    "def bitwise_encode(img: np.array):\n",
    "    height, width = img.shape\n",
    "    assert img.dtype == np.uint8\n",
    "\n",
    "    bits = bitarray()\n",
    "    zero_probabilities = [] # values are 1 through 255\n",
    "    pred = Predictor()\n",
    "    num_missed_predictions = 0\n",
    "    for y, row in enumerate(img):\n",
    "        for x, val in enumerate(row):\n",
    "            neighbors = [31] * 4 # 11111 means \"no neighbor\", 0xxxx means \"neighbor is that color\"\n",
    "            if y > 0:\n",
    "                if x > 0:\n",
    "                    neighbors[0] = img[y - 1, x - 1] # NW\n",
    "                neighbors[1] = img[y - 1, x] # N\n",
    "                if x < width - 1:\n",
    "                    neighbors[2] = img[y - 1, x + 1] # NE\n",
    "            if x > 0:\n",
    "                neighbors[3] = img[y, x - 1] # W\n",
    "            # We'll store context as 1(5 neighbor bits)(0-3 partial pixel bits)\n",
    "            for i in range(4):\n",
    "                neighbors[i] = int(neighbors[i]) | (1 << 5)\n",
    "            for i in range(3, -1, -1):\n",
    "                bit = int((val >> i) & 1)\n",
    "                bits.append(bit)\n",
    "                prediction = pred.predict(neighbors)\n",
    "                missed_prediction = (prediction < 16 and bit == 0) or (prediction > 255 - 16 and bit == 1)\n",
    "                if missed_prediction:\n",
    "                    num_missed_predictions += 1\n",
    "                zero_probabilities.append(prediction)\n",
    "                pred.update(neighbors, bit)\n",
    "                for i in range(4):\n",
    "                    neighbors[i] = (neighbors[i] << 1) | bit\n",
    "    #print(f\"Missed predictions: {num_missed_predictions} out of {len(bits)}\")\n",
    "\n",
    "    return bans_encode(bits, zero_probabilities, n=8)\n",
    "\n",
    "\n",
    "def corpus_bitwise_encode():\n",
    "    byte_size = 0\n",
    "    for p in walk_corpus():\n",
    "        img = np.array(Image.open(p))\n",
    "        t = len(bitwise_encode(img))\n",
    "        byte_size += t\n",
    "        print(f\"{p=}, bytes={t} (percent={t/8192*100:.3}%, bits per pixel={t*8/128**2})\")\n",
    "    print(\"Total bytes:\", byte_size)\n",
    "\n",
    "corpus_bitwise_encode()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Things learned so far:\n",
    "\n",
    "- ANS is promising!\n",
    "- Diagonal neighbors currently count for 30% relative to orthogonal neighbors.\n",
    "- It's important to bound counts (e.g., with a decay) so they don't go out of control.\n",
    "    - It looks like we had been confidently wrong (1-15 or 240-255) about 3-4% of the time?\n",
    "    - Assuming a wrong guess costs 6 bits and 2500 wrong guesses/image: this costs ~2k per image or ~40k for the corpus\n",
    "    - We shaved off about half that (~20k) by adding decay and fine-tuning it\n",
    "- If we don't need the precision, we could reduce probability granularities.\n",
    "- Compression bit-by-bit (no RLE, etc) gets us 23.2% compression on average. This sometimes beats PX9!"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}