#!/usr/bin/env python3
"""
cdna_stream.py - TRUE Compressed Domain Compute from HXZ Seeds

This is the REAL CDC: Weight matrix W is NEVER decoded to float32.

Architecture:
  HXZ seed contains:
    - indices: [M, K] uint8 - which codebook entry each weight uses
    - codebook: [256] float32 - the 256 possible weight values

  Compute path (per row block):
    1. Stream index block [B, K] uint8 from seed
    2. For each row: bincount(indices[row], weights=x) → weighted sums per bin
    3. Dot product with codebook → y[row]
    4. W is NEVER materialized as float32!

Memory comparison:
  - Dense: M × K × 4 bytes (float32)
  - CDNA:  B × K bytes (uint8) + 256 × 4 bytes (codebook) + K × 4 bytes (x)
  - Savings: ~4× from uint8 + streaming = 32-128× total

Work Order: WO-COMP-FROM-CDNA-STREAM
"""

from __future__ import annotations

import hashlib
import json
import struct
import zlib
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, Generator, Literal, Optional, Tuple

import numpy as np

if TYPE_CHECKING:
    from helix_cdc.cdr_state_manager import CDRStateManager


# CDR trace constants
CDR_MSG_PREFIX = b"CDNA"  # 4 bytes
CDR_MSG_SIZE = 64  # Fixed message size for CDR tape


# HXZ-CDNA format constants
CDNA_MAGIC = b"HXZC"  # HXZ-CDNA magic bytes
CDNA_VERSION = 1
CODEBOOK_SIZE = 256


def _build_cdr_trace_msg(
    layer_id: str,
    start_row: int,
    block_rows: int,
    x_sample: bytes,
    indices_sample: bytes,
    codebook_sample: bytes,
) -> bytes:
    """
    Build a 64-byte CDR trace message for CDNA matmul.

    Format (64 bytes total):
      [0:4]   - prefix "CDNA"
      [4:12]  - layer_id hash (8 bytes)
      [12:16] - start_row (uint32 big-endian)
      [16:18] - block_rows (uint16 big-endian)
      [18:26] - x hash sample (8 bytes)
      [26:34] - indices hash sample (8 bytes)
      [34:42] - codebook hash sample (8 bytes)
      [42:64] - reserved zeros (22 bytes)

    All hashes are truncated SHA256 for fingerprinting.
    """
    msg = bytearray(CDR_MSG_SIZE)

    # Prefix (4 bytes)
    msg[0:4] = CDR_MSG_PREFIX

    # Layer ID hash (8 bytes) - hash the string
    layer_hash = hashlib.sha256(layer_id.encode()).digest()[:8]
    msg[4:12] = layer_hash

    # Start row (4 bytes, big-endian uint32)
    msg[12:16] = struct.pack(">I", start_row & 0xFFFFFFFF)

    # Block rows (2 bytes, big-endian uint16)
    msg[16:18] = struct.pack(">H", block_rows & 0xFFFF)

    # X sample hash (8 bytes)
    x_hash = hashlib.sha256(x_sample).digest()[:8]
    msg[18:26] = x_hash

    # Indices sample hash (8 bytes)
    idx_hash = hashlib.sha256(indices_sample).digest()[:8]
    msg[26:34] = idx_hash

    # Codebook sample hash (8 bytes)
    cb_hash = hashlib.sha256(codebook_sample).digest()[:8]
    msg[34:42] = cb_hash

    # Bytes 42-64 remain zeros (reserved)

    return bytes(msg)


class CDNAStreamReader:
    """
    Streaming reader for CDNA-format HXZ seeds.

    The key property: W is NEVER decoded to float32.
    We stream uint8 indices and compute directly via bincount + dot.
    """

    def __init__(self, hxz_path: str):
        """
        Initialize CDNA stream reader.

        Args:
            hxz_path: Path to .hxz file containing CDNA data
        """
        self.hxz_path = Path(hxz_path)

        # Parse header and load codebook (small, always in RAM)
        self._parse_header()

        # Indices blob (decompressed bytes, NOT converted to float32)
        self._indices_blob: Optional[bytes] = None

    def _parse_header(self) -> None:
        """Parse CDNA HXZ header and load codebook."""
        with open(self.hxz_path, "rb") as f:
            # Read magic
            magic = f.read(4)
            if magic != CDNA_MAGIC:
                raise ValueError(f"Not a CDNA HXZ file (magic: {magic!r})")

            # Read version
            version = struct.unpack(">H", f.read(2))[0]
            if version != CDNA_VERSION:
                raise ValueError(f"Unsupported CDNA version: {version}")

            # Read header length
            header_len = struct.unpack(">H", f.read(2))[0]

            # Read JSON header
            header_json = f.read(header_len)
            self.header = json.loads(header_json.decode("utf-8"))

            # Extract metadata
            self.tensor_name = self.header.get("tensor_name", "weight")
            self.shape = tuple(self.header["shape"])  # [M, K]
            self.rows, self.cols = self.shape

            # Read codebook (always float32, 256 entries)
            codebook_bytes = f.read(CODEBOOK_SIZE * 4)
            self.codebook = np.frombuffer(codebook_bytes, dtype=np.float32).copy()

            # Record blob offset
            self.blob_offset = f.tell()

            # Check encoding
            self.encoding = self.header.get("encoding", "zlib")

    def _ensure_indices_blob(self) -> bytes:
        """Load and decompress indices blob if not already done."""
        if self._indices_blob is not None:
            return self._indices_blob

        with open(self.hxz_path, "rb") as f:
            f.seek(self.blob_offset)
            compressed = f.read()

        if self.encoding == "zlib":
            self._indices_blob = zlib.decompress(compressed)
        else:
            self._indices_blob = compressed

        return self._indices_blob

    def stream_index_rows(
        self,
        block_rows: int = 32,
    ) -> Generator[Tuple[np.ndarray, int], None, None]:
        """
        Stream index rows in blocks.

        CRITICAL: This yields uint8 indices, NOT decoded float32!

        Args:
            block_rows: Number of rows per block

        Yields:
            (index_block [B, K] uint8, start_row_idx)
        """
        blob = self._ensure_indices_blob()

        # Bytes per row (uint8, so 1 byte per element)
        bytes_per_row = self.cols

        for start_row in range(0, self.rows, block_rows):
            end_row = min(start_row + block_rows, self.rows)
            num_rows = end_row - start_row

            # Extract bytes for this block
            byte_start = start_row * bytes_per_row
            byte_end = end_row * bytes_per_row
            block_bytes = blob[byte_start:byte_end]

            # Convert to uint8 array (NOT float32!)
            indices_block = np.frombuffer(block_bytes, dtype=np.uint8)
            indices_block = indices_block.reshape(num_rows, self.cols)

            yield indices_block, start_row

    @property
    def memory_stats(self) -> Dict[str, Any]:
        """Get memory statistics."""
        # If we decoded W fully
        full_tensor_bytes = self.rows * self.cols * 4  # float32

        # What CDNA actually uses (codebook always in RAM)
        codebook_bytes = CODEBOOK_SIZE * 4  # 1 KB

        return {
            "shape": list(self.shape),
            "rows": self.rows,
            "cols": self.cols,
            "full_tensor_bytes": full_tensor_bytes,
            "codebook_bytes": codebook_bytes,
            "codebook_entries": CODEBOOK_SIZE,
        }


def cdna_matmul_row(
    indices_row: np.ndarray,
    codebook: np.ndarray,
    x: np.ndarray,
) -> float:
    """
    Compute one row of y = W @ x using CDNA (no decoding).

    This is the CORE CDC operation:
      y[row] = sum over k: W[row, k] * x[k]
             = sum over k: codebook[indices[row, k]] * x[k]
             = sum over bin b: codebook[b] * (sum of x[k] where indices[row, k] == b)

    The bincount trick:
      weighted_x = bincount(indices[row], weights=x, minlength=256)
      y[row] = codebook @ weighted_x

    W is NEVER decoded to float32!

    Args:
        indices_row: [K] uint8 - codebook indices for this row
        codebook: [256] float32 - the codebook
        x: [K] float32 - input vector

    Returns:
        Scalar y value for this row
    """
    n_bins = len(codebook)

    # bincount: for each bin b, sum x[k] where indices_row[k] == b
    weighted_x = np.bincount(
        indices_row,
        weights=x,
        minlength=n_bins
    ).astype(np.float32)

    # Dot product with codebook
    return float(np.dot(codebook, weighted_x))


def cdna_matmul_block(
    indices_block: np.ndarray,
    codebook: np.ndarray,
    x: np.ndarray,
) -> np.ndarray:
    """
    Compute block of y = W @ x using CDNA (no decoding).

    Args:
        indices_block: [B, K] uint8 - codebook indices
        codebook: [256] float32 - the codebook
        x: [K] float32 - input vector

    Returns:
        y_block: [B] float32 - output for this block
    """
    B = indices_block.shape[0]
    y_block = np.zeros(B, dtype=np.float32)

    for i in range(B):
        y_block[i] = cdna_matmul_row(indices_block[i], codebook, x)

    return y_block


def stream_cdna_matmul(
    hxz_path: str,
    x: np.ndarray,
    block_rows: int = 32,
    *,
    cdr_mgr: Optional["CDRStateManager"] = None,
    layer_id: str = "",
    cdr_granularity: Literal["none", "layer", "block"] = "layer",
    max_cdr_blocks: int = 512,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Compute y = W @ x where W is CDNA-compressed in HXZ seed.

    This is TRUE compressed-domain compute:
      - W is stored as indices [M, K] uint8 + codebook [256] float32
      - We stream uint8 indices (NOT float32 weights)
      - Compute via bincount + dot (W NEVER decoded)
      - Full float32 tensor NEVER exists

    CDR Integration (optional):
      When cdr_mgr is provided, absorbs trace messages to CDR sponge state.
      This enables audit trails and deterministic replay verification.

      NOTE: CDR traces the COMPUTATION, not the values. It does NOT make
      matmul reversible (matmul isn't bijective). Use for:
      - Audit / trace chain at matmul granularity
      - Deterministic receipt-like progression
      - Searchable checkpoints for replay

    Args:
        hxz_path: Path to CDNA .hxz file
        x: Input vector [K]
        block_rows: Number of rows per block
        cdr_mgr: Optional CDR state manager for trace integration
        layer_id: Layer identifier for CDR trace (e.g., "L12", "attn.q")
        cdr_granularity: Trace granularity:
          - "none": No CDR tracing (default overhead-free mode)
          - "layer": One absorb per layer (cheap, safe)
          - "block": One absorb per block (fine-grained, limited by max_cdr_blocks)
        max_cdr_blocks: Max block absorbs in "block" mode (prevents tape overflow)

    Returns:
        (y [M], stats)
    """
    reader = CDNAStreamReader(hxz_path)

    # Ensure x is float32 and 1D
    x = np.asarray(x, dtype=np.float32).ravel()

    if len(x) != reader.cols:
        raise ValueError(
            f"Input dimension mismatch: W is [{reader.rows}, {reader.cols}], "
            f"x has {len(x)} elements"
        )

    # CDR setup: pre-compute samples for layer-level tracing
    cdr_active = cdr_mgr is not None and cdr_granularity != "none"
    cdr_absorbs = 0

    # Sample bytes for hashing (small samples, not full tensors)
    x_sample = x[:64].tobytes() if len(x) >= 64 else x.tobytes()
    codebook_sample = reader.codebook[:64].tobytes()

    # For layer-level, accumulate a rolling hash of all blocks
    layer_rolling_hash = hashlib.sha256() if cdr_active and cdr_granularity == "layer" else None

    # Compute in blocks
    y_chunks = []
    blocks_processed = 0
    max_block_bytes = 0

    for indices_block, start_idx in reader.stream_index_rows(block_rows):
        # CDNA kernel (W never decoded!)
        y_chunk = cdna_matmul_block(indices_block, reader.codebook, x)
        y_chunks.append(y_chunk)

        blocks_processed += 1
        # Block memory: uint8 indices only!
        block_bytes = indices_block.nbytes
        max_block_bytes = max(max_block_bytes, block_bytes)

        # CDR tracing
        if cdr_active:
            # Sample indices (first 256 bytes or less)
            indices_sample = indices_block.tobytes()[:256]

            if cdr_granularity == "block" and cdr_absorbs < max_cdr_blocks:
                # Per-block absorb
                msg = _build_cdr_trace_msg(
                    layer_id=layer_id or reader.tensor_name,
                    start_row=start_idx,
                    block_rows=indices_block.shape[0],
                    x_sample=x_sample,
                    indices_sample=indices_sample,
                    codebook_sample=codebook_sample,
                )
                cdr_mgr.cdr_absorb(msg, store_message=True)
                cdr_absorbs += 1

            elif cdr_granularity == "layer" and layer_rolling_hash is not None:
                # Accumulate for layer-level absorb
                layer_rolling_hash.update(indices_sample)
                layer_rolling_hash.update(struct.pack(">I", start_idx))

    # Layer-level CDR absorb (once at end)
    if cdr_active and cdr_granularity == "layer" and layer_rolling_hash is not None:
        # Build layer summary message
        msg = _build_cdr_trace_msg(
            layer_id=layer_id or reader.tensor_name,
            start_row=0,
            block_rows=reader.rows,  # Total rows
            x_sample=x_sample,
            indices_sample=layer_rolling_hash.digest()[:256],  # Rolling hash as "indices sample"
            codebook_sample=codebook_sample,
        )
        cdr_mgr.cdr_absorb(msg, store_message=True)
        cdr_absorbs += 1

    y = np.concatenate(y_chunks)

    # Memory stats
    mem = reader.memory_stats
    full_tensor_bytes = mem["full_tensor_bytes"]

    # What we actually touched in float32:
    # - codebook: 256 * 4 = 1024 bytes
    # - x: K * 4 bytes
    # - y: M * 4 bytes (output, unavoidable)
    float32_touched = CODEBOOK_SIZE * 4 + len(x) * 4 + len(y) * 4

    # What dense would touch: M * K * 4 (the full W matrix)
    float32_dense = reader.rows * reader.cols * 4

    stats = {
        "tensor_name": reader.tensor_name,
        "tensor_shape": list(reader.shape),
        "block_rows": block_rows,
        "blocks_processed": blocks_processed,
        "max_block_bytes_uint8": max_block_bytes,
        "codebook_bytes": CODEBOOK_SIZE * 4,
        "full_tensor_bytes_if_decoded": full_tensor_bytes,
        "float32_touched_cdna": float32_touched,
        "float32_touched_dense": float32_dense,
        "memory_reduction_factor": float32_dense / float32_touched,
        "w_decoded_to_float32": False,  # THE KEY CLAIM
        # CDR trace stats
        "cdr_enabled": cdr_active,
        "cdr_granularity": cdr_granularity if cdr_active else "none",
        "cdr_absorbs": cdr_absorbs,
    }

    return y, stats


def stream_cdna_matmul_batch(
    hxz_path: str,
    x: np.ndarray,
    block_rows: int = 32,
) -> Tuple[np.ndarray, Dict[str, Any]]:
    """
    Batch version of stream_cdna_matmul.

    Accepts a matrix of inputs shaped [K, B] (or [B, K] which will be transposed),
    streams the CDNA shard in row blocks, and returns [M, B] without ever
    materializing W as float32.
    """
    reader = CDNAStreamReader(hxz_path)

    x_arr = np.asarray(x, dtype=np.float32)
    if x_arr.ndim == 1:
        x_arr = x_arr.reshape(-1, 1)
    elif x_arr.shape[0] != reader.cols and x_arr.shape[1] == reader.cols:
        x_arr = x_arr.T

    if x_arr.shape[0] != reader.cols:
        raise ValueError(
            f"Input dimension mismatch: W is [{reader.rows}, {reader.cols}], "
            f"x has {x_arr.shape[0]} elements on the feature axis"
        )

    batch_size = x_arr.shape[1]
    y = np.zeros((reader.rows, batch_size), dtype=np.float32)

    blocks_processed = 0
    max_block_bytes = 0

    for indices_block, start_idx in reader.stream_index_rows(block_rows):
        rows_in_block = indices_block.shape[0]
        for b in range(batch_size):
            y_block = cdna_matmul_block(indices_block, reader.codebook, x_arr[:, b])
            y[start_idx : start_idx + rows_in_block, b] = y_block

        blocks_processed += 1
        max_block_bytes = max(max_block_bytes, indices_block.nbytes)

    float32_touched = (
        reader.codebook.nbytes + x_arr.size * 4 + y.size * 4
    )
    float32_dense = reader.rows * reader.cols * 4 * batch_size

    stats = {
        "tensor_name": reader.tensor_name,
        "tensor_shape": list(reader.shape),
        "batch_size": batch_size,
        "block_rows": block_rows,
        "blocks_processed": blocks_processed,
        "max_block_bytes_uint8": max_block_bytes,
        "codebook_bytes": reader.codebook.nbytes,
        "full_tensor_bytes_if_decoded": reader.rows * reader.cols * 4,
        "float32_touched_cdna": float32_touched,
        "float32_touched_dense": float32_dense,
        "memory_reduction_factor": float32_dense / float32_touched,
        "w_decoded_to_float32": False,
    }

    return y, stats


def vq_linear_forward(
    indices: np.ndarray,
    codebook: np.ndarray,
    x: np.ndarray,
) -> np.ndarray:
    """
    Compute y = W @ x directly from VQ indices + codebook.

    Args:
        indices: [out_dim, in_dim] uint8 codebook indices
        codebook: [K] float32 codebook values
        x: [in_dim] or [in_dim, batch] input matrix (can also be [batch, in_dim])

    Returns:
        y with shape [out_dim] (if x 1D) or [out_dim, batch]
    """
    indices_arr = np.asarray(indices, dtype=np.uint8)
    codebook_arr = np.asarray(codebook, dtype=np.float32)
    x_arr = np.asarray(x, dtype=np.float32)

    out_dim, in_dim = indices_arr.shape

    if x_arr.ndim == 1:
        x_arr = x_arr.reshape(in_dim, 1)
    elif x_arr.shape[0] != in_dim and x_arr.shape[1] == in_dim:
        x_arr = x_arr.T

    if x_arr.shape[0] != in_dim:
        raise ValueError(
            f"Input dimension mismatch: indices expect {in_dim} features, got {x_arr.shape[0]}"
        )

    batch_size = x_arr.shape[1]
    y = np.zeros((out_dim, batch_size), dtype=np.float32)

    for b in range(batch_size):
        y[:, b] = cdna_matmul_block(indices_arr, codebook_arr, x_arr[:, b])

    return y[:, 0] if batch_size == 1 else y


def create_cdna_hxz(
    indices: np.ndarray,
    codebook: np.ndarray,
    tensor_name: str,
    output_path: str,
    extra_meta: Optional[Dict] = None,
) -> Dict[str, Any]:
    """
    Create CDNA HXZ seed from indices and codebook.

    Args:
        indices: [M, K] uint8 - codebook indices
        codebook: [256] float32 - the codebook
        tensor_name: Name of the tensor
        output_path: Where to write .hxz file
        extra_meta: Additional metadata

    Returns:
        Receipt with creation stats
    """
    import datetime

    M, K = indices.shape
    assert indices.dtype == np.uint8
    assert len(codebook) == CODEBOOK_SIZE
    assert codebook.dtype == np.float32

    # Build header
    header = {
        "schema": "cdna_hxz_v1",
        "tensor_name": tensor_name,
        "shape": [M, K],
        "dtype": "uint8",  # indices are uint8
        "codebook_size": CODEBOOK_SIZE,
        "encoding": "zlib",
        "created_utc": datetime.datetime.utcnow().isoformat() + "Z",
    }
    if extra_meta:
        header.update(extra_meta)

    header_bytes = json.dumps(header, separators=(",", ":")).encode("utf-8")

    # Compress indices
    indices_bytes = indices.tobytes()
    compressed = zlib.compress(indices_bytes, level=9)

    # Build file
    output = bytearray()
    output.extend(CDNA_MAGIC)
    output.extend(struct.pack(">H", CDNA_VERSION))
    output.extend(struct.pack(">H", len(header_bytes)))
    output.extend(header_bytes)
    output.extend(codebook.tobytes())  # Always uncompressed (only 1KB)
    output.extend(compressed)

    # Write
    Path(output_path).write_bytes(bytes(output))

    # Receipt
    receipt = {
        "schema": "create_cdna_hxz_v1",
        "tensor_name": tensor_name,
        "shape": [M, K],
        "indices_bytes": len(indices_bytes),
        "codebook_bytes": CODEBOOK_SIZE * 4,
        "compressed_bytes": len(compressed),
        "total_bytes": len(output),
        "compression_ratio": len(indices_bytes) / len(compressed),
        "output_path": str(output_path),
    }

    return receipt


def quantize_tensor_to_cdna(
    W: np.ndarray,
    n_clusters: int = 256,
    seed: int = 42,
) -> Tuple[np.ndarray, np.ndarray, Dict[str, Any]]:
    """
    Quantize float32 tensor to CDNA format (indices + codebook).

    Uses k-means to find 256 representative values.

    Args:
        W: [M, K] float32 tensor
        n_clusters: Number of codebook entries (default 256)
        seed: Random seed for k-means

    Returns:
        (indices [M, K] uint8, codebook [256] float32, stats)
    """
    from sklearn.cluster import MiniBatchKMeans

    M, K = W.shape

    # Flatten for k-means
    W_flat = W.reshape(-1, 1).astype(np.float32)

    # Run k-means
    kmeans = MiniBatchKMeans(
        n_clusters=n_clusters,
        random_state=seed,
        batch_size=min(10000, len(W_flat)),
        n_init=3,
        max_iter=100,
    )

    indices_flat = kmeans.fit_predict(W_flat).astype(np.uint8)
    codebook = kmeans.cluster_centers_.flatten().astype(np.float32)

    indices = indices_flat.reshape(M, K)

    # Quantization error
    W_reconstructed = codebook[indices]
    mse = float(np.mean((W - W_reconstructed) ** 2))
    max_error = float(np.max(np.abs(W - W_reconstructed)))

    stats = {
        "shape": [M, K],
        "n_clusters": n_clusters,
        "quantization_mse": mse,
        "quantization_max_error": max_error,
        "compression_ratio": (M * K * 4) / (M * K + CODEBOOK_SIZE * 4),
    }

    return indices, codebook, stats


def inspect_cdna_hxz_header(hxz_path: str) -> Dict[str, Any]:
    """
    Inspect CDNA HXZ header without loading indices.

    Useful for building manifests and checking file validity.

    Args:
        hxz_path: Path to .cdna.hxz file

    Returns:
        Dict with tensor_name, shape, dtype, codebook_size, encoding, etc.
    """
    import os

    path = Path(hxz_path)
    file_size = path.stat().st_size

    with open(path, "rb") as f:
        # Read magic
        magic = f.read(4)
        if magic != CDNA_MAGIC:
            return {
                "valid": False,
                "error": f"Not a CDNA HXZ file (magic: {magic!r})",
                "path": str(path),
            }

        # Read version
        version = struct.unpack(">H", f.read(2))[0]

        # Read header length
        header_len = struct.unpack(">H", f.read(2))[0]

        # Read JSON header
        header_json = f.read(header_len)
        header = json.loads(header_json.decode("utf-8"))

        # Read codebook stats (256 * 4 = 1024 bytes)
        codebook_bytes = f.read(CODEBOOK_SIZE * 4)
        codebook = np.frombuffer(codebook_bytes, dtype=np.float32)

        # Blob offset
        blob_offset = f.tell()

        # Compressed blob size
        compressed_size = file_size - blob_offset

    # Extract shape
    shape = header.get("shape", [0, 0])
    M, K = shape if len(shape) == 2 else (0, 0)

    # Memory calculations - TWO different ratios:
    # 1. Representation ratio: dense storage vs CDNA storage (indices + codebook)
    # 2. Runtime working set: measured during actual matmul (see receipts)
    full_tensor_bytes = M * K * 4  # float32 dense
    cdna_rep_bytes = M * K * 1 + CODEBOOK_SIZE * 4  # uint8 indices + float32 codebook
    rep_memory_reduction = full_tensor_bytes / cdna_rep_bytes if cdna_rep_bytes > 0 else 0

    return {
        "valid": True,
        "path": str(path),
        "file_size_bytes": file_size,
        "version": version,
        "tensor_name": header.get("tensor_name", "unknown"),
        "shape": shape,
        "rows": M,
        "cols": K,
        "indices_dtype": "uint8",
        "codebook_dtype": "float32",
        "codebook_size": CODEBOOK_SIZE,
        "codebook_min": float(codebook.min()),
        "codebook_max": float(codebook.max()),
        "codebook_mean": float(codebook.mean()),
        "encoding": header.get("encoding", "raw"),
        "compressed_blob_bytes": compressed_size,
        "full_tensor_bytes": full_tensor_bytes,
        "cdna_rep_bytes": cdna_rep_bytes,
        "rep_memory_reduction": round(rep_memory_reduction, 2),  # ~4x (storage ratio, not runtime)
        "header": header,
    }


# CLI for testing
if __name__ == "__main__":
    import sys

    print("=== CDNA Stream Test ===")
    print()
    print("This module provides TRUE compressed-domain compute:")
    print("  - W is stored as indices [M,K] uint8 + codebook [256] float32")
    print("  - Stream uint8 indices (NOT float32 weights)")
    print("  - Compute via bincount + dot (W NEVER decoded)")
    print()
    print("Usage:")
    print("  from helix_cdc.regrow.cdna_stream import stream_cdna_matmul")
    print("  y, stats = stream_cdna_matmul('weight.cdna.hxz', x)")
    print()
    print("  # stats['w_decoded_to_float32'] == False  # THE KEY CLAIM")
