"""BitNetBase: 1-bit and ternary inference on hyperdimensional primitives.

The observation this module is built on is that a 1-bit dot product and an
HDC similarity are *the same operation*.  For sign vectors ``a, b`` in
{-1, +1}^n, packed one bit per weight with +1 -> 1 and -1 -> 0:

    disagree = popcount(a_bits ^ b_bits)          # XOR is HDC's bind
    dot      = n - 2 * disagree                   # popcount is HDC's sim

so ``sim(a, b) = dot(a, b) / n``.  A BitNet layer's entire forward pass is
therefore XOR-and-bitcount over packed registers — no multiply, no float
accumulate — and :func:`dot_1bit_packed` verifies that identity against the
plain integer dot at runtime rather than asserting it in prose.

Two things live here:

``BitLinear``    a ternary {-1, 0, +1} linear layer in the spirit of BitNet
                 b1.58: per-tensor absmean scaling with a dead zone, so small
                 weights collapse to 0 instead of costing compute
``PrototypeNet`` a 1-bit classifier whose weights *are* hypervectors in an
                 :class:`~hdc.memory.HDCStorage` cell — training is bundling,
                 inference is a similarity argmax, and nothing is ever a
                 float array

Standard library only; vectors are Python ints, so the kernels below are
bitwise ops over whole machine words.
"""

from __future__ import annotations

import json
import random
from typing import Sequence

from .hypervec import DEFAULT_DIM, bundle, hvc, mask, popcount, sim
from .memory import HDCStorage

__all__ = [
    "pack_signs",
    "unpack_signs",
    "dot_1bit",
    "dot_1bit_packed",
    "ternarize",
    "quantize_absmax",
    "ternary_dot",
    "ternary_matvec",
    "BitLinear",
    "PrototypeNet",
    "selfcheck",
]


# --------------------------------------------------------------------------
# 1-bit kernels
# --------------------------------------------------------------------------
def pack_signs(vec: Sequence[int]) -> int:
    """Pack a {-1, +1} sequence into a bitmask, +1 -> 1, -1 -> 0.

    Bit ``i`` of the result is weight ``i``.  Any non-positive entry counts as
    -1, which keeps zero weights meaningful: a packed zero contributes -1 to
    the dot product, so ternary matrices are packed with an explicit sparse
    form (see :func:`ternary_dot`) rather than through this function.
    """
    out = 0
    for i, v in enumerate(vec):
        if v > 0:
            out |= 1 << i
    return out


def unpack_signs(bits: int, n: int) -> list[int]:
    """Inverse of :func:`pack_signs`: ``n`` weights as +1/-1."""
    return [1 if (bits >> i) & 1 else -1 for i in range(n)]


def dot_1bit(a: Sequence[int], b: Sequence[int]) -> int:
    """Plain integer dot product of two +/-1 sequences (the reference form)."""
    if len(a) != len(b):
        raise ValueError("length mismatch: %d vs %d" % (len(a), len(b)))
    return sum(x * y for x, y in zip(a, b))


def dot_1bit_packed(a_bits: int, b_bits: int, n: int) -> int:
    """Bitwise dot product of two packed sign vectors.

    ``n - 2 * popcount(a ^ b)`` — one XOR and one bit-count for the whole
    vector, independent of ``n`` in loop terms.  This is the identity the
    module exists to demonstrate: bind, then count disagreement.
    """
    disagree = popcount((a_bits ^ b_bits) & mask(n))
    return n - 2 * disagree


# --------------------------------------------------------------------------
# ternary quantization
# --------------------------------------------------------------------------
def ternarize(weights: Sequence[float], ratio: float = 0.7) -> list[int]:
    """Quantize to {-1, 0, +1} with a per-tensor absmean scale.

    ``scale = mean(|w|)``, and weights inside ``ratio * scale`` of zero are
    mapped to 0.  With ``ratio = 1.0`` this is nearest-rounding of BitNet
    b1.58's absmean scheme; lower ratios keep more weights non-zero (denser,
    slightly more accurate), higher ratios sparsify the layer.
    """
    items = [float(w) for w in weights]
    if not items:
        return []
    scale = sum(abs(w) for w in items) / len(items)
    if scale == 0.0:
        return [0] * len(items)
    dead = ratio * scale
    out: list[int] = []
    for w in items:
        if w > dead:
            out.append(1)
        elif w < -dead:
            out.append(-1)
        else:
            out.append(0)
    return out


def quantize_absmax(x: Sequence[float], bits: int = 8) -> tuple[list[int], float]:
    """Symmetric activation quantization to signed integers.

    Returns ``(integers, scale)`` such that ``x ~= integer * scale``.  ``scale``
    is 1.0 for an all-zero input, which keeps the call total instead of
    dividing by zero.
    """
    if not 1 <= bits <= 16:
        raise ValueError("bits must be in 1..16")
    qmax = (1 << (bits - 1)) - 1
    peak = max((abs(float(v)) for v in x), default=0.0)
    if peak == 0.0:
        return [0] * len(x), 1.0
    scale = peak / qmax
    out = []
    for v in x:
        q = round(float(v) / scale)
        out.append(max(-qmax - 1, min(qmax, q)))
    return out, scale


def ternary_dot(w: Sequence[int], x: Sequence[int]) -> int:
    """Integer dot product of a ternary weight row with an integer activation."""
    if len(w) != len(x):
        raise ValueError("length mismatch: %d vs %d" % (len(w), len(x)))
    total = 0
    for wi, xi in zip(w, x):
        if wi == 1:
            total += xi
        elif wi == -1:
            total -= xi
        # wi == 0 contributes nothing: the point of sparsity
    return total


def ternary_matvec(matrix: Sequence[Sequence[int]], x: Sequence[int]) -> list[int]:
    """Apply a ternary matrix to an integer vector."""
    return [ternary_dot(row, x) for row in matrix]


# --------------------------------------------------------------------------
# a ternary linear layer
# --------------------------------------------------------------------------
class BitLinear:
    """A ternary {-1, 0, +1} linear layer with additive-only accumulation.

    Weights are drawn once from a seeded normal distribution and ternarized at
    construction, so two layers with the same shape and seed are identical.
    ``forward`` accumulates with adds and subtracts only: a weight of -1 is a
    subtraction and a weight of 0 is skipped, which is what makes 1-bit
    inference cheap on hardware that has no multiplier.

    :meth:`scale` exposes the absmean scale discarded at ternarization time, so
    callers that want a dequantized output can multiply it back in; ``forward``
    itself returns the raw integer accumulators.
    """

    def __init__(self, in_features: int, out_features: int, seed: int = 0):
        if in_features <= 0 or out_features <= 0:
            raise ValueError("layer dimensions must be positive")
        self.in_features = in_features
        self.out_features = out_features
        self.seed = seed
        rng = random.Random(seed)
        raw = [[rng.gauss(0.0, 1.0) for _ in range(in_features)] for _ in range(out_features)]
        self.ternary_weights: list[list[int]] = [ternarize(row) for row in raw]
        self._scales = [
            sum(abs(w) for w in row) / in_features for row in raw
        ]

    def scale(self, row: int) -> float:
        """Absmean scale of output row ``row`` (multiply an output by this)."""
        return self._scales[row]

    def forward(self, x: Sequence[int]) -> list[int]:
        """Integer accumulators for one input vector."""
        if len(x) != self.in_features:
            raise ValueError(
                "expected %d inputs, got %d" % (self.in_features, len(x))
            )
        return ternary_matvec(self.ternary_weights, x)

    def density(self) -> float:
        """Fraction of non-zero weights (measured, not assumed)."""
        total = self.in_features * self.out_features
        live = sum(1 for row in self.ternary_weights for w in row if w != 0)
        return live / total if total else 0.0


# --------------------------------------------------------------------------
# 1-bit classifier whose weights live in HDC storage
# --------------------------------------------------------------------------
class PrototypeNet:
    """Nearest-prototype classifier with class weights stored as hypervectors.

    Training a class is bundling its samples into a single majority vector, so
    the class weight is a superposition — the HDC centroid.  Inference compares
    a probe against each class cell by Hamming similarity and takes the argmax.
    The weights are ordinary :class:`HDCStorage` cells, so they persist with
    ``HDCStorage.save`` and are addressable by content like everything else.
    """

    def __init__(
        self,
        storage: HDCStorage | None = None,
        d: int = DEFAULT_DIM,
        prefix: str = "proto/",
    ):
        self.d = d
        self.prefix = prefix
        self.storage = storage if storage is not None else HDCStorage(d=d)
        self._samples: dict[str, list[int]] = {}

    def _cell(self, label: str) -> str:
        return self.prefix + label

    def fit(self, label: str, sample: int) -> None:
        """Add one hypervector sample to a class and rewrite its weight cell.

        The class weight is recomputed as the majority bundle of every sample
        seen so far rather than folded in pairwise: pairwise bundling of two
        vectors is AND-like and would forget the first sample quickly.
        """
        samples = self._samples.setdefault(label, [])
        samples.append(sample & mask(self.d))
        weight = bundle(samples, self.d)
        self.storage.write(self._cell(label), weight, symbol=True)

    def classes(self) -> list[str]:
        """Class labels, in the order they were first fitted."""
        return list(self._samples)

    def weights(self) -> list[tuple[str, int]]:
        """``(label, weight vector)`` pairs as currently stored."""
        return [(label, self.storage.read(self._cell(label))) for label in self._samples]

    def predict(self, sample: int) -> tuple[str, float]:
        """Best matching class and its similarity; ``('', 0.0)`` when untrained."""
        best_label = ""
        best = -2.0
        for label in self._samples:
            s = sim(sample, self.storage.read(self._cell(label)), self.d)
            if s > best:
                best, best_label = s, label
        if not best_label:
            return "", 0.0
        return best_label, best

    def accuracy(self, pairs: Sequence[tuple[str, int]]) -> float:
        """Fraction of ``(expected_label, sample)`` pairs classified correctly."""
        if not pairs:
            return 0.0
        hits = 0
        for label, sample in pairs:
            got, _ = self.predict(sample)
            if got == label:
                hits += 1
        return hits / len(pairs)


# --------------------------------------------------------------------------
# self-verification
# --------------------------------------------------------------------------
def _noisy(base: int, rate: float, rng: random.Random, d: int) -> int:
    """Flip each bit of ``base`` with probability ``rate``."""
    v = base
    for i in range(d):
        if rng.random() < rate:
            v ^= 1 << i
    return v & mask(d)


def selfcheck() -> dict:
    """Exercise the kernels and return measured values (nothing hardcoded).

    Covers the packed/plain dot-product equivalence, the sparsity a ternary
    layer actually reaches, and end-to-end prototype classification accuracy
    under 2% bit noise.
    """
    rng = random.Random(20240501)

    # 1. the identity: packed XOR dot == integer dot
    trials = 25
    mismatches = 0
    for _ in range(trials):
        n = rng.choice((16, 64, 256, 1000))
        signs_a = [rng.choice((-1, 1)) for _ in range(n)]
        signs_b = [rng.choice((-1, 1)) for _ in range(n)]
        plain = dot_1bit(signs_a, signs_b)
        packed = dot_1bit_packed(pack_signs(signs_a), pack_signs(signs_b), n)
        if plain != packed:
            mismatches += 1
    xor_matches_dot = mismatches == 0

    # 2. measured sparsity of a ternarized gaussian layer
    layer = BitLinear(128, 32, seed=7)
    zeros = sum(1 for row in layer.ternary_weights for w in row if w == 0)
    ternary_sparsity = zeros / (128 * 32)

    # 3. prototype accuracy: 4 classes, 6 training samples, 8 noisy test samples
    d = 2000
    net = PrototypeNet(d=d)
    test_pairs: list[tuple[str, int]] = []
    for c in range(4):
        base = hvc("selfcheck/class/%d" % c, d)
        for _ in range(6):
            net.fit("class%d" % c, _noisy(base, 0.02, rng, d))
        for _ in range(2):
            test_pairs.append(("class%d" % c, _noisy(base, 0.02, rng, d)))

    return {
        "xor_matches_dot": xor_matches_dot,
        "xor_trials": trials,
        "ternary_sparsity": round(ternary_sparsity, 4),
        "proto_density": round(layer.density(), 4),
        "prototype_accuracy": round(net.accuracy(test_pairs), 4),
        "prototype_classes": len(net.classes()),
        "prototype_dim": d,
        "d": DEFAULT_DIM,
    }


if __name__ == "__main__":  # pragma: no cover - manual entry point
    print(json.dumps(selfcheck(), indent=2, sort_keys=True))