"""
QuarterBit Functional Operations
================================

Copyright (c) 2026 Clouthier Simulation Labs. All rights reserved.
"""

import torch
import ctypes
from .utils import get_lib, is_available

def _ptr(tensor):
    """Get ctypes pointer to tensor data."""
    if tensor is None:
        return None
    return ctypes.cast(tensor.data_ptr(), ctypes.POINTER(ctypes.c_float))


def eft_accumulate(dst: torch.Tensor, src: torch.Tensor, comp: torch.Tensor = None):
    """
    Accumulate src into dst with EFT precision: dst += src

    If comp (compensation buffer) is not provided, creates one.
    Returns the compensation buffer for continued accumulation.

    Args:
        dst: Destination tensor (modified in-place)
        src: Source tensor to add
        comp: Compensation buffer (same shape as dst)

    Returns:
        comp: Compensation buffer

    Example:
        >>> dst = torch.zeros(1000).cuda()
        >>> comp = None
        >>> for batch in data:
        ...     comp = eft_accumulate(dst, batch, comp)
    """
    if not is_available():
        # Fallback: standard addition
        dst.add_(src)
        return torch.zeros_like(dst) if comp is None else comp

    if comp is None:
        comp = torch.zeros_like(dst)

    if not (dst.is_cuda and src.is_cuda and dst.dtype == torch.float32):
        # Fallback for non-CUDA or non-float32
        dst.add_(src)
        return comp

    lib = get_lib()
    n = dst.numel()

    lib.eft_accumulate(
        _ptr(dst), _ptr(comp), _ptr(src),
        ctypes.c_int(n)
    )

    return comp


def eft_accumulate_scaled(dst: torch.Tensor, src: torch.Tensor,
                          scale: float, comp: torch.Tensor = None):
    """
    Scaled accumulation with EFT: dst += scale * src

    Args:
        dst: Destination tensor
        src: Source tensor
        scale: Scaling factor
        comp: Compensation buffer

    Returns:
        comp: Compensation buffer
    """
    if not is_available():
        dst.add_(src, alpha=scale)
        return torch.zeros_like(dst) if comp is None else comp

    if comp is None:
        comp = torch.zeros_like(dst)

    if not (dst.is_cuda and src.is_cuda and dst.dtype == torch.float32):
        dst.add_(src, alpha=scale)
        return comp

    lib = get_lib()
    n = dst.numel()

    lib.eft_accumulate_scaled(
        _ptr(dst), _ptr(comp), _ptr(src),
        ctypes.c_float(scale),
        ctypes.c_int(n)
    )

    return comp


def eft_sum(tensor: torch.Tensor) -> torch.Tensor:
    """
    Sum all elements with EFT precision.

    More accurate than tensor.sum() for large tensors.

    Args:
        tensor: Input tensor

    Returns:
        Scalar tensor with sum

    Example:
        >>> x = torch.randn(1000000).cuda()
        >>> standard_sum = x.sum()
        >>> eft_precise_sum = eft_sum(x)
    """
    if not is_available() or not tensor.is_cuda or tensor.dtype != torch.float32:
        return tensor.sum()

    lib = get_lib()
    n = tensor.numel()

    output = torch.zeros(1, device=tensor.device, dtype=torch.float32)
    output_comp = torch.zeros(1, device=tensor.device, dtype=torch.float32)

    # Flatten for contiguous access
    flat = tensor.contiguous().view(-1)

    lib.eft_reduce_sum(
        _ptr(flat), _ptr(output), _ptr(output_comp),
        ctypes.c_int(n)
    )

    # Final compensation
    return output + output_comp


def eft_matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Matrix multiplication with EFT precision: C = A @ B

    More accurate than torch.mm for ill-conditioned matrices.

    Args:
        a: First matrix [M, K]
        b: Second matrix [K, N]

    Returns:
        Result matrix [M, N]

    Example:
        >>> A = torch.randn(1000, 500).cuda()
        >>> B = torch.randn(500, 1000).cuda()
        >>> C = eft_matmul(A, B)  # More precise than A @ B
    """
    if not is_available():
        return torch.mm(a, b)

    if not (a.is_cuda and b.is_cuda and a.dtype == torch.float32 and b.dtype == torch.float32):
        return torch.mm(a, b)

    if a.dim() != 2 or b.dim() != 2:
        raise ValueError("eft_matmul requires 2D tensors")

    M, K = a.shape
    K2, N = b.shape
    if K != K2:
        raise ValueError(f"Matrix dimensions don't match: {a.shape} @ {b.shape}")

    lib = get_lib()

    c = torch.empty(M, N, device=a.device, dtype=torch.float32)

    # Ensure contiguous
    a_contig = a.contiguous()
    b_contig = b.contiguous()

    lib.eft_matmul(
        _ptr(a_contig), _ptr(b_contig), _ptr(c),
        ctypes.c_int(M), ctypes.c_int(K), ctypes.c_int(N)
    )

    return c


def eft_dot(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
    """
    Dot product with EFT precision.

    Args:
        a: First vector
        b: Second vector

    Returns:
        Scalar tensor with dot product
    """
    return eft_sum(a * b)


def eft_mean(tensor: torch.Tensor) -> torch.Tensor:
    """
    Mean with EFT precision.

    Args:
        tensor: Input tensor

    Returns:
        Scalar tensor with mean
    """
    return eft_sum(tensor) / tensor.numel()


def eft_var(tensor: torch.Tensor, unbiased: bool = True) -> torch.Tensor:
    """
    Variance with EFT precision.

    Args:
        tensor: Input tensor
        unbiased: Use Bessel's correction (N-1 denominator)

    Returns:
        Scalar tensor with variance
    """
    mean = eft_mean(tensor)
    diff = tensor - mean
    ss = eft_sum(diff * diff)
    n = tensor.numel()
    return ss / (n - 1 if unbiased else n)
