#!/usr/bin/env python3
"""
helixcode_shared_v0.py - HelixCode Shared Codebook Format

WO-HELIXCODE-02: Multi-Tensor Shared Codebooks

The Insight:
  A single codebook can be shared across tensors with the SAME d_in (cols of W).
  Target: Q/K/V attention projections from same layer (all have cols == hidden_size).

Format:
  seeds/helixcode_shared_v0/<group_name>/
  ├── codebook.f32            # Shared [256] float32 (1024 bytes)
  ├── manifest.json           # Group metadata + per-tensor info
  └── tensors/
      ├── blk_0_attn_q_weight/
      │   └── indices.bin     # Raw mmap-able indices
      ├── blk_0_attn_k_weight/
      │   └── indices.bin
      └── blk_0_attn_v_weight/
          └── indices.bin

Key Claims:
  - codebook_loaded_count == 1 (proves sharing works)
  - Per-tensor cosine >= 0.98 (shared codebook doesn't hurt fidelity)
  - Per-tensor entropy_tripwire_passed: true
  - Per-tensor scratch_limit_passed: true
"""

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


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

HELIXCODE_SHARED_V0_SCHEMA = "helixcode_shared_v0"
CODEBOOK_SIZE = 256


# =============================================================================
# Tensor View (shares codebook with group)
# =============================================================================

@dataclass
class HelixCodeTensorView:
    """
    View into shared group for one tensor.

    The codebook is a reference to the shared codebook (NOT a copy).
    This proves codebook sharing at the memory level.
    """

    tensor_name: str
    codebook: np.ndarray          # Reference to shared codebook (NOT owned)
    indices_path: Path
    shape: Tuple[int, int]        # (d_out, d_in)

    # 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]

    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. View into mmap, NOT a copy."""
        indices = self._ensure_mmap()
        return indices[start_row:end_row]

    def close(self) -> None:
        """Close mmap if open."""
        self._indices_view = None
        if self._mmap is not None:
            try:
                self._mmap.close()
            except BufferError:
                pass
            self._mmap = None

    def matmul(
        self,
        X: np.ndarray,
        transpose_w: bool = True,
        block_rows: int = 256,
    ) -> np.ndarray:
        """
        Compute Y = X @ W (or X @ W.T) without materializing dense W.

        Uses shared codebook from group.
        """
        Y, _ = self.matmul_with_stats(X, transpose_w, block_rows)
        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
        indices = self._ensure_mmap()
        codebook = self.codebook  # Shared codebook reference

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

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

                # 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]]
                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
        else:
            # Less common case: Y = X @ W
            for start_row in range(0, self.rows, block_rows):
                end_row = min(start_row + block_rows, self.rows)

                idx_page = indices[start_row:end_row]
                W_block = codebook[idx_page]

                max_scratch_bytes = max(max_scratch_bytes, W_block.nbytes)

                Y += X[:, start_row:end_row] @ W_block

        # 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 = {
            "tensor_name": self.tensor_name,
            "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,
            "shape": list(self.shape),
        }

        # 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


# =============================================================================
# Shared Codebook Group
# =============================================================================

@dataclass
class HelixCodeSharedV0:
    """
    Shared codebook with per-tensor index pages.

    Key proof point: codebook is loaded ONCE, shared by all tensor views.
    """

    group_name: str
    group_dir: Path
    d_in: int                            # Common d_in for all tensors
    codebook: np.ndarray                 # Shared [256] float32
    codebook_sha256: str
    encoder_seed: int
    samples_per_tensor: int
    tensor_info: Dict[str, Dict[str, Any]]  # Per-tensor metadata

    # Track codebook load count for proof
    _codebook_load_count: int = field(default=1, repr=False)  # Loaded once at construction

    # Cached tensor views
    _tensor_views: Dict[str, HelixCodeTensorView] = field(default_factory=dict, repr=False)

    @property
    def tensor_names(self) -> List[str]:
        """List of tensor names in this group."""
        return list(self.tensor_info.keys())

    @property
    def codebook_load_count(self) -> int:
        """Number of times codebook was loaded (should be 1)."""
        return self._codebook_load_count

    @property
    def shared_codebook_bytes(self) -> int:
        """Size of shared codebook in bytes."""
        return self.codebook.nbytes

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

    @classmethod
    def load(cls, group_dir: Union[str, Path]) -> 'HelixCodeSharedV0':
        """
        Load shared codebook group from directory.

        Args:
            group_dir: Directory containing manifest.json, codebook.f32, tensors/

        Returns:
            HelixCodeSharedV0 instance (tensor indices loaded lazily on first access)
        """
        group_dir = Path(group_dir)
        manifest_path = group_dir / "manifest.json"

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

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

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

        # Verify codebook hash
        actual_hash = hashlib.sha256(codebook.tobytes()).hexdigest()
        expected_hash = manifest.get("codebook_sha256", "")
        if expected_hash and actual_hash != expected_hash:
            raise ValueError(f"Codebook hash mismatch: {actual_hash} != {expected_hash}")

        return cls(
            group_name=manifest["group_name"],
            group_dir=group_dir,
            d_in=manifest["d_in"],
            codebook=codebook,
            codebook_sha256=actual_hash,
            encoder_seed=manifest.get("encoder_seed", 0),
            samples_per_tensor=manifest.get("samples_per_tensor", 200000),
            tensor_info=manifest.get("tensors", {}),
            _codebook_load_count=1,
        )

    # -------------------------------------------------------------------------
    # Tensor Access
    # -------------------------------------------------------------------------

    def get_tensor(self, tensor_name: str) -> HelixCodeTensorView:
        """
        Get view for specific tensor (shares codebook).

        Tensor indices are loaded lazily on first matmul.
        """
        if tensor_name not in self.tensor_info:
            available = list(self.tensor_info.keys())
            raise KeyError(f"Tensor '{tensor_name}' not in group. Available: {available}")

        # Return cached view if exists
        if tensor_name in self._tensor_views:
            return self._tensor_views[tensor_name]

        # Create new view (shares codebook reference)
        info = self.tensor_info[tensor_name]

        # Build indices path from tensor name
        # Convert tensor name to filesystem-safe format
        safe_name = tensor_name.replace(".", "_").replace("/", "_")
        indices_path = self.group_dir / "tensors" / safe_name / "indices.bin"

        view = HelixCodeTensorView(
            tensor_name=tensor_name,
            codebook=self.codebook,  # SHARED reference, not copy
            indices_path=indices_path,
            shape=tuple(info["shape"]),
        )

        self._tensor_views[tensor_name] = view
        return view

    def matmul(
        self,
        tensor_name: str,
        X: np.ndarray,
        transpose_w: bool = True,
        block_rows: int = 256,
    ) -> np.ndarray:
        """
        Compute X @ W using shared codebook.

        Convenience method that gets tensor view and calls matmul.
        """
        view = self.get_tensor(tensor_name)
        return view.matmul(X, transpose_w, block_rows)

    def matmul_with_stats(
        self,
        tensor_name: str,
        X: np.ndarray,
        transpose_w: bool = True,
        block_rows: int = 256,
        scratch_limit_ratio: float = 0.25,
    ) -> Tuple[np.ndarray, Dict[str, Any]]:
        """
        Compute X @ W with stats using shared codebook.
        """
        view = self.get_tensor(tensor_name)
        return view.matmul_with_stats(X, transpose_w, block_rows, scratch_limit_ratio)

    def close(self) -> None:
        """Close all tensor mmaps."""
        for view in self._tensor_views.values():
            view.close()
        self._tensor_views.clear()

    def __repr__(self) -> str:
        return (
            f"HelixCodeSharedV0(group={self.group_name}, "
            f"d_in={self.d_in}, tensors={len(self.tensor_info)})"
        )


# =============================================================================
# Writer (Create shared format from indices + codebook)
# =============================================================================

def create_helixcode_shared_v0(
    codebook: np.ndarray,
    tensor_data: Dict[str, Dict[str, Any]],
    output_dir: Union[str, Path],
    group_name: str,
    d_in: int,
    encoder_seed: int = 0,
    samples_per_tensor: int = 200000,
    extra_meta: Optional[Dict] = None,
) -> Dict[str, Any]:
    """
    Create HelixCode shared v0 format.

    Args:
        codebook: [256] float32 shared codebook values
        tensor_data: Dict mapping tensor_name to {
            "indices": np.ndarray [d_out, d_in] uint8,
            "shape": (d_out, d_in),
            "source_cdna": Optional[str],  # Source path for manifest
        }
        output_dir: Directory to write format files
        group_name: Name for this tensor group (e.g., "layer0_attn")
        d_in: Common input dimension for all tensors
        encoder_seed: Seed used for encoding (for reproducibility)
        samples_per_tensor: Samples per tensor used to build codebook
        extra_meta: Additional metadata for manifest

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

    codebook = np.asarray(codebook, dtype=np.float32)
    assert len(codebook) == CODEBOOK_SIZE, f"Codebook must be 256 entries, got {len(codebook)}"

    # Write shared codebook
    codebook_path = output_dir / "codebook.f32"
    codebook.tofile(codebook_path)
    codebook_sha256 = hashlib.sha256(codebook.tobytes()).hexdigest()

    # Create tensors directory
    tensors_dir = output_dir / "tensors"
    tensors_dir.mkdir(exist_ok=True)

    # Write per-tensor indices
    tensor_manifest = {}
    total_indices_bytes = 0

    for tensor_name, data in tensor_data.items():
        indices = np.asarray(data["indices"], dtype=np.uint8)
        shape = data.get("shape", indices.shape)

        # Verify d_in matches
        if shape[1] != d_in:
            raise ValueError(f"Tensor {tensor_name} has d_in={shape[1]}, expected {d_in}")

        # Create tensor directory
        safe_name = tensor_name.replace(".", "_").replace("/", "_")
        tensor_dir = tensors_dir / safe_name
        tensor_dir.mkdir(exist_ok=True)

        # Write raw indices
        indices_path = tensor_dir / "indices.bin"
        indices.tofile(indices_path)
        indices_sha256 = hashlib.sha256(indices.tobytes()).hexdigest()

        total_indices_bytes += indices.nbytes

        tensor_manifest[tensor_name] = {
            "shape": list(shape),
            "indices_sha256": indices_sha256,
            "indices_bytes": indices.nbytes,
        }
        if data.get("source_cdna"):
            tensor_manifest[tensor_name]["source_cdna"] = data["source_cdna"]

    # Build manifest
    manifest = {
        "schema": HELIXCODE_SHARED_V0_SCHEMA,
        "group_name": group_name,
        "d_in": d_in,
        "codebook_sha256": codebook_sha256,
        "codebook_k": CODEBOOK_SIZE,
        "encoder_seed": encoder_seed,
        "samples_per_tensor": samples_per_tensor,
        "tensors": tensor_manifest,
        "created_utc": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    }
    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_shared_v0_creation_receipt",
        "work_order": "WO-HELIXCODE-02",
        "group_name": group_name,
        "d_in": d_in,
        "output_dir": str(output_dir),
        "tensor_count": len(tensor_data),
        "tensor_names": list(tensor_data.keys()),
        "files": {
            "manifest": str(manifest_path),
            "codebook": str(codebook_path),
            "tensors_dir": str(tensors_dir),
        },
        "sizes": {
            "codebook_bytes": codebook.nbytes,
            "total_indices_bytes": total_indices_bytes,
            "total_bytes": codebook.nbytes + total_indices_bytes,
        },
        "hashes": {
            "codebook_sha256": codebook_sha256,
        },
        "encoder_seed": encoder_seed,
        "samples_per_tensor": samples_per_tensor,
        "compression": "none",
        "entropy_decode_in_hot_path": False,
        "created_utc": manifest["created_utc"],
    }

    return receipt


# =============================================================================
# Entropy Tripwire (proof that gather is from real codec, not zeros)
# =============================================================================

def compute_entropy_tripwire(indices: np.ndarray) -> Dict[str, Any]:
    """
    Compute entropy metrics to prove indices are real codec output, not zeros.

    Returns:
        Dict with entropy metrics
    """
    indices_flat = indices.flatten()

    # Count unique values
    unique_values = len(np.unique(indices_flat))

    # Compute histogram
    counts = np.bincount(indices_flat, minlength=256)
    probabilities = counts / counts.sum()

    # Shannon entropy
    nonzero_probs = probabilities[probabilities > 0]
    entropy = -np.sum(nonzero_probs * np.log2(nonzero_probs))

    # Max entropy for 256 symbols
    max_entropy = np.log2(256)  # 8.0 bits

    # Normalized entropy (0 = constant, 1 = uniform)
    normalized_entropy = entropy / max_entropy

    # Tripwire thresholds:
    # - unique_values > 100 (uses significant portion of codebook)
    # - normalized_entropy > 0.5 (not degenerate)
    tripwire_passed = unique_values > 100 and normalized_entropy > 0.5

    return {
        "unique_values": unique_values,
        "entropy_bits": float(entropy),
        "max_entropy_bits": float(max_entropy),
        "normalized_entropy": float(normalized_entropy),
        "tripwire_passed": tripwire_passed,
    }


if __name__ == "__main__":
    print("=== HelixCode Shared V0 Format ===")
    print()
    print(f"Schema: {HELIXCODE_SHARED_V0_SCHEMA}")
    print(f"Codebook size: {CODEBOOK_SIZE}")
    print()
    print("Usage:")
    print("  # Load shared codebook group")
    print("  group = HelixCodeSharedV0.load('seeds/helixcode_shared_v0/layer0_attn/')")
    print()
    print("  # Get tensor view (shares codebook)")
    print("  q_view = group.get_tensor('blk.0.attn_q.weight')")
    print()
    print("  # Compute matmul")
    print("  Y = q_view.matmul(X, transpose_w=True)")
    print()
    print("  # Verify sharing")
    print(f"  assert group.codebook_load_count == 1")
