#!/usr/bin/env python3
"""
helixcode_v0.py - HelixCode v0: "Decode = Compute" Format

WO-HELIXCODE-01: A weight codec where the stored form is an executable factor/page
program, and `X @ W` is computed by interpreting that program—never materializing
dense W.

The Core Insight:
  "If you can't do random access at the granularity you compute on,
   you're not in the substrate—you're in storage."

HelixCode v0 = Raw page-indexed VQ with NO entropy codec in hot path.

Storage Format:
  tensor_name/
  ├── manifest.json       # Metadata, shapes, hashes
  ├── codebook.bin        # [256] float32 (1024 bytes)
  └── indices.bin         # [M, K] uint8 raw (mmap-able)

vs CDNA:
  - CDNA: brotli-compressed indices → decompress → compute (hot path includes decompression)
  - HelixCode: raw indices → mmap → compute (NO decompression in hot path)

Trade-off:
  - File size: ~4x larger than brotli-compressed CDNA
  - Compute: Zero decompression overhead in hot path

Key Claims (to be proven with receipts):
  - w_materialized: false (never build full W as fp32)
  - entropy_decode_in_hot_path: false (no brotli/zstd/lz4)
"""

from __future__ import annotations

import hashlib
import json
import mmap
import os
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import numpy as np


# =============================================================================
# Dataclass for matmul statistics
# =============================================================================

@dataclass
class MatmulStats:
    """Statistics from matmul operation for proof of 'no dense W'."""
    max_scratch_bytes: int
    w_dense_bytes: int
    scratch_ratio: float
    scratch_limit_ratio: float
    scratch_limit_passed: bool
    block_rows: int
    kernel_used: str = "standard"
    k_tile: Optional[int] = None


# =============================================================================
# Format Constants
# =============================================================================

HELIXCODE_V0_SCHEMA = "helixcode_v0"
CODEBOOK_SIZE = 256


# =============================================================================
# HelixCode v0 Format
# =============================================================================

@dataclass
class HelixCodeV0:
    """
    Executable weight codec - decode IS compute.

    The stored form (raw indices + codebook) is directly interpretable
    for matmul without any decompression step.
    """

    codebook: np.ndarray          # [256] float32
    indices_path: Path            # Path to raw .bin file
    shape: Tuple[int, int]        # (rows, cols) of original W
    tensor_name: str
    manifest_path: Path

    # Layout contract fields (WO-HELIXCODE-BLOCK-02)
    storage_layout: Optional[str] = None      # "out_in" or "in_out"
    expected_matmul: Optional[str] = None     # "X@W.T" or "X@W"
    source_layout: Optional[str] = None       # "gguf", "cdna"
    transposed_during_conversion: Optional[bool] = None

    # Cached mmap (lazy)
    _mmap: Optional[mmap.mmap] = field(default=None, repr=False)
    _indices_view: Optional[np.ndarray] = field(default=None, repr=False)

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

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

    @property
    def has_layout_contract(self) -> bool:
        """Check if this HelixCode has explicit layout contract fields."""
        return self.expected_matmul is not None

    def get_transpose_from_manifest(self) -> Optional[bool]:
        """
        Get transpose decision from manifest layout contract.

        Returns:
            True if should use X @ W.T
            False if should use X @ W
            None if no layout contract (must infer)
        """
        if self.expected_matmul == "X@W.T":
            return True
        elif self.expected_matmul == "X@W":
            return False
        else:
            return None

    # -------------------------------------------------------------------------
    # Loading
    # -------------------------------------------------------------------------

    @classmethod
    def load(cls, path: Union[str, Path], validate: bool = True) -> 'HelixCodeV0':
        """
        Load HelixCode v0 from directory.

        Args:
            path: Directory containing manifest.json, codebook.bin, indices.bin
            validate: If True (default), run entropy tripwire check on indices

        Returns:
            HelixCodeV0 instance (indices not loaded until first access)

        Raises:
            ValueError: If schema doesn't match or validation fails
        """
        path = Path(path)
        manifest_path = path / "manifest.json"

        with open(manifest_path) as f:
            manifest = json.load(f)

        # Validate schema
        if manifest.get("schema") != HELIXCODE_V0_SCHEMA:
            raise ValueError(f"Not a HelixCode v0 format: {manifest.get('schema')}")

        # Load codebook
        codebook_path = path / "codebook.bin"
        codebook = np.fromfile(codebook_path, dtype=np.float32)
        assert len(codebook) == CODEBOOK_SIZE, f"Expected 256 codebook entries, got {len(codebook)}"

        # Indices path (loaded lazily via mmap)
        indices_path = path / "indices.bin"

        instance = cls(
            codebook=codebook,
            indices_path=indices_path,
            shape=tuple(manifest["shape"]),
            tensor_name=manifest["tensor_name"],
            manifest_path=manifest_path,
            # Layout contract fields (WO-HELIXCODE-BLOCK-02)
            storage_layout=manifest.get("storage_layout"),
            expected_matmul=manifest.get("expected_matmul"),
            source_layout=manifest.get("source_layout"),
            transposed_during_conversion=manifest.get("transposed_during_conversion"),
        )

        # Run entropy tripwire if validation requested
        if validate:
            instance._tripwire_indices()

        return instance

    # -------------------------------------------------------------------------
    # Entropy/Degeneracy Tripwire (Phase 2)
    # -------------------------------------------------------------------------

    def _tripwire_indices(
        self,
        min_unique: int = 8,
        max_dominance: float = 0.95,
        sample_n: int = 8192,
        raise_on_fail: bool = True,
    ) -> bool:
        """
        Sample-based entropy check on indices.

        Catches:
        - All-zero indices (degenerate)
        - Truncated files that still parse
        - Indices with suspiciously low entropy (e.g., all same value)
        - Single-symbol collapse (one value dominates >95% of sample)

        Args:
            min_unique: Minimum unique values required in sample (default 8).
                        Lower than 50 because structured tensors can have
                        localized patterns that appear low-entropy in strided samples.
            max_dominance: Maximum fraction any single value can occupy (default 0.95).
                           Catches "mostly zeros" and similar collapse patterns.
            sample_n: Number of samples to check (default 8192)
            raise_on_fail: If True, raise ValueError on failure

        Returns:
            True if sane, raises ValueError if not (when raise_on_fail=True).
        """
        indices = self._ensure_mmap()
        flat = indices.ravel()

        if flat.size == 0:
            if raise_on_fail:
                raise ValueError(f"Empty indices in {self.tensor_name}")
            return False

        # Check for file size mismatch (truncation detection)
        expected_size = self.rows * self.cols
        if flat.size != expected_size:
            if raise_on_fail:
                raise ValueError(
                    f"Size mismatch in {self.tensor_name}: "
                    f"expected {expected_size} indices, got {flat.size}"
                )
            return False

        # Random sample for better coverage (avoid strided sampling artifacts)
        # Use multiple random regions instead of uniform stride
        rng = np.random.default_rng(42)  # Deterministic for reproducibility
        sample_indices = rng.choice(flat.size, min(sample_n, flat.size), replace=False)
        sample = flat[sample_indices]

        # Check 1: Minimum unique values
        unique_vals, counts = np.unique(sample, return_counts=True)
        unique_count = len(unique_vals)

        if unique_count < min_unique:
            if raise_on_fail:
                raise ValueError(
                    f"Degenerate indices in {self.tensor_name}: "
                    f"only {unique_count} unique values in sample (min={min_unique})"
                )
            return False

        # Check 2: Max-frequency dominance (catches "mostly zeros" collapse)
        top_frequency = counts.max() / len(sample)
        if top_frequency > max_dominance:
            dominant_val = unique_vals[counts.argmax()]
            if raise_on_fail:
                raise ValueError(
                    f"Single-symbol collapse in {self.tensor_name}: "
                    f"value {dominant_val} is {top_frequency:.1%} of sample (max={max_dominance:.0%})"
                )
            return False

        # Check 3: Codebook/indices sanity (catch wrong/corrupted codebook)
        # Sample a small tile, gather W, check variance and NaN
        tile_size = min(64, self.rows, self.cols)
        idx_tile = indices[:tile_size, :tile_size]
        W_tile = self.codebook[idx_tile]

        w_std = np.std(W_tile)
        has_nan = np.isnan(W_tile).any()
        has_inf = np.isinf(W_tile).any()

        if w_std < 1e-10 or has_nan or has_inf:
            if raise_on_fail:
                raise ValueError(
                    f"Codebook/indices mismatch in {self.tensor_name}: "
                    f"W_tile std={w_std:.2e}, has_nan={has_nan}, has_inf={has_inf}"
                )
            return False

        return True

    # -------------------------------------------------------------------------
    # Memory-mapped Index Access
    # -------------------------------------------------------------------------

    def _ensure_mmap(self) -> np.ndarray:
        """Ensure indices are memory-mapped. Returns view into mmap."""
        if self._indices_view is not None:
            return self._indices_view

        # Open file and create mmap
        fd = os.open(str(self.indices_path), os.O_RDONLY)
        try:
            self._mmap = mmap.mmap(fd, 0, access=mmap.ACCESS_READ)
        finally:
            os.close(fd)

        # Create numpy view (no copy!)
        self._indices_view = np.frombuffer(self._mmap, dtype=np.uint8).reshape(self.shape)
        return self._indices_view

    def get_indices(self) -> np.ndarray:
        """Get full indices array (mmap view, no copy)."""
        return self._ensure_mmap()

    def get_page(self, start_row: int, end_row: int) -> np.ndarray:
        """
        Get a page (row slice) of indices.

        This is a view into the mmap, NOT a copy.
        """
        indices = self._ensure_mmap()
        return indices[start_row:end_row]

    def close(self) -> None:
        """Close mmap if open."""
        # Clear the numpy view first to release the buffer
        self._indices_view = None
        if self._mmap is not None:
            try:
                self._mmap.close()
            except BufferError:
                # Buffer still has references, let GC handle it
                pass
            self._mmap = None

    def __del__(self):
        # Don't call close() in __del__ - it can cause issues with GC order
        pass

    # -------------------------------------------------------------------------
    # Context Manager Protocol
    # -------------------------------------------------------------------------

    def __enter__(self) -> 'HelixCodeV0':
        """Support context manager protocol."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
        """Close mmap on context exit."""
        self.close()
        return False  # Don't suppress exceptions

    # -------------------------------------------------------------------------
    # Matmul Kernels
    # -------------------------------------------------------------------------

    def matmul(
        self,
        X: np.ndarray,
        transpose_w: Optional[bool] = None,
        block_rows: int = 256,
        track_stats: bool = False,
    ) -> Union[np.ndarray, Tuple[np.ndarray, MatmulStats]]:
        """
        Compute Y = X @ W (or X @ W.T) without materializing dense W.

        This is the core "decode = compute" operation.

        Args:
            X: Input activations. Shape [K] (1D), [seq, K] (2D), or [batch, seq, K] (3D).
            transpose_w: If True, compute Y = X @ W.T (PyTorch convention).
                         If False, compute Y = X @ W.
                         If None (default), auto-resolve from manifest's expected_matmul,
                         or fall back to shape inference (transpose if cols == X last dim).
            block_rows: Process this many rows at a time (memory/speed tradeoff).
            track_stats: If True, return (Y, MatmulStats) tuple with scratch tracking.

        Returns:
            Y with shape [..., d_out] if track_stats=False
            (Y, MatmulStats) if track_stats=True
        """
        X = np.asarray(X, dtype=np.float32)
        original_shape = X.shape

        # Normalize to 2D [N, K]
        if X.ndim == 1:
            X = X.reshape(1, -1)
            squeeze_output = True
        elif X.ndim == 2:
            squeeze_output = False
        elif X.ndim == 3:
            batch, seq, K = X.shape
            X = X.reshape(batch * seq, K)
            squeeze_output = False
        else:
            raise ValueError(f"X must be 1D, 2D, or 3D, got shape {original_shape}")

        N, K_x = X.shape

        # Auto-resolve transpose from manifest if not specified
        if transpose_w is None:
            transpose_w = self.get_transpose_from_manifest()
            if transpose_w is None:
                # Fallback to shape inference: transpose if cols match X last dim
                transpose_w = (self.cols == K_x)

        # Validate dimensions
        if transpose_w:
            # Y = X @ W.T where W is [d_out, d_in]
            # X has [N, d_in], W.T is [d_in, d_out], Y is [N, d_out]
            if K_x != self.cols:
                raise ValueError(f"X last dim ({K_x}) must match W cols ({self.cols}) for X @ W.T")
            d_out = self.rows
        else:
            # Y = X @ W where W is [d_out, d_in]
            # This would require X [N, d_out] @ W [d_out, d_in] = Y [N, d_in]
            # Less common, but supported
            if K_x != self.rows:
                raise ValueError(f"X last dim ({K_x}) must match W rows ({self.rows}) for X @ W")
            d_out = self.cols

        # Determine kernel selection: use K-tiled kernel for large d_in
        L2_CACHE_BUDGET = 2 * 1024 * 1024  # 2MB
        scratch_estimate = block_rows * self.cols * 4  # float32
        use_ktiled = scratch_estimate > L2_CACHE_BUDGET

        # Compute using gather kernel
        if transpose_w:
            if use_ktiled and track_stats:
                Y, max_scratch_bytes, kernel_used, k_tile = self._matmul_gather_transpose_ktiled_tracked(X, block_rows)
            elif use_ktiled:
                Y = self._matmul_gather_transpose_ktiled(X, block_rows)
                max_scratch_bytes = None
                kernel_used = "ktiled"
                k_tile = 512
            elif track_stats:
                Y, max_scratch_bytes = self._matmul_gather_transpose_tracked(X, block_rows)
                kernel_used = "standard"
                k_tile = None
            else:
                Y = self._matmul_gather_transpose(X, block_rows)
                max_scratch_bytes = None
                kernel_used = "standard"
                k_tile = None
        else:
            if track_stats:
                Y, max_scratch_bytes = self._matmul_gather_direct_tracked(X, block_rows)
                kernel_used = "direct"
                k_tile = None
            else:
                Y = self._matmul_gather_direct(X, block_rows)
                max_scratch_bytes = None
                kernel_used = "direct"
                k_tile = None

        # Restore original batch dimensions
        if squeeze_output:
            Y = Y.squeeze(0)
        elif len(original_shape) == 3:
            Y = Y.reshape(original_shape[0], original_shape[1], -1)

        if track_stats:
            w_dense_bytes = self.rows * self.cols * 4
            scratch_ratio = max_scratch_bytes / w_dense_bytes if max_scratch_bytes else 0.0
            stats = MatmulStats(
                max_scratch_bytes=max_scratch_bytes or 0,
                w_dense_bytes=w_dense_bytes,
                scratch_ratio=scratch_ratio,
                scratch_limit_ratio=0.25,
                scratch_limit_passed=scratch_ratio <= 0.25,
                block_rows=block_rows,
                kernel_used=kernel_used,
                k_tile=k_tile,
            )
            return Y, stats

        return Y

    def matmul_with_stats(
        self,
        X: np.ndarray,
        transpose_w: bool = True,
        block_rows: int = 256,
        scratch_limit_ratio: float = 0.25,
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Compute Y = X @ W with scratch tracking for proof of "no dense W".

        Returns:
            (Y, stats) where stats includes max_scratch_bytes, scratch_ratio, etc.
        """
        X = np.asarray(X, dtype=np.float32)
        original_shape = X.shape

        # Normalize to 2D [N, K]
        if X.ndim == 1:
            X = X.reshape(1, -1)
            squeeze_output = True
        elif X.ndim == 2:
            squeeze_output = False
        elif X.ndim == 3:
            batch, seq, K = X.shape
            X = X.reshape(batch * seq, K)
            squeeze_output = False
        else:
            raise ValueError(f"X must be 1D, 2D, or 3D, got shape {original_shape}")

        N, K_x = X.shape

        # Validate dimensions
        if transpose_w:
            if K_x != self.cols:
                raise ValueError(f"X last dim ({K_x}) must match W cols ({self.cols}) for X @ W.T")
            d_out = self.rows
        else:
            if K_x != self.rows:
                raise ValueError(f"X last dim ({K_x}) must match W rows ({self.rows}) for X @ W")
            d_out = self.cols

        # Calculate W dense size for comparison
        w_dense_bytes = self.rows * self.cols * 4  # float32

        # Run matmul with scratch tracking
        if transpose_w:
            Y, max_scratch_bytes = self._matmul_gather_transpose_tracked(X, block_rows)
        else:
            Y, max_scratch_bytes = self._matmul_gather_direct_tracked(X, block_rows)

        # Restore original batch dimensions
        if squeeze_output:
            Y = Y.squeeze(0)
        elif len(original_shape) == 3:
            Y = Y.reshape(original_shape[0], original_shape[1], -1)

        # Compute scratch ratio
        scratch_ratio = max_scratch_bytes / w_dense_bytes

        # Build stats
        stats = {
            "max_scratch_bytes": max_scratch_bytes,
            "w_dense_bytes": w_dense_bytes,
            "scratch_ratio": scratch_ratio,
            "scratch_limit_ratio": scratch_limit_ratio,
            "scratch_limit_passed": scratch_ratio <= scratch_limit_ratio,
            "block_rows": block_rows,
        }

        # Enforce guardrail
        if scratch_ratio > scratch_limit_ratio:
            raise RuntimeError(
                f"Scratch exceeds limit: {scratch_ratio:.2%} > {scratch_limit_ratio:.2%}. "
                f"Max scratch: {max_scratch_bytes:,} bytes, W dense: {w_dense_bytes:,} bytes"
            )

        return Y, stats

    def _matmul_gather_transpose(
        self,
        X: np.ndarray,
        block_rows: int,
    ) -> np.ndarray:
        """Compute Y = X @ W.T using gather (without tracking)."""
        Y, _ = self._matmul_gather_transpose_tracked(X, block_rows)
        return Y

    def _matmul_gather_transpose_tracked(
        self,
        X: np.ndarray,
        block_rows: int,
    ) -> Tuple[np.ndarray, int]:
        """
        Compute Y = X @ W.T using gather, tracking max scratch bytes.

        W is [d_out, d_in], stored as indices [d_out, d_in].
        W.T is [d_in, d_out].

        Y[n, i] = sum_j X[n, j] * W[i, j]
                = sum_j X[n, j] * codebook[indices[i, j]]

        Using gather: for each output row i, gather codebook[indices[i, :]] → W_row[d_in]
        Then: Y[:, i] = X @ W_row

        This processes in blocks of output rows for memory efficiency.

        Returns:
            (Y, max_scratch_bytes)
        """
        indices = self._ensure_mmap()
        codebook = self.codebook

        N = X.shape[0]
        d_out = self.rows

        Y = np.zeros((N, d_out), dtype=np.float32)
        max_scratch_bytes = 0

        # Process blocks of output rows
        for start_row in range(0, d_out, block_rows):
            end_row = min(start_row + block_rows, d_out)

            # Get index page for this block [block_size, d_in]
            idx_page = indices[start_row:end_row]

            # Gather: W_block[i, j] = codebook[idx_page[i, j]]
            # This is the KEY operation - codebook[indices] does the gather
            W_block = codebook[idx_page]  # [block_size, d_in] float32

            # Track scratch size
            max_scratch_bytes = max(max_scratch_bytes, W_block.nbytes)

            # Matmul: Y[:, start:end] = X @ W_block.T
            Y[:, start_row:end_row] = X @ W_block.T

        return Y, max_scratch_bytes

    def _matmul_gather_direct(
        self,
        X: np.ndarray,
        block_rows: int,
    ) -> np.ndarray:
        """Compute Y = X @ W using gather (without tracking)."""
        Y, _ = self._matmul_gather_direct_tracked(X, block_rows)
        return Y

    def _matmul_gather_direct_tracked(
        self,
        X: np.ndarray,
        block_rows: int,
    ) -> Tuple[np.ndarray, int]:
        """
        Compute Y = X @ W using gather (less common case), with scratch tracking.

        W is [d_out, d_in], X is [N, d_out], Y is [N, d_in].

        Returns:
            (Y, max_scratch_bytes)
        """
        indices = self._ensure_mmap()
        codebook = self.codebook

        N = X.shape[0]
        d_in = self.cols

        Y = np.zeros((N, d_in), dtype=np.float32)
        max_scratch_bytes = 0

        # Process blocks of W rows (which are input rows for this operation)
        for start_row in range(0, self.rows, block_rows):
            end_row = min(start_row + block_rows, self.rows)

            # Get index page [block_size, d_in]
            idx_page = indices[start_row:end_row]

            # Gather: W_block[i, j] = codebook[idx_page[i, j]]
            W_block = codebook[idx_page]  # [block_size, d_in]

            # Track scratch size
            max_scratch_bytes = max(max_scratch_bytes, W_block.nbytes)

            # For Y = X @ W: Y[n, j] += sum_i X[n, i] * W[i, j]
            # X[:, start:end] @ W_block
            Y += X[:, start_row:end_row] @ W_block

        return Y, max_scratch_bytes

    # -------------------------------------------------------------------------
    # K-Tiled Kernels (Phase 3: Cache-Friendly for Large d_in)
    # -------------------------------------------------------------------------

    def _matmul_gather_transpose_ktiled(
        self,
        X: np.ndarray,
        block_rows: int = 256,
        k_tile: int = 512,
    ) -> np.ndarray:
        """Compute Y = X @ W.T using K-tiled gather (without tracking)."""
        Y, _, _, _ = self._matmul_gather_transpose_ktiled_tracked(X, block_rows, k_tile)
        return Y

    def _matmul_gather_transpose_ktiled_tracked(
        self,
        X: np.ndarray,
        block_rows: int = 256,
        k_tile: int = 512,
    ) -> Tuple[np.ndarray, int, str, int]:
        """
        K-tiled gather matmul for cache-friendly large d_in.

        Instead of: Y[:, start:end] = X @ W_block.T
        Do:         Y[:, start:end] += X[:, k:k+k_tile] @ W_block[:, k:k+k_tile].T

        This reduces peak scratch memory from (block_rows * d_in * 4) to
        (block_rows * k_tile * 4), which is critical when d_in >= 2048.

        Args:
            X: Input [N, d_in]
            block_rows: Output rows per block
            k_tile: K-dimension tile size (controls scratch vs compute tradeoff)

        Returns:
            (Y, max_scratch_bytes, kernel_name, k_tile_used)
        """
        indices = self._ensure_mmap()
        codebook = self.codebook

        N = X.shape[0]
        d_out = self.rows
        d_in = self.cols

        Y = np.zeros((N, d_out), dtype=np.float32)
        max_scratch_bytes = 0

        # Process blocks of output rows
        for start_row in range(0, d_out, block_rows):
            end_row = min(start_row + block_rows, d_out)
            actual_block_rows = end_row - start_row

            # Get full index page for this block [actual_block_rows, d_in]
            idx_page = indices[start_row:end_row]

            # K-tile: process k-dimension in tiles
            for k_start in range(0, d_in, k_tile):
                k_end = min(k_start + k_tile, d_in)

                # Get K-tile of indices [actual_block_rows, k_tile_actual]
                idx_ktile = idx_page[:, k_start:k_end]

                # Gather: W_ktile[i, j] = codebook[idx_ktile[i, j]]
                W_ktile = codebook[idx_ktile]  # [actual_block_rows, k_tile_actual]

                # Track scratch size
                max_scratch_bytes = max(max_scratch_bytes, W_ktile.nbytes)

                # Partial matmul: Y[:, start:end] += X[:, k_start:k_end] @ W_ktile.T
                Y[:, start_row:end_row] += X[:, k_start:k_end] @ W_ktile.T

        return Y, max_scratch_bytes, "ktiled", k_tile

    # -------------------------------------------------------------------------
    # Routing Capability Checks (Phase 4: Integration with Track A)
    # -------------------------------------------------------------------------

    def gpu_available(self) -> bool:
        """
        Check if GPU is available and tensor fits in GPU memory.

        NOTE: This checks CUDA runtime, not GPU service liveness.
        Use GpuServiceClient.ping() to verify service is alive.

        Checks CuPy first (Track A worker uses CuPy), falls back to PyTorch.

        Returns:
            True if GPU is available and tensor fits in GPU memory with 80% headroom.
        """
        tensor_bytes = self.rows * self.cols * 4  # float32

        # Try CuPy first (Track A GPU worker uses CuPy)
        try:
            import cupy as cp
            if cp.cuda.runtime.getDeviceCount() > 0:
                free_mem, total_mem = cp.cuda.runtime.memGetInfo()
                return tensor_bytes < free_mem * 0.8
        except ImportError:
            pass
        except Exception:
            # CuPy installed but CUDA runtime issue
            pass

        # Fall back to PyTorch
        try:
            import torch
            if torch.cuda.is_available():
                free_mem = torch.cuda.mem_get_info()[0]
                return tensor_bytes < free_mem * 0.8
        except ImportError:
            pass

        return False

    # Backwards compatibility alias
    can_use_gpu_service = gpu_available

    def needs_transpose(self) -> bool:
        """
        Return whether this tensor needs transpose for X @ W.T.

        Returns:
            True if should use X @ W.T, False if X @ W.
            Defaults to True (GGUF convention: out_in layout).
        """
        result = self.get_transpose_from_manifest()
        if result is None:
            # Fallback: assume GGUF convention (out_in layout)
            return True
        return result

    def indices_sane(self) -> bool:
        """
        Quick entropy check without raising exception.

        Returns:
            True if indices pass entropy sanity check.
        """
        return self._tripwire_indices(raise_on_fail=False)

    def get_artifact_hashes(self) -> Dict[str, str]:
        """
        Return SHA256 hashes of underlying artifacts for provenance tracking.

        WO-VAULT-PROOF-HARNESS-01: Artifact-linked receipts.

        Returns:
            Dict with codebook_sha256, indices_sha256, manifest_sha256 (all truncated to 16 chars).
        """
        hashes = {}

        # Codebook hash
        codebook_path = self.indices_path.parent / "codebook.bin"
        if codebook_path.exists():
            hashes["codebook_sha256"] = hashlib.sha256(
                codebook_path.read_bytes()
            ).hexdigest()[:16]

        # Indices hash
        if self.indices_path.exists():
            hashes["indices_sha256"] = hashlib.sha256(
                self.indices_path.read_bytes()
            ).hexdigest()[:16]

        # Manifest hash
        if self.manifest_path.exists():
            hashes["manifest_sha256"] = hashlib.sha256(
                self.manifest_path.read_bytes()
            ).hexdigest()[:16]

        return hashes

    def fits_buffers(self, batch_size: int, block_rows: int = 256) -> bool:
        """
        Check if matmul will fit in reasonable memory.

        Args:
            batch_size: Expected batch size
            block_rows: Block rows for matmul kernel

        Returns:
            True if total buffer requirements < 4GB limit.
        """
        # X: [batch_size, d_in]
        # W_block: [block_rows, d_in]
        # Y: [batch_size, d_out]
        x_bytes = batch_size * self.cols * 4
        w_block_bytes = block_rows * self.cols * 4
        y_bytes = batch_size * self.rows * 4
        total = x_bytes + w_block_bytes + y_bytes
        return total < 4 * 1024 * 1024 * 1024  # 4GB limit


# =============================================================================
# Writer (Create HelixCode from indices + codebook)
# =============================================================================

def create_helixcode_v0(
    indices: np.ndarray,
    codebook: np.ndarray,
    output_dir: Union[str, Path],
    tensor_name: str,
    extra_meta: Optional[Dict] = None,
    # Layout contract fields (WO-HELIXCODE-BLOCK-02)
    storage_layout: Optional[str] = None,  # "out_in" or "in_out"
    expected_matmul: Optional[str] = None,  # "X@W.T" or "X@W"
    source_layout: Optional[str] = None,   # "gguf", "cdna", etc.
    transposed_during_conversion: Optional[bool] = None,
) -> Dict[str, Any]:
    """
    Create HelixCode v0 format from indices and codebook.

    Args:
        indices: [M, K] uint8 codebook indices
        codebook: [256] float32 codebook values
        output_dir: Directory to write format files
        tensor_name: Name of the tensor
        extra_meta: Additional metadata for manifest
        storage_layout: How indices are laid out - "out_in" means [d_out, d_in]
        expected_matmul: What matmul form to use - "X@W.T" or "X@W"
        source_layout: Where indices came from - "gguf", "cdna"
        transposed_during_conversion: Was tensor transposed during conversion?

    Returns:
        Receipt with creation stats
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    indices = np.asarray(indices, dtype=np.uint8)
    codebook = np.asarray(codebook, dtype=np.float32)

    if indices.ndim == 1:
        indices = indices.reshape(1, -1)

    M, K = indices.shape
    assert len(codebook) == CODEBOOK_SIZE, f"Codebook must be 256 entries, got {len(codebook)}"

    # Write codebook
    codebook_path = output_dir / "codebook.bin"
    codebook.tofile(codebook_path)

    # Write raw indices (NO compression - this is the point!)
    indices_path = output_dir / "indices.bin"
    indices.tofile(indices_path)

    # Compute hashes
    codebook_sha256 = hashlib.sha256(codebook.tobytes()).hexdigest()
    indices_sha256 = hashlib.sha256(indices.tobytes()).hexdigest()

    # Build manifest
    manifest = {
        "schema": HELIXCODE_V0_SCHEMA,
        "tensor_name": tensor_name,
        "shape": [M, K],
        "dtype": "uint8",
        "codebook_size": CODEBOOK_SIZE,
        "codebook_sha256": codebook_sha256,
        "indices_sha256": indices_sha256,
        "created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
        "compression": "none",  # THE KEY CLAIM
    }

    # Add layout contract fields (WO-HELIXCODE-BLOCK-02)
    if storage_layout is not None:
        manifest["storage_layout"] = storage_layout
    if expected_matmul is not None:
        manifest["expected_matmul"] = expected_matmul
    if source_layout is not None:
        manifest["source_layout"] = source_layout
    if transposed_during_conversion is not None:
        manifest["transposed_during_conversion"] = transposed_during_conversion

    if extra_meta:
        manifest.update(extra_meta)

    manifest_path = output_dir / "manifest.json"
    with open(manifest_path, "w") as f:
        json.dump(manifest, f, indent=2)

    # Build receipt
    receipt = {
        "schema": "helixcode_v0_creation_receipt",
        "work_order": "WO-HELIXCODE-01",
        "tensor_name": tensor_name,
        "shape": [M, K],
        "output_dir": str(output_dir),
        "files": {
            "manifest": str(manifest_path),
            "codebook": str(codebook_path),
            "indices": str(indices_path),
        },
        "sizes": {
            "codebook_bytes": codebook.nbytes,
            "indices_bytes": indices.nbytes,
            "total_bytes": codebook.nbytes + indices.nbytes,
        },
        "hashes": {
            "codebook_sha256": codebook_sha256,
            "indices_sha256": indices_sha256,
        },
        "compression": "none",  # THE KEY CLAIM
        "entropy_decode_in_hot_path": False,  # THE OTHER KEY CLAIM
        "created_utc": manifest["created_utc"],
    }

    return receipt


# =============================================================================
# Converter: CDNA to HelixCode
# =============================================================================

def convert_cdna_to_helixcode(
    cdna_path: Union[str, Path],
    output_dir: Union[str, Path],
    sidecar_path: Optional[Union[str, Path]] = None,
    # Layout contract fields (WO-HELIXCODE-BLOCK-02)
    storage_layout: Optional[str] = None,
    expected_matmul: Optional[str] = None,
    source_layout: Optional[str] = None,
    transposed_during_conversion: Optional[bool] = None,
) -> Dict[str, Any]:
    """
    Convert existing CDNA (v1 or v2) to HelixCode v0.

    This extracts indices and codebook from CDNA, stripping the brotli/zstd
    compression to create a raw mmap-able format.

    Args:
        cdna_path: Path to .cdna.hxz or .cdna2.hxz file
        output_dir: Directory to write HelixCode files
        sidecar_path: Optional path to .hxzo sidecar (stored in manifest but not applied)
        storage_layout: How indices are laid out - "out_in" means [d_out, d_in]
        expected_matmul: What matmul form to use - "X@W.T" or "X@W"
        source_layout: Where indices came from - "gguf", "cdna"
        transposed_during_conversion: Was tensor transposed during conversion?

    Returns:
        Conversion receipt
    """
    cdna_path = Path(cdna_path)
    output_dir = Path(output_dir)

    # Detect CDNA version and use appropriate reader
    is_v2 = "cdna2" in str(cdna_path)

    if is_v2:
        from helix_cdc.regrow.cdna_stream_v2 import CDNAv2Reader
        reader = CDNAv2Reader(cdna_path)

        # Extract all indices by decompressing all blocks (one-time cost)
        all_indices = []
        with open(cdna_path, "rb") as f:
            for indices_block, start_row in reader.stream_blocks(file_handle=f, verify=True):
                all_indices.append(indices_block.copy())

        indices = np.vstack(all_indices)
        codebook = reader.codebook.copy()
        source_format = "cdna2"
        source_codec = reader.codec
        source_blocks = reader.num_blocks
        source_block_rows = reader.block_rows
    else:
        # CDNA v1 format
        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            from helix_cdc.regrow.cdna_stream import CDNAStreamReader

        reader = CDNAStreamReader(cdna_path)
        indices = reader._ensure_indices_blob()
        indices = np.frombuffer(indices, dtype=np.uint8).reshape(reader.rows, reader.cols)
        codebook = reader.codebook.copy()
        source_format = "cdna"
        source_codec = "zlib"  # v1 uses zlib
        source_blocks = 1  # v1 is single block
        source_block_rows = reader.rows

    # Create HelixCode
    extra_meta = {
        "source_cdna": str(cdna_path),
        "source_cdna_sha256": hashlib.sha256(cdna_path.read_bytes()).hexdigest(),
    }
    if sidecar_path:
        extra_meta["sidecar_path"] = str(sidecar_path)

    receipt = create_helixcode_v0(
        indices=indices,
        codebook=codebook,
        output_dir=output_dir,
        tensor_name=reader.tensor_name,
        extra_meta=extra_meta,
        # Layout contract fields (WO-HELIXCODE-BLOCK-02)
        storage_layout=storage_layout,
        expected_matmul=expected_matmul,
        source_layout=source_layout,
        transposed_during_conversion=transposed_during_conversion,
    )

    # Add conversion-specific info
    receipt["conversion"] = {
        "source_format": source_format,
        "source_codec": source_codec,
        "source_blocks": source_blocks,
        "source_block_rows": source_block_rows,
    }

    return receipt
