code/wal_tat/fv5_conditional.py

"""Lossless conditional-high packing for five-valued row weights.

The public FV5 layout uses three dense bitplanes.  The third plane only says
whether a non-zero value uses the high magnitude, so storing it densely costs
one bit for every weight even when high values are uncommon.  This reference
format instead stores two dense tag planes and one compact sign stream only
for high-magnitude entries::

    00 ->  0
    01 -> +1
    10 -> -1
    11 -> high magnitude; consume one compact sign bit (+2/-2)

The exact code cost is therefore ``2*N + N_high`` bits.  Two FP32 row scales
are appended unchanged.  Packing is deterministic, byte padding must be zero,
and decoding rejects trailing data and inconsistent high counts.
"""
from __future__ import annotations

from dataclasses import dataclass
import hashlib
import math
import struct
from typing import Any

import numpy as np


MAGIC = b"WALFV5C1"
HEADER = struct.Struct("<8sIIQQQQ")
ALPHABET = np.asarray((-2, -1, 0, 1, 2), dtype=np.int8)


def _as_numpy(value: Any) -> np.ndarray:
    if hasattr(value, "detach"):
        value = value.detach().cpu().numpy()
    return np.asarray(value)


def _codes_2d(value: Any) -> np.ndarray:
    codes = _as_numpy(value)
    if codes.ndim != 2 or codes.shape[0] < 1 or codes.shape[1] < 1:
        raise ValueError("FV5 codes must be a non-empty rank-2 matrix")
    if not np.isin(codes, ALPHABET).all():
        raise ValueError("FV5 codes escape {-2,-1,0,+1,+2}")
    return np.asarray(codes, dtype=np.int8, order="C")


def _scales_1d(value: Any, *, rows: int, name: str) -> np.ndarray:
    scales = _as_numpy(value)
    if scales.shape == (rows, 1):
        scales = scales[:, 0]
    if scales.shape != (rows,):
        raise ValueError(f"{name} scales must have shape [{rows}] or [{rows},1]")
    return np.asarray(scales, dtype="<f4", order="C")


def _packed_size(bit_count: int) -> int:
    if bit_count < 0:
        raise ValueError("bit count must be non-negative")
    return (int(bit_count) + 7) // 8


def _pack_bits(bits: np.ndarray) -> bytes:
    flat = np.asarray(bits, dtype=np.uint8, order="C").reshape(-1)
    if flat.size and np.any(flat > 1):
        raise ValueError("bit array contains a non-binary value")
    return np.packbits(flat, bitorder="little").tobytes()


def _unpack_bits(payload: memoryview, *, count: int, label: str) -> np.ndarray:
    expected = _packed_size(count)
    if len(payload) != expected:
        raise ValueError(f"{label} byte count is inconsistent")
    if not payload and count == 0:
        return np.empty(0, dtype=np.uint8)
    packed = np.frombuffer(payload, dtype=np.uint8)
    bits = np.unpackbits(packed, bitorder="little")
    if np.any(bits[count:]):
        raise ValueError(f"{label} has non-zero padding bits")
    return bits[:count].copy()


@dataclass(frozen=True)
class ConditionalFV5Accounting:
    rows: int
    columns: int
    logical_weights: int
    high_codes: int
    tag_plane_bytes: int
    high_sign_bytes: int
    scale_bytes: int
    serialized_bytes: int

    @property
    def code_bits_exact(self) -> int:
        return 2 * self.logical_weights + self.high_codes

    @property
    def serialized_bpw(self) -> float:
        return self.serialized_bytes * 8 / self.logical_weights

    def as_dict(self) -> dict[str, int | float]:
        return {
            "rows": self.rows,
            "columns": self.columns,
            "logical_weights": self.logical_weights,
            "high_codes": self.high_codes,
            "code_bits_exact": self.code_bits_exact,
            "tag_plane_bytes_each": self.tag_plane_bytes,
            "high_sign_bytes": self.high_sign_bytes,
            "scale_bytes": self.scale_bytes,
            "header_bytes": HEADER.size,
            "serialized_bytes": self.serialized_bytes,
            "serialized_bpw": self.serialized_bpw,
        }


def conditional_fv5_accounting(codes: Any) -> ConditionalFV5Accounting:
    values = _codes_2d(codes)
    rows, columns = map(int, values.shape)
    weights = rows * columns
    high = int(np.count_nonzero(np.abs(values) == 2))
    plane_bytes = _packed_size(weights)
    high_bytes = _packed_size(high)
    scale_bytes = rows * 2 * np.dtype("<f4").itemsize
    total = HEADER.size + 2 * plane_bytes + high_bytes + scale_bytes
    return ConditionalFV5Accounting(
        rows=rows,
        columns=columns,
        logical_weights=weights,
        high_codes=high,
        tag_plane_bytes=plane_bytes,
        high_sign_bytes=high_bytes,
        scale_bytes=scale_bytes,
        serialized_bytes=total,
    )


def encode_conditional_fv5(
    codes: Any,
    low_scales: Any,
    high_scales: Any,
) -> bytes:
    """Encode one exact FV5 matrix and its two FP32 row scales."""

    values = _codes_2d(codes)
    accounting = conditional_fv5_accounting(values)
    low = _scales_1d(low_scales, rows=accounting.rows, name="low")
    high = _scales_1d(high_scales, rows=accounting.rows, name="high")
    if (
        not np.isfinite(low).all()
        or not np.isfinite(high).all()
        or np.any(low <= 0)
        or np.any(high <= low)
    ):
        raise ValueError("row scales must satisfy finite 0 < low < high")

    flat = values.reshape(-1)
    high_mask = np.abs(flat) == 2
    # 00 zero, 01 +low, 10 -low, 11 high.  High sign is compacted.
    left = np.logical_or(flat < 0, high_mask)
    right = np.logical_or(flat > 0, high_mask)
    high_negative = flat[high_mask] < 0
    left_payload = _pack_bits(left)
    right_payload = _pack_bits(right)
    high_payload = _pack_bits(high_negative)
    if (
        len(left_payload) != accounting.tag_plane_bytes
        or len(right_payload) != accounting.tag_plane_bytes
        or len(high_payload) != accounting.high_sign_bytes
    ):
        raise RuntimeError("conditional FV5 packing length mismatch")

    header = HEADER.pack(
        MAGIC,
        accounting.rows,
        accounting.columns,
        accounting.tag_plane_bytes,
        accounting.high_codes,
        accounting.high_sign_bytes,
        accounting.scale_bytes,
    )
    encoded = b"".join(
        (
            header,
            left_payload,
            right_payload,
            high_payload,
            low.tobytes(),
            high.tobytes(),
        )
    )
    if len(encoded) != accounting.serialized_bytes:
        raise RuntimeError("conditional FV5 serialized length mismatch")
    return encoded


def decode_conditional_fv5(
    encoded: bytes | bytearray | memoryview,
) -> dict[str, np.ndarray | int]:
    """Decode and fully validate one conditional-high FV5 frame."""

    view = memoryview(encoded)
    if len(view) < HEADER.size:
        raise ValueError("conditional FV5 frame is shorter than its header")
    magic, rows, columns, plane_bytes, high_count, high_bytes, scale_bytes = (
        HEADER.unpack_from(view)
    )
    if magic != MAGIC:
        raise ValueError("unsupported conditional FV5 frame magic")
    if rows < 1 or columns < 1:
        raise ValueError("conditional FV5 geometry must be positive")
    weights = int(rows) * int(columns)
    if plane_bytes != _packed_size(weights):
        raise ValueError("conditional FV5 tag-plane length is inconsistent")
    if high_count > weights or high_bytes != _packed_size(int(high_count)):
        raise ValueError("conditional FV5 high-sign length is inconsistent")
    if scale_bytes != int(rows) * 2 * np.dtype("<f4").itemsize:
        raise ValueError("conditional FV5 scale length is inconsistent")
    expected = HEADER.size + 2 * plane_bytes + high_bytes + scale_bytes
    if len(view) != expected:
        raise ValueError("conditional FV5 frame has trailing or truncated data")

    offset = HEADER.size
    left = _unpack_bits(
        view[offset : offset + plane_bytes], count=weights, label="left tag plane"
    ).astype(bool)
    offset += plane_bytes
    right = _unpack_bits(
        view[offset : offset + plane_bytes], count=weights, label="right tag plane"
    ).astype(bool)
    offset += plane_bytes
    high_mask = np.logical_and(left, right)
    observed_high = int(np.count_nonzero(high_mask))
    if observed_high != int(high_count):
        raise ValueError("conditional FV5 high count disagrees with tag planes")
    high_negative = _unpack_bits(
        view[offset : offset + high_bytes],
        count=int(high_count),
        label="high sign stream",
    ).astype(bool)
    offset += high_bytes

    codes = np.zeros(weights, dtype=np.int8)
    codes[np.logical_and(~left, right)] = 1
    codes[np.logical_and(left, ~right)] = -1
    high_indices = np.flatnonzero(high_mask)
    codes[high_indices] = np.where(high_negative, -2, 2).astype(np.int8)

    scale_payload = view[offset:]
    low = np.frombuffer(scale_payload[: int(rows) * 4], dtype="<f4").copy()
    high = np.frombuffer(scale_payload[int(rows) * 4 :], dtype="<f4").copy()
    if (
        not np.isfinite(low).all()
        or not np.isfinite(high).all()
        or np.any(low <= 0)
        or np.any(high <= low)
    ):
        raise ValueError("decoded row scales violate finite 0 < low < high")
    return {
        "rows": int(rows),
        "columns": int(columns),
        "codes": codes.reshape(int(rows), int(columns)),
        "low_scales": low,
        "high_scales": high,
    }


def conditional_fv5_sha256(encoded: bytes | bytearray | memoryview) -> str:
    return hashlib.sha256(encoded).hexdigest()


def aggregate_conditional_fv5_bpw(
    accountings: list[ConditionalFV5Accounting],
) -> float:
    if not accountings:
        raise ValueError("at least one FV5 matrix accounting is required")
    weights = sum(item.logical_weights for item in accountings)
    serialized = sum(item.serialized_bytes for item in accountings)
    if weights < 1 or not math.isfinite(serialized / weights):
        raise ValueError("invalid aggregate FV5 accounting")
    return serialized * 8 / weights