#!/usr/bin/env python3
"""
helix_cdc/backends/cdna_cuda_backend.py

CUDA-accelerated CDNA matmul using PyTorch operations.

This keeps indices and codebook on GPU, avoiding CPU-GPU transfers.
The key property: W is NEVER decoded to float32 (no full weight matrix).
"""

import torch
from typing import Tuple, Optional

# Check CUDA availability
_HAS_CUDA = torch.cuda.is_available()


def is_available() -> bool:
    """Check if CUDA backend is available."""
    return _HAS_CUDA


def cdna_matmul_cuda(
    indices: torch.Tensor,
    codebook: torch.Tensor,
    x: torch.Tensor,
) -> torch.Tensor:
    """
    Compute y = W @ x using CUDA-accelerated CDNA matmul.

    Uses bincount trick on GPU - W is NEVER decoded to float32.

    Args:
        indices: [M, K] uint8 (or int64) tensor on CUDA
        codebook: [B] float32 tensor on CUDA
        x: [K] float32 tensor on CUDA

    Returns:
        y: [M] float32 tensor on CUDA
    """
    M, K = indices.shape
    B = codebook.shape[0]

    # Ensure on same device
    device = x.device

    # Convert indices to int64 for scatter_add
    indices_long = indices.long()

    # Allocate output
    y = torch.zeros(M, dtype=torch.float32, device=device)

    # Process each row
    # For each row i: y[i] = codebook @ bincount(indices[i], weights=x)
    for i in range(M):
        # bincount with weights - accumulates x values by index
        weighted = torch.bincount(
            indices_long[i],
            weights=x,
            minlength=B
        )
        # Dot product with codebook
        y[i] = torch.dot(codebook, weighted.float())

    return y


def cdna_matmul_cuda_batch(
    indices: torch.Tensor,
    codebook: torch.Tensor,
    x_batch: torch.Tensor,
) -> torch.Tensor:
    """
    Batch CDNA matmul on CUDA: y_batch[n] = W @ x_batch[n].

    Uses vectorized operations where possible.

    Args:
        indices: [M, K] uint8/int64 tensor on CUDA
        codebook: [B] float32 tensor on CUDA
        x_batch: [batch, K] float32 tensor on CUDA

    Returns:
        y_batch: [batch, M] float32 tensor on CUDA
    """
    batch_size, K = x_batch.shape
    M = indices.shape[0]
    B = codebook.shape[0]
    device = x_batch.device

    # Convert indices to int64
    indices_long = indices.long()

    # Allocate output
    y_batch = torch.zeros(batch_size, M, dtype=torch.float32, device=device)

    # Process batch - can be parallelized on GPU
    for n in range(batch_size):
        x = x_batch[n]
        for i in range(M):
            weighted = torch.bincount(
                indices_long[i],
                weights=x,
                minlength=B
            )
            y_batch[n, i] = torch.dot(codebook, weighted.float())

    return y_batch


def cdna_matmul_cuda_vectorized(
    indices: torch.Tensor,
    codebook: torch.Tensor,
    x_batch: torch.Tensor,
) -> torch.Tensor:
    """
    Fully vectorized CDNA matmul on CUDA using scatter_add.

    This is the fastest GPU implementation - no Python loops.

    Args:
        indices: [M, K] uint8/int64 tensor on CUDA
        codebook: [B] float32 tensor on CUDA
        x_batch: [batch, K] float32 tensor on CUDA

    Returns:
        y_batch: [batch, M] float32 tensor on CUDA
    """
    batch_size, K = x_batch.shape
    M = indices.shape[0]
    B = codebook.shape[0]
    device = x_batch.device

    # Convert indices to int64
    indices_long = indices.long()  # [M, K]

    # Allocate accumulator: [batch, M, B]
    # For each (n, i, b): acc[n, i, b] = sum of x[n, k] where indices[i, k] == b
    acc = torch.zeros(batch_size, M, B, dtype=torch.float32, device=device)

    # Expand indices for batch dimension: [batch, M, K]
    indices_expanded = indices_long.unsqueeze(0).expand(batch_size, -1, -1)

    # Expand x for row dimension: [batch, M, K]
    x_expanded = x_batch.unsqueeze(1).expand(-1, M, -1)

    # scatter_add: acc[n, i, indices[i, k]] += x[n, k]
    acc.scatter_add_(2, indices_expanded, x_expanded)

    # Dot product with codebook: y[n, i] = sum_b(codebook[b] * acc[n, i, b])
    # = (acc @ codebook)
    y_batch = torch.matmul(acc, codebook)  # [batch, M]

    return y_batch


# Export the best implementation
cdna_matmul_cuda_best = cdna_matmul_cuda_vectorized
