"""Lossless BF16 weight storage with a 12-bit common path and sparse escapes.

BF16's low byte is stored verbatim. A nibble encodes the sign and one of eight
adjacent high-byte exponents. Rare values outside that window use an exact
uint16 escape, located by a sorted-index binary search only on escape lanes.
The GPU reconstructs the original BF16 bits before the ordinary Metal GEMM.
This reduces resident weights without quantization or a new matrix algorithm.
"""

from __future__ import annotations

from dataclasses import dataclass
from functools import lru_cache

import mlx.core as mx
import numpy as np


@lru_cache(maxsize=1)
def _decoder():
    return mx.fast.metal_kernel(
        name="h3_unpack_bf16_palette12_v1",
        input_names=["low", "nibbles", "escape_indices", "escape_values", "base"],
        output_names=["out"],
        source=r"""
        uint i = thread_position_in_grid.x;
        if (i >= low_shape[0]) return;
        uint code = (nibbles[i / 2] >> ((i % 2) * 4)) & 15;
        uint count = escape_indices_shape[0];
        if (code == 0 && count != 0) {
            uint left = 0, right = count;
            while (left < right) {
                uint mid = (left + right) / 2;
                if (escape_indices[mid] < i) left = mid + 1;
                else right = mid;
            }
            out[i] = escape_values[left];
        } else {
            out[i] = ushort(low[i]) | ushort(((code & 7) + base[0]) << 8)
                   | ushort((code & 8) << 12);
        }
        """,
    )


@dataclass
class PackedBF16:
    shape: tuple[int, ...]
    low: mx.array
    nibbles: mx.array
    escape_indices: mx.array
    escape_values: mx.array
    base: mx.array

    @classmethod
    def pack(cls, value):
        if value.dtype != mx.bfloat16:
            raise ValueError("Lossless palette12 stores BF16 tensors only")
        raw = np.asarray(mx.contiguous(value).view(mx.uint16)).reshape(-1)
        high = ((raw >> 8) & 127).astype(np.uint8)
        base = max(0, int(high.max(initial=0)) - 7)
        sign = ((raw >> 15) << 3).astype(np.uint8)
        codes = np.clip(high.astype(np.int16) - base, 0, 7).astype(np.uint8) | sign
        # Nibble zero is the escape code. When base==0 the whole representable
        # high-byte range fits and zero can retain its ordinary interpretation.
        mask = high < base
        if base:
            mask |= (high == base) & (sign == 0)
        indices = np.flatnonzero(mask).astype(np.uint32)
        codes[indices] = 0
        if codes.size % 2:
            codes = np.pad(codes, (0, 1))
        values = [
            mx.array(raw.astype(np.uint8)),
            mx.array(codes[::2] | (codes[1::2] << 4)),
            mx.array(indices),
            mx.array(raw[indices]),
            mx.array([base], mx.uint32),
        ]
        mx.eval(values)
        return cls(tuple(value.shape), *values)

    @property
    def nbytes(self):
        return sum(
            x.nbytes
            for x in (
                self.low,
                self.nibbles,
                self.escape_indices,
                self.escape_values,
                self.base,
            )
        )

    def state(self):
        return (
            self.low,
            self.nibbles,
            self.escape_indices,
            self.escape_values,
            self.base,
        )

    def unpack(self, state=None):
        raw = _decoder()(
            inputs=self.state() if state is None else state,
            grid=(self.low.size, 1, 1),
            threadgroup=(256, 1, 1),
            output_shapes=[(self.low.size,)],
            output_dtypes=[mx.uint16],
        )[0]
        return raw.view(mx.bfloat16).reshape(self.shape)

    def tensors(self, prefix):
        return {
            f"{prefix}.{name}": getattr(self, name)
            for name in ("low", "nibbles", "escape_indices", "escape_values", "base")
        }

    @classmethod
    def from_tensors(cls, tensors, prefix, shape):
        result = cls(
            tuple(shape),
            *(
                tensors[f"{prefix}.{name}"]
                for name in (
                    "low",
                    "nibbles",
                    "escape_indices",
                    "escape_values",
                    "base",
                )
            ),
        )
        size = int(np.prod(shape))
        if result.low.shape != (size,) or result.nibbles.shape != ((size + 1) // 2,):
            raise ValueError("Invalid palette12 tensor dimensions")
        if result.escape_indices.shape != result.escape_values.shape:
            raise ValueError("Invalid palette12 escape arrays")
        return result


class PackedLinear:
    def __init__(self, weight, bias=None):
        self.weight, self.bias = weight, bias

    @classmethod
    def from_linear(cls, linear):
        if linear.scales is not None:
            raise ValueError("Lossless packing cannot be combined with quantization")
        return cls(PackedBF16.pack(linear.weight), linear.bias)

    @property
    def nbytes(self):
        return self.weight.nbytes + (0 if self.bias is None else self.bias.nbytes)

    @property
    def dtype(self):
        return mx.bfloat16

    def __call__(self, x, *, state=None):
        x = x.astype(mx.bfloat16)
        w = self.weight.unpack(state)
        return x @ w.T if self.bias is None else mx.addmm(self.bias, x, w.T)