#!/usr/bin/env python3
"""
helixcode_trainable.py - Trainable HelixCode with Procedural Indices

WO-HELIXCODE-TRAIN-01: A trainable VQ layer where:
- Indices are generated deterministically from a seed (SHAKE256 expansion)
- Only the 256-entry codebook is trained to create useful weights

This proves: "seed provides structure, training adjusts the 256 values"

NOT the disproven claim: "seed reconstructs pretrained weights"
IS the viable claim: "seed -> deterministic structure -> training creates useful weights"

Key Insight:
  For a 4096x4096 layer:
  - Dense Linear: 16M trainable params
  - HelixCode Trainable: 256 trainable params (62,500:1 compression)

  The codebook gather is differentiable: W = codebook[indices] allows gradients to flow.

Usage:
    from helix_cdc.regrow.helixcode_trainable import HelixCodeTrainableLinear

    # Create layer with procedural indices from seed
    layer = HelixCodeTrainableLinear(784, 128, seed=b"echo_v1", tensor_name="fc1")

    # Forward pass - gradients flow to codebook
    out = layer(x)

    # Only codebook + bias are trainable (< 300 params for this layer)
    trainable = sum(p.numel() for p in layer.parameters() if p.requires_grad)
"""

from __future__ import annotations

import hashlib
from typing import Optional

import numpy as np
import torch
import torch.nn as nn


CODEBOOK_SIZE = 256


class HelixCodeTrainableLinear(nn.Module):
    """
    VQ Linear with procedural indices and trainable codebook.

    - indices = SHAKE256(seed + tensor_name) -> uint8 [d_out, d_in] (buffer, not trained)
    - codebook = nn.Parameter([256] float32) (the ONLY trainable weights)
    - forward: W = codebook[indices], return X @ W.T

    The indices are deterministic given (seed, tensor_name) and are NOT trained.
    Only the 256-entry codebook learns during backprop.
    """

    def __init__(
        self,
        in_features: int,
        out_features: int,
        seed: bytes,
        tensor_name: str,
        bias: bool = True,
        codebook_init: str = "uniform",
    ):
        """
        Initialize trainable HelixCode layer.

        Args:
            in_features: Input dimension
            out_features: Output dimension
            seed: Deterministic seed for index generation
            tensor_name: Unique tensor name for index derivation
            bias: Whether to include bias
            codebook_init: Initialization method ("uniform", "normal", "linspace")
        """
        super().__init__()

        self.in_features = in_features
        self.out_features = out_features
        self.seed = seed
        self.tensor_name = tensor_name

        # Generate indices deterministically (buffer, NOT Parameter)
        indices = self._generate_indices(seed, tensor_name, (out_features, in_features))
        self.register_buffer("indices", indices)

        # Trainable codebook (256 floats) - THE ONLY TRAINED WEIGHTS
        self.codebook = nn.Parameter(torch.zeros(CODEBOOK_SIZE, dtype=torch.float32))
        self._init_codebook(codebook_init)

        # Optional bias (trainable)
        if bias:
            self.bias = nn.Parameter(torch.zeros(out_features, dtype=torch.float32))
        else:
            self.register_parameter("bias", None)

    def _generate_indices(
        self, seed: bytes, tensor_name: str, shape: tuple[int, int]
    ) -> torch.Tensor:
        """
        Generate deterministic uint8 indices via SHAKE256 expansion.

        This is the core "procedural" aspect - same (seed, tensor_name) always
        produces the same indices, so the structure is fixed.

        Args:
            seed: Base seed bytes
            tensor_name: Tensor identifier string
            shape: (out_features, in_features)

        Returns:
            torch.Tensor of uint8 indices with shape [out_features, in_features]
        """
        rows, cols = shape
        num_bytes = rows * cols

        # SHAKE256 expansion: seed || tensor_name -> raw bytes
        h = hashlib.shake_256(seed + tensor_name.encode("utf-8"))
        raw_bytes = h.digest(num_bytes)

        # Convert to uint8 array
        indices_np = np.frombuffer(raw_bytes, dtype=np.uint8).reshape(shape).copy()
        return torch.from_numpy(indices_np)

    def _init_codebook(self, method: str = "uniform"):
        """
        Initialize codebook for gradient-friendly training.

        Args:
            method: "uniform" (default), "normal", or "linspace"
        """
        with torch.no_grad():
            if method == "uniform":
                # Small uniform range for stable initial gradients
                nn.init.uniform_(self.codebook, -0.1, 0.1)
            elif method == "normal":
                nn.init.normal_(self.codebook, mean=0.0, std=0.02)
            elif method == "linspace":
                # Spread values across range (good for seeing gradient flow)
                self.codebook.copy_(
                    torch.linspace(-0.1, 0.1, CODEBOOK_SIZE)
                )
            else:
                raise ValueError(f"Unknown codebook_init: {method}")

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass: gather from codebook, then matmul.

        The gather operation `codebook[indices]` is differentiable - gradients
        flow back through the indexing to the codebook entries.

        Args:
            x: Input tensor [..., in_features]

        Returns:
            Output tensor [..., out_features]
        """
        # Gather: W[i, j] = codebook[indices[i, j]]
        # indices is uint8 buffer, convert to long for indexing
        W = self.codebook[self.indices.long()]  # [out_features, in_features]

        # Matmul: x @ W.T (standard linear layer convention)
        out = x @ W.t()

        # Add bias if present
        if self.bias is not None:
            out = out + self.bias

        return out

    def extra_repr(self) -> str:
        """String representation with trainable param count."""
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
        dense_equiv = self.in_features * self.out_features
        if self.bias is not None:
            dense_equiv += self.out_features
        ratio = dense_equiv / trainable if trainable > 0 else 0

        return (
            f"in_features={self.in_features}, out_features={self.out_features}, "
            f"bias={self.bias is not None}, "
            f"trainable={trainable}, dense_equiv={dense_equiv}, ratio={ratio:.1f}:1"
        )

    def indices_hash(self) -> str:
        """SHA256 of indices for determinism verification."""
        return hashlib.sha256(self.indices.numpy().tobytes()).hexdigest()


class HelixCodeTrainableMLP(nn.Module):
    """
    Multi-layer perceptron using HelixCodeTrainableLinear layers.

    Demonstrates that entire networks can use procedural indices with
    trained codebooks.
    """

    def __init__(
        self,
        seed: bytes,
        in_dim: int = 784,
        hidden_dim: int = 128,
        out_dim: int = 10,
        bias: bool = True,
    ):
        super().__init__()

        self.fc1 = HelixCodeTrainableLinear(
            in_dim, hidden_dim, seed, "fc1", bias=bias
        )
        self.fc2 = HelixCodeTrainableLinear(
            hidden_dim, out_dim, seed, "fc2", bias=bias
        )
        self.relu = nn.ReLU()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Forward pass through MLP."""
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

    def count_params(self) -> dict:
        """Count trainable vs dense-equivalent parameters."""
        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)

        # Dense equivalent
        dense_fc1 = self.fc1.in_features * self.fc1.out_features
        dense_fc2 = self.fc2.in_features * self.fc2.out_features
        if self.fc1.bias is not None:
            dense_fc1 += self.fc1.out_features
        if self.fc2.bias is not None:
            dense_fc2 += self.fc2.out_features
        dense_equiv = dense_fc1 + dense_fc2

        return {
            "trainable": trainable,
            "dense_equivalent": dense_equiv,
            "compression_ratio": dense_equiv / trainable if trainable > 0 else 0,
        }


def verify_index_determinism(
    seed: bytes, tensor_name: str, in_features: int, out_features: int, trials: int = 3
) -> bool:
    """
    Verify that indices are deterministic across multiple instantiations.

    Args:
        seed: Seed bytes
        tensor_name: Tensor name
        in_features: Input dimension
        out_features: Output dimension
        trials: Number of trials to verify

    Returns:
        True if all trials produce identical indices
    """
    hashes = []
    for _ in range(trials):
        layer = HelixCodeTrainableLinear(
            in_features, out_features, seed, tensor_name, bias=False
        )
        hashes.append(layer.indices_hash())

    return len(set(hashes)) == 1


if __name__ == "__main__":
    # Quick sanity check
    print("HelixCodeTrainableLinear sanity check:")

    layer = HelixCodeTrainableLinear(784, 128, b"test_seed", "test_layer")
    print(f"  Layer: {layer}")

    x = torch.randn(4, 784)
    y = layer(x)
    print(f"  Input: {x.shape}, Output: {y.shape}")

    # Verify gradients flow
    loss = y.sum()
    loss.backward()
    assert layer.codebook.grad is not None, "Codebook gradient is None!"
    print(f"  Codebook grad norm: {layer.codebook.grad.norm().item():.6f}")

    # Verify determinism
    deterministic = verify_index_determinism(b"test", "fc1", 784, 128)
    print(f"  Index determinism: {deterministic}")

    print("\nHelixCodeTrainableMLP sanity check:")
    mlp = HelixCodeTrainableMLP(b"echo_v1")
    params = mlp.count_params()
    print(f"  Trainable params: {params['trainable']}")
    print(f"  Dense equivalent: {params['dense_equivalent']}")
    print(f"  Compression ratio: {params['compression_ratio']:.1f}:1")
