#!/usr/bin/env python3
"""
helix_cdc/cdna_linear.py

CDNALinear - A drop-in replacement for nn.Linear that uses CDNA compression.

This enables end-to-end CDNA vs dense comparison by swapping Linear layers.

Two modes:
  1. from_dense(): Quantize a dense weight matrix to CDNA format (for testing)
  2. from_cdna_file(): Load pre-quantized CDNA shard (for production)

The key property: W is NEVER decoded to float32 during forward pass.
"""

import tempfile
from pathlib import Path
from typing import Optional, Tuple, Dict, Any

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

# Import CDNA infrastructure
from helix_cdc.regrow.cdna_stream import (
    CDNAStreamReader,
    stream_cdna_matmul,
    CDNA_MAGIC,
    CDNA_VERSION,
    CODEBOOK_SIZE,
)

# Try to import C backend
try:
    from helix_cdc.backends.cdna_c_backend import (
        is_available as c_backend_available,
        has_batch as c_has_batch,
        has_parallel as c_has_parallel,
        cdna_matmul_c,
        cdna_matmul_c_batch,
        get_backend_info,
    )
    _HAS_C_BACKEND = c_backend_available()
    _HAS_C_BATCH = c_has_batch()
    _HAS_C_PARALLEL = c_has_parallel()
except ImportError:
    _HAS_C_BACKEND = False
    _HAS_C_BATCH = False
    _HAS_C_PARALLEL = False
    cdna_matmul_c = None
    cdna_matmul_c_batch = None
    get_backend_info = lambda: {"available": False}

# Try to import CUDA backend
try:
    from helix_cdc.backends.cdna_cuda_backend import (
        is_available as cuda_backend_available,
        cdna_matmul_cuda_vectorized,
    )
    _HAS_CUDA_BACKEND = cuda_backend_available()
except ImportError:
    _HAS_CUDA_BACKEND = False
    cdna_matmul_cuda_vectorized = None


def quantize_to_cdna(
    weight: np.ndarray,
    n_bins: int = 256,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Quantize a dense weight matrix to CDNA format.

    Uses k-means style uniform quantization to 256 bins.

    Args:
        weight: [M, K] float32 weight matrix
        n_bins: Number of codebook entries (default 256 for uint8)

    Returns:
        indices: [M, K] uint8 - codebook indices
        codebook: [256] float32 - the codebook values
    """
    # Flatten for histogram
    flat = weight.flatten().astype(np.float32)

    # Create uniform bins across weight range
    vmin, vmax = flat.min(), flat.max()

    # Small epsilon to avoid edge cases
    eps = 1e-7
    vmin -= eps
    vmax += eps

    # Create codebook (bin centers)
    bin_edges = np.linspace(vmin, vmax, n_bins + 1)
    codebook = (bin_edges[:-1] + bin_edges[1:]) / 2
    codebook = codebook.astype(np.float32)

    # Quantize: assign each weight to nearest bin
    indices = np.digitize(flat, bin_edges[1:-1]).astype(np.uint8)
    indices = indices.reshape(weight.shape)

    return indices, codebook


def write_cdna_file(
    indices: np.ndarray,
    codebook: np.ndarray,
    output_path: Path,
    tensor_name: str = "weight",
) -> None:
    """
    Write CDNA format HXZ file.

    Format:
      - 4 bytes: magic "HXZC"
      - 2 bytes: version (big-endian)
      - 2 bytes: header length (big-endian)
      - N bytes: JSON header
      - 1024 bytes: codebook (256 * float32)
      - compressed indices blob
    """
    import json
    import struct
    import zlib

    M, K = indices.shape

    # Build header
    header = {
        "tensor_name": tensor_name,
        "shape": [M, K],
        "encoding": "zlib",
        "dtype": "uint8",
        "codebook_size": CODEBOOK_SIZE,
    }
    header_json = json.dumps(header).encode("utf-8")

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

    with open(output_path, "wb") as f:
        # Magic
        f.write(CDNA_MAGIC)

        # Version
        f.write(struct.pack(">H", CDNA_VERSION))

        # Header length
        f.write(struct.pack(">H", len(header_json)))

        # Header
        f.write(header_json)

        # Codebook
        f.write(codebook.astype(np.float32).tobytes())

        # Compressed indices
        f.write(compressed)


class CDNALinear(nn.Module):
    """
    Drop-in replacement for nn.Linear using CDNA compression.

    Key property: Weight matrix W is NEVER decoded to float32.
    Compute is done via bincount + codebook dot product.

    Usage:
        # From existing Linear layer
        cdna_linear = CDNALinear.from_dense(existing_linear)

        # Or from pre-quantized CDNA file
        cdna_linear = CDNALinear.from_cdna_file("path/to/weight.cdna.hxz")
    """

    def __init__(
        self,
        cdna_path: str,
        bias: Optional[torch.Tensor] = None,
        device: torch.device = torch.device("cpu"),
    ):
        """
        Initialize from CDNA file path.

        Args:
            cdna_path: Path to .cdna.hxz file
            bias: Optional bias tensor
            device: Device for output tensors
        """
        super().__init__()

        self.cdna_path = cdna_path
        self.device = device

        # Load reader to get shape info
        self.reader = CDNAStreamReader(cdna_path)
        self.out_features, self.in_features = self.reader.shape

        # Bias (optional)
        if bias is not None:
            self.register_buffer("bias", bias.to(device))
        else:
            self.bias = None

        # Track stats
        self.forward_count = 0
        self.w_decoded_to_float32 = False  # Key CDC claim

        # Determine best backend
        self.using_cuda_backend = _HAS_CUDA_BACKEND and device.type == "cuda"
        self.using_c_backend = _HAS_C_BACKEND and not self.using_cuda_backend

        # Pre-load indices and codebook
        indices_blob = self.reader._ensure_indices_blob()
        indices_np = np.frombuffer(indices_blob, dtype=np.uint8).reshape(
            self.out_features, self.in_features
        ).copy()
        codebook_np = self.reader.codebook.copy()

        if self.using_cuda_backend:
            # Store as CUDA tensors for GPU compute
            self._indices_cuda = torch.from_numpy(indices_np).to(device)
            self._codebook_cuda = torch.from_numpy(codebook_np).to(device)
            self._indices = None
            self._codebook = None
        elif self.using_c_backend:
            # Store as numpy for C backend
            self._indices = indices_np
            self._codebook = codebook_np
            self._indices_cuda = None
            self._codebook_cuda = None
        else:
            self._indices = None
            self._codebook = None
            self._indices_cuda = None
            self._codebook_cuda = None

    @classmethod
    def from_dense(
        cls,
        linear: nn.Linear,
        temp_dir: Optional[str] = None,
        tensor_name: str = "weight",
    ) -> "CDNALinear":
        """
        Create CDNALinear from an existing nn.Linear layer.

        This quantizes the dense weights to CDNA format.

        Args:
            linear: Source nn.Linear layer
            temp_dir: Directory for temporary CDNA file (default: system temp)
            tensor_name: Name for the tensor in CDNA metadata

        Returns:
            CDNALinear instance
        """
        # Extract weight
        weight = linear.weight.detach().cpu().numpy()  # [out, in]

        # Quantize to CDNA
        indices, codebook = quantize_to_cdna(weight)

        # Write to temporary file
        if temp_dir is None:
            temp_dir = tempfile.gettempdir()
        temp_path = Path(temp_dir) / f"{tensor_name}.cdna.hxz"
        write_cdna_file(indices, codebook, temp_path, tensor_name)

        # Extract bias if present
        bias = None
        if linear.bias is not None:
            bias = linear.bias.detach().clone()

        # Create instance
        device = next(linear.parameters()).device
        return cls(str(temp_path), bias=bias, device=device)

    @classmethod
    def from_cdna_file(
        cls,
        cdna_path: str,
        bias: Optional[torch.Tensor] = None,
        device: torch.device = torch.device("cpu"),
    ) -> "CDNALinear":
        """
        Create CDNALinear from pre-quantized CDNA file.

        Args:
            cdna_path: Path to .cdna.hxz file
            bias: Optional bias tensor
            device: Device for output tensors

        Returns:
            CDNALinear instance
        """
        return cls(cdna_path, bias=bias, device=device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass using CDNA matmul.

        W is NEVER decoded to float32!

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

        Returns:
            Output tensor [..., out_features]
        """
        self.forward_count += 1

        # Handle batched input
        orig_shape = x.shape
        if x.dim() > 2:
            # Flatten to [batch, in_features]
            x_flat = x.view(-1, self.in_features)
        else:
            x_flat = x

        batch_size = x_flat.shape[0]

        # Choose best backend (priority: CUDA > C > Python)
        if self.using_cuda_backend and self._indices_cuda is not None:
            # CUDA backend - fastest, stays on GPU
            y = cdna_matmul_cuda_vectorized(
                self._indices_cuda, self._codebook_cuda, x_flat
            )
            # CUDA backend never decodes W to float32
        elif self.using_c_backend and self._indices is not None:
            # C backend - fast, but needs CPU transfer
            x_np = x_flat.detach().cpu().numpy().astype(np.float32)
            if _HAS_C_BATCH and cdna_matmul_c_batch is not None:
                y_np = cdna_matmul_c_batch(
                    self._indices, self._codebook, x_np,
                    use_parallel=_HAS_C_PARALLEL
                )
            else:
                y_np = np.zeros((batch_size, self.out_features), dtype=np.float32)
                for i in range(batch_size):
                    y_np[i] = cdna_matmul_c(self._indices, self._codebook, x_np[i])
            y = torch.from_numpy(y_np).to(self.device)
            # C backend never decodes W to float32
        else:
            # Python reference backend - slowest
            x_np = x_flat.detach().cpu().numpy().astype(np.float32)
            y_np = np.zeros((batch_size, self.out_features), dtype=np.float32)
            for i in range(batch_size):
                y_row, stats = stream_cdna_matmul(self.cdna_path, x_np[i])
                y_np[i] = y_row
                if stats.get("w_decoded_to_float32", False):
                    self.w_decoded_to_float32 = True
            y = torch.from_numpy(y_np).to(self.device)

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

        # Restore original batch shape
        if len(orig_shape) > 2:
            y = y.view(*orig_shape[:-1], self.out_features)

        return y

    def get_stats(self) -> Dict[str, Any]:
        """Get CDNA layer statistics."""
        if self.using_cuda_backend:
            backend_str = "cuda"
        elif self.using_c_backend:
            if _HAS_C_PARALLEL:
                backend_str = "c_parallel"
            elif _HAS_C_BATCH:
                backend_str = "c_batch"
            else:
                backend_str = "c"
        else:
            backend_str = "python"
        return {
            "cdna_path": self.cdna_path,
            "shape": [self.out_features, self.in_features],
            "forward_count": self.forward_count,
            "w_decoded_to_float32": self.w_decoded_to_float32,
            "has_bias": self.bias is not None,
            "backend": backend_str,
            "memory_stats": self.reader.memory_stats,
        }

    def extra_repr(self) -> str:
        return f"in_features={self.in_features}, out_features={self.out_features}, bias={self.bias is not None}"


def swap_linear_to_cdna(
    module: nn.Module,
    temp_dir: Optional[str] = None,
    prefix: str = "",
) -> Dict[str, CDNALinear]:
    """
    Recursively swap all nn.Linear layers in a module with CDNALinear.

    Args:
        module: PyTorch module to modify
        temp_dir: Directory for temporary CDNA files
        prefix: Name prefix for tensor naming

    Returns:
        Dict mapping layer names to CDNALinear instances
    """
    swapped = {}

    for name, child in list(module.named_children()):
        full_name = f"{prefix}.{name}" if prefix else name

        if isinstance(child, nn.Linear):
            # Swap this Linear layer
            cdna_layer = CDNALinear.from_dense(
                child,
                temp_dir=temp_dir,
                tensor_name=full_name.replace(".", "_"),
            )
            setattr(module, name, cdna_layer)
            swapped[full_name] = cdna_layer
        else:
            # Recurse into children
            child_swapped = swap_linear_to_cdna(child, temp_dir, full_name)
            swapped.update(child_swapped)

    return swapped
