n_dot_braille_moment_encoding_python_reference_skeleton.py

"""
N-dot Braille Moment Encoding — Reference Implementation Skeleton

This module provides:
  • Cell/sequence representation with arbitrary n-dot ordering
  • Braille→vector pipelines: linear projection, codebook lookup, simple encoders
  • Weighted moments (μ, Σ, M3, M4) with streaming support for μ, Σ
  • Packing/round-trip back into n-dot cells (bit packing, VQ tiling, simple framing)
  • A façade op: NEW_DIST_MOMENTS[n, d]

Notes
-----
- This is a skeleton intended for extension. Some components (e.g., GI compression,
  robust CRC, higher-moment online updates) are stubbed with clear TODOs.
- Dependencies: numpy only (optional torch hooks are marked but not required).

Author: <you>
License: MIT
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterable, Optional, Tuple, List, Dict, Any
import numpy as np

Array = np.ndarray

# -----------------------------------------------------------------------------
# 1) n-dot Braille cell & sequences
# -----------------------------------------------------------------------------

@dataclass(frozen=True)
class BitOrder:
    """Fixed bit ordering π over n positions.

    π maps logical bit index j∈{0..n-1} (ranked by significance) to physical bit
    position in the cell. By convention j=0 corresponds to LSB.
    """
    pi: Tuple[int, ...]  # length n, a permutation of range(n)

    @property
    def n(self) -> int:
        return len(self.pi)

    def encode_cell(self, b: Array) -> int:
        """Encode binary cell vector b∈{0,1}^n into integer code c∈[0,2^n).
        Expects b shape (n,), dtype {bool,uint8,int}.
        """
        b = np.asarray(b, dtype=np.uint8).reshape(-1)
        assert b.size == self.n, "b must have length n"
        # Place bits by π: code = Σ b[π(j)] * 2^j
        j = np.arange(self.n, dtype=np.uint64)
        return int(np.sum((b[np.array(self.pi, dtype=int)] & 1).astype(np.uint64) << j))

    def decode_code(self, c: int) -> Array:
        """Decode integer code to binary vector b∈{0,1}^n using π inverse."""
        c = int(c)
        bits = ((c >> np.arange(self.n)) & 1).astype(np.uint8)  # logical order j
        # invert π: place logical bit j into physical index pi[j]
        b = np.zeros(self.n, dtype=np.uint8)
        b[np.array(self.pi, dtype=int)] = bits
        return b

@dataclass
class BrailleSequence:
    order: BitOrder
    cells: Array  # shape (n, L), dtype {0,1}

    @property
    def n(self) -> int:
        return self.order.n

    @property
    def L(self) -> int:
        return int(self.cells.shape[1])

    @classmethod
    def from_codes(cls, order: BitOrder, codes: Iterable[int]) -> "BrailleSequence":
        codes = list(map(int, codes))
        n = order.n
        L = len(codes)
        cells = np.stack([order.decode_code(c) for c in codes], axis=1).astype(np.uint8)
        assert cells.shape == (n, L)
        return cls(order=order, cells=cells)

    def to_codes(self) -> List[int]:
        return [self.order.encode_cell(self.cells[:, i]) for i in range(self.L)]

# -----------------------------------------------------------------------------
# 2) From n-dot Braille to semantic vectors
# -----------------------------------------------------------------------------

@dataclass
class LinearCellProjector:
    """Per-cell linear projection: v_i = E b_i, E∈R^{d×n}."""
    E: Array  # shape (d, n)

    @property
    def d(self) -> int:
        return int(self.E.shape[0])

    def __call__(self, B: BrailleSequence) -> Array:
        b = B.cells.astype(np.float32)  # (n, L)
        v = self.E @ b  # (d, L)
        return v

@dataclass
class CodebookLookup:
    """Codebook lookup: v_i = C[c_i], C∈R^{2^n × d}."""
    C: Array  # shape (2^n, d)
    order: BitOrder

    @property
    def d(self) -> int:
        return int(self.C.shape[1])

    def __call__(self, B: BrailleSequence) -> Array:
        codes = np.array(B.to_codes(), dtype=np.int64)  # (L,)
        v = self.C[codes]  # (L, d)
        return v.T  # (d, L)

@dataclass
class SequenceEncoder:
    """Simple sequence encoders (pooling/attention-lite/RNN-stub).

    Currently implements weighted pooling with optional mask weights α_i.
    """
    mode: str = "pool"  # "pool"

    def __call__(self, V: Array, alpha: Optional[Array] = None) -> Array:
        """V: (d, L); alpha: (L,) nonnegative or None.
        Returns v∈R^d.
        """
        d, L = V.shape
        if self.mode == "pool":
            if alpha is None:
                alpha = np.ones(L, dtype=np.float32)
            alpha = np.asarray(alpha, dtype=np.float32).reshape(1, L)
            wsum = float(alpha.sum())
            if wsum == 0.0:
                return np.zeros(d, dtype=np.float32)
            v = (V * alpha).sum(axis=1) / wsum
            return v
        else:
            raise NotImplementedError(f"SequenceEncoder mode={self.mode} not implemented")

# -----------------------------------------------------------------------------
# 3) Weighted semantic centroid and moments (batch + streaming)
# -----------------------------------------------------------------------------

@dataclass
class MomentsAccumulator:
    """Streaming (weighted) moments up to 2nd central moment.

    Maintains (W, μ, M2) for stable online covariance. Provides finalize() → (μ, Σ).
    For higher moments M3/M4, batch utilities are provided below.
    """
    d: int
    W: float = 0.0
    mu: Array = field(default_factory=lambda: np.zeros(0, dtype=np.float64))
    M2: Array = field(default_factory=lambda: np.zeros((0, 0), dtype=np.float64))

    def __post_init__(self):
        if self.mu.size == 0:
            object.__setattr__(self, 'mu', np.zeros(self.d, dtype=np.float64))
        if self.M2.size == 0:
            object.__setattr__(self, 'M2', np.zeros((self.d, self.d), dtype=np.float64))

    def add(self, v: Array, w: float = 1.0):
        v = np.asarray(v, dtype=np.float64).reshape(self.d)
        w = float(max(0.0, w))
        if w == 0.0:
            return
        delta = v - self.mu
        Wp = self.W + w
        mu_p = self.mu + (w / max(Wp, np.finfo(float).eps)) * delta
        # Welford-weighted update for M2
        self.M2 += w * np.outer(delta, (v - mu_p))
        self.mu = mu_p
        self.W = Wp

    def merge(self, other: "MomentsAccumulator"):
        assert other.d == self.d
        if other.W == 0.0:
            return
        if self.W == 0.0:
            self.W, self.mu, self.M2 = other.W, other.mu.copy(), other.M2.copy()
            return
        delta = other.mu - self.mu
        Wp = self.W + other.W
        mu_p = (self.W * self.mu + other.W * other.mu) / Wp
        self.M2 = self.M2 + other.M2 + (self.W * other.W / Wp) * np.outer(delta, delta)
        self.mu = mu_p
        self.W = Wp

    def finalize(self) -> Tuple[Array, Array]:
        if self.W == 0.0:
            return self.mu.copy(), np.zeros_like(self.M2)
        Sigma = self.M2 / self.W
        return self.mu.copy(), Sigma

# ---- Higher-order moments (batch utilities) ---------------------------------

def batch_moments(vs: Array, ws: Optional[Array] = None, order: int = 4) -> Dict[str, Array]:
    """Compute weighted central moments up to order in batch.

    vs: (N, d), ws: (N,) nonnegative or None.
    Returns dict keys: 'mu', 'Sigma', and optionally 'M3', 'M4'.
    """
    vs = np.asarray(vs, dtype=np.float64)
    N, d = vs.shape
    if ws is None:
        ws = np.ones(N, dtype=np.float64)
    else:
        ws = np.asarray(ws, dtype=np.float64).reshape(N)
        ws = np.maximum(ws, 0.0)
    W = float(ws.sum())
    if W == 0.0:
        mu = np.zeros(d, dtype=np.float64)
        Sigma = np.zeros((d, d), dtype=np.float64)
        out = dict(mu=mu, Sigma=Sigma)
        if order >= 3:
            out['M3'] = np.zeros((d, d, d), dtype=np.float64)
        if order >= 4:
            out['M4'] = np.zeros((d, d, d, d), dtype=np.float64)
        return out
    mu = (ws[:, None] * vs).sum(axis=0) / W
    xc = vs - mu
    Sigma = (ws[:, None, None] * np.einsum('ni,nj->nij', xc, xc)).sum(axis=0) / W
    out = dict(mu=mu, Sigma=Sigma)
    if order >= 3:
        M3 = (ws[:, None, None, None] * np.einsum('ni,nj,nk->nijk', xc, xc, xc)).sum(axis=0) / W
        out['M3'] = M3
    if order >= 4:
        M4 = (ws[:, None, None, None, None] * np.einsum('ni,nj,nk,nl->nijkl', xc, xc, xc, xc)).sum(axis=0) / W
        out['M4'] = M4
    return out

# -----------------------------------------------------------------------------
# 4) Packing back to n-dot Braille
# -----------------------------------------------------------------------------

@dataclass
class Framing:
    n: int
    d: int
    scale: float = 1.0
    payload_kind: int = 0  # 0: μ only, 1: μ+Σ, 2: μ+Σ+M3, 3: μ+Σ+M3+M4
    version: int = 1
    group_size: int = 128  # grouping for weights/tiles
    mk: int = 2            # Mk payload level aligned with moments (≥2 means μ+Σ)
    crc16: int = 0         # simple CRC placeholder (sum mod 2^16)

    def header_bits(self) -> List[int]:
        """Minimal illustrative header (not a final spec):
        [version(4) | n(8) | d(16) | payload_kind(2) | mk(3) | group_size(16) |
         scale exp2 log2-scale signed(8) | crc16(16)]
        Returns list of bits LSB-first.
        """
        bits: List[int] = []
        fields = [
            (self.version, 4),
            (self.n, 8),
            (self.d, 16),
            (self.payload_kind, 2),
            (max(0, int(self.mk)) & 0x7, 3),
            (max(0, int(self.group_size)) & 0xFFFF, 16),
            (int(np.clip(np.round(np.log2(max(self.scale, 1e-12))), -64, 63)) & 0xFF, 8),
            (max(0, int(self.crc16)) & 0xFFFF, 16),
        ]
        for val, width in fields:
            for j in range(width):
                bits.append((val >> j) & 1)
        return bits


def quantize_vector(x: Array, step: float) -> Tuple[Array, float]:
    step = float(step)
    q = np.round(x / step).astype(np.int64)
    return q, step


def dequantize_vector(q: Array, step: float) -> Array:
    return q.astype(np.float64) * float(step)


def int_to_bits(vals: Array, width: int) -> List[int]:
    vals = np.asarray(vals).reshape(-1)
    bits: List[int] = []
    for v in vals:
        vv = int(v)
        # two's complement for negative values
        if vv < 0:
            vv = (1 << width) + vv
        for j in range(width):
            bits.append((vv >> j) & 1)
    return bits


def bits_to_ints(bits: List[int], width: int, count: int) -> Array:
    out = np.zeros(count, dtype=np.int64)
    for i in range(count):
        v = 0
        for j in range(width):
            v |= (int(bits[i * width + j]) & 1) << j
        # sign recover
        if (v >> (width - 1)) & 1:
            v -= (1 << width)
        out[i] = v
    return out


def pack_bits_to_cells(bits: List[int], order: BitOrder) -> BrailleSequence:
    n = order.n
    L = (len(bits) + n - 1) // n
    cells = np.zeros((n, L), dtype=np.uint8)
    for i in range(L):
        chunk = bits[i * n : (i + 1) * n]
        chunk = chunk + [0] * (n - len(chunk)) if len(chunk) < n else chunk
        cells[:, i] = order.decode_code(sum(int(bit) << j for j, bit in enumerate(chunk)))
    return BrailleSequence(order=order, cells=cells)


def unpack_cells_to_bits(B: BrailleSequence) -> List[int]:
    bits: List[int] = []
    for code in B.to_codes():
        for j in range(B.n):
            bits.append((code >> j) & 1)
    return bits

# -----------------------------------------------------------------------------
# 5) Vector quantization (VQ) tiling
# -----------------------------------------------------------------------------

@dataclass
class Codebook:
    table: Array  # (K, dblk)

    @property
    def K(self) -> int:
        return int(self.table.shape[0])

    @property
    def dblk(self) -> int:
        return int(self.table.shape[1])

    @classmethod
    def random(cls, K: int, dblk: int, seed: int = 0) -> "Codebook":
        rng = np.random.default_rng(seed)
        return cls(table=rng.standard_normal((K, dblk)).astype(np.float32))

    def nearest(self, x: Array) -> int:
        # x: (dblk,)
        dif = self.table - x.reshape(1, -1)
        i = int(np.argmin(np.einsum('ij,ij->i', dif, dif)))
        return i


def vq_tile(mu: Array, C: Codebook) -> Tuple[List[int], Array]:
    """Blockwise VQ of μ using codebook C. Returns (codes, recon)."""
    mu = np.asarray(mu, dtype=np.float32)
    d = mu.size
    B = C.dblk
    blocks = int(np.ceil(d / B))
    codes: List[int] = []
    recon = np.zeros_like(mu)
    for k in range(blocks):
        sl = slice(k * B, min((k + 1) * B, d))
        x = np.zeros(B, dtype=np.float32)
        x[: mu[sl].size] = mu[sl]
        idx = C.nearest(x)
        codes.append(idx)
        recon_block = C.table[idx]
        recon[sl] = recon_block[: mu[sl].size]
    return codes, recon

# -----------------------------------------------------------------------------
# 6) Facade: NEW.DIST.MOMENTS[n, d]
# -----------------------------------------------------------------------------

@dataclass
class NewDistMoments:
    n: int
    d: int
    order: BitOrder
    projector: Optional[LinearCellProjector] = None
    codebook_lookup: Optional[CodebookLookup] = None
    encoder: SequenceEncoder = field(default_factory=SequenceEncoder)

    def encode_seq(self, B: BrailleSequence, weights: Optional[Array] = None,
                   return_per_cell: bool = False) -> Tuple[Array, Optional[Array]]:
        """Braille sequence → semantic vector (pooled) with optional per-cell.
        Returns (v, V) where v∈R^d and V∈R^{d×L} if return_per_cell.
        """
        V_parts: List[Array] = []
        if self.projector is not None:
            V_parts.append(self.projector(B))  # (d, L)
        if self.codebook_lookup is not None:
            V_parts.append(self.codebook_lookup(B))  # (d, L)
        if not V_parts:
            raise ValueError("No embedding path provided (projector/codebook)")
        V = np.sum(np.stack(V_parts, axis=0), axis=0)  # simple combine: sum
        alpha = None
        if weights is not None:
            alpha = np.asarray(weights, dtype=np.float32).reshape(-1)
            assert alpha.size == B.L
        v = self.encoder(V, alpha=alpha)
        return (v, V) if return_per_cell else (v, None)

    def accumulate(self, vecs: Iterable[Array], ws: Optional[Iterable[float]] = None) -> MomentsAccumulator:
        acc = MomentsAccumulator(d=self.d)
        if ws is None:
            ws = [1.0] * sum(1 for _ in vecs)  # consumes iterator; avoid in practice
        # Better: convert to list once
        vecs_list = list(vecs)
        ws_list = list(ws)
        for v, w in zip(vecs_list, ws_list):
            acc.add(v, float(w))
        return acc

    # ---- Packing APIs --------------------------------------------------------
    def pack_mu(self, mu: Array, step: float = 1/256, payload_kind: int = 0) -> BrailleSequence:
        frame = Framing(n=self.n, d=self.d, scale=step, payload_kind=payload_kind)
        hdr_bits = frame.header_bits()
        q, st = quantize_vector(mu, step)
        width = 16  # bits per component (two's complement)
        payload_bits = int_to_bits(q, width)
        bits = hdr_bits + payload_bits
        return pack_bits_to_cells(bits, self.order)

    def unpack_mu(self, B: BrailleSequence) -> Tuple[Array, Framing]:
        bits = unpack_cells_to_bits(B)
        # Recover toy header
        offs = 0
        def take(width: int) -> int:
            nonlocal offs
            v = 0
            for j in range(width):
                v |= (bits[offs] & 1) << j
                offs += 1
            return v
        version = take(4)
        n = take(8)
        d = take(16)
        payload_kind = take(2)
        scale_log2 = np.int8(take(8)).item()
        scale = float(np.exp2(scale_log2))
        frame = Framing(n=n, d=d, scale=scale, payload_kind=payload_kind, version=version)
        width = 16
        count = d
        remain = bits[offs: offs + width * count]
        q = bits_to_ints(remain + [0] * (width * count - len(remain)), width, count)
        mu = dequantize_vector(q, scale)
        return mu, frame

    def pack_full(self, mu: Array, Sigma: Array,
                  M3: Optional[Array] = None, M4: Optional[Array] = None,
                  step_mu: float = 1/256, step_S: float = 1/512,
                  payload_kind: int = 3) -> BrailleSequence:
        """Illustrative packing of μ, Σ (upper-triangular), and optional M3/M4.
        This is not a final bitstream spec; intended for experimentation.
        """
        frame = Framing(n=self.n, d=self.d, scale=step_mu, payload_kind=payload_kind)
        bits = frame.header_bits()
        # μ
        q_mu, _ = quantize_vector(mu, step_mu)
        bits += int_to_bits(q_mu, 16)
        # Σ upper triangular (row-major i<=j)
        iu = np.triu_indices(self.d)
        q_S, _ = quantize_vector(Sigma[iu], step_S)
        bits += int_to_bits(q_S, 16)
        # TODO: symmetric tensor bases for M3/M4; placeholder pack as full tensors if provided
        if M3 is not None:
            q_M3, _ = quantize_vector(M3.reshape(-1), step_S)
            bits += int_to_bits(q_M3, 16)
        if M4 is not None:
            q_M4, _ = quantize_vector(M4.reshape(-1), step_S)
            bits += int_to_bits(q_M4, 16)
        return pack_bits_to_cells(bits, self.order)

# -----------------------------------------------------------------------------
# 7) Sanity checks
# -----------------------------------------------------------------------------

def is_psd(Sigma: Array, tol: float = 1e-8) -> bool:
    # check eigenvalues ≥ -tol
    w = np.linalg.eigvalsh((Sigma + Sigma.T) * 0.5)
    return bool(np.all(w >= -tol))


def round_trip_error(mu: Array, sys: NewDistMoments, step: float = 1/256) -> float:
    B = sys.pack_mu(mu, step=step)
    mu2, _ = sys.unpack_mu(B)
    return float(np.linalg.norm(mu - mu2))

# -----------------------------------------------------------------------------
# 8) Demonstration (usage skeleton)
# -----------------------------------------------------------------------------

if __name__ == "__main__":
    # Define a 8-dot machine channel with a simple bit order π = identity
    n = 8
    order = BitOrder(pi=tuple(range(n)))

    # Build a toy sequence from random codes
    rng = np.random.default_rng(42)
    codes = rng.integers(0, 2**n, size=32)
    B = BrailleSequence.from_codes(order, codes)

    # Embedding pipeline: linear projection + codebook lookup (combined by sum)
    d = 16
    E = rng.standard_normal((d, n)).astype(np.float32) * 0.2
    projector = LinearCellProjector(E=E)

    Ctab = rng.standard_normal((2**n, d)).astype(np.float32) * 0.3
    cbook = CodebookLookup(C=Ctab, order=order)

    sys = NewDistMoments(n=n, d=d, order=order, projector=projector, codebook_lookup=cbook)

    # Encode sequence → pooled vector
    v, V = sys.encode_seq(B, return_per_cell=True)

    # Accumulate moments from multiple sequences (here just reuse jittered copies)
    acc = MomentsAccumulator(d=d)
    for k in range(10):
        noise = rng.standard_normal(d).astype(np.float32) * 0.1
        acc.add(v + noise, w=1.0)
    mu, Sigma = acc.finalize()

    # Batch higher moments (optional)
    vs = np.stack([v + rng.standard_normal(d) * 0.1 for _ in range(32)], axis=0)
    Ms = batch_moments(vs, order=4)

    # Packing μ back to n-dot cells and round-trip
    Bout = sys.pack_mu(mu, step=1/512)
    mu_rt, frame = sys.unpack_mu(Bout)
    err = np.linalg.norm(mu - mu_rt)

    print("μ‖: ", float(np.linalg.norm(mu)))
    print("round-trip error: ", err)
    print("Σ PSD? ", is_psd(Sigma))

    # VQ tiling example
    C = Codebook.random(K=2**4, dblk=4)
    codes_vq, recon = vq_tile(mu, C)
    print("VQ codes (len):", len(codes_vq))

    # Build a final Braille payload with μ + Σ (upper-triangular only)
    payload = sys.pack_full(mu, Sigma, step_mu=1/256, step_S=1/1024, payload_kind=1)
    print("Payload cells:", payload.cells.shape)