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

Python ctypes binding to libcdna_matmul.so - the optimized C CDNA matmul.

Supports:
  - cdna_matmul_v0: Single vector matmul
  - cdna_matmul_batch: Batch processing (multiple vectors)
  - cdna_matmul_batch_parallel: Batch + OpenMP parallel
"""

import ctypes
import os
from pathlib import Path
from typing import Optional, Tuple, Dict, Any

import numpy as np

# Try to find the library (prefer optimized version with OpenMP)
_LIB_PATHS = [
    # Optimized version with OpenMP (preferred)
    Path(__file__).parent.parent.parent / "c" / "bin" / "libcdna_matmul_opt.so",
    # Original version
    Path(__file__).parent.parent.parent / "proof" / "HELIX_CDNA_V1" / "libcdna_matmul.so",
    Path(__file__).parent.parent.parent / "c" / "bin" / "libhelix_cdc.so",
]

_lib: Optional[ctypes.CDLL] = None
_has_batch: bool = False
_has_batch_parallel: bool = False
_lib_path: Optional[str] = None


def _load_library() -> Optional[ctypes.CDLL]:
    """Load the CDNA matmul shared library."""
    global _lib, _lib_path
    if _lib is not None:
        return _lib

    for path in _LIB_PATHS:
        if path.exists():
            try:
                _lib = ctypes.CDLL(str(path))
                _lib_path = str(path)
                print(f"[cdna_c_backend] Loaded: {path}")
                return _lib
            except OSError as e:
                print(f"[cdna_c_backend] Failed to load {path}: {e}")
                continue

    print("[cdna_c_backend] No library found, C backend unavailable")
    return None


def _setup_functions(lib: ctypes.CDLL) -> None:
    """Set up ctypes function signatures."""
    global _has_batch, _has_batch_parallel

    # cdna_matmul_v0 - single vector
    lib.cdna_matmul_v0.argtypes = [
        ctypes.POINTER(ctypes.c_uint8),  # indices_W
        ctypes.c_size_t,                  # M
        ctypes.c_size_t,                  # K
        ctypes.POINTER(ctypes.c_float),   # codebook_W
        ctypes.c_size_t,                  # B
        ctypes.POINTER(ctypes.c_float),   # x
        ctypes.POINTER(ctypes.c_float),   # y_out
    ]
    lib.cdna_matmul_v0.restype = ctypes.c_int

    # cdna_matmul_batch - batch processing
    if hasattr(lib, 'cdna_matmul_batch'):
        lib.cdna_matmul_batch.argtypes = [
            ctypes.POINTER(ctypes.c_uint8),  # indices_W
            ctypes.c_size_t,                  # M
            ctypes.c_size_t,                  # K
            ctypes.POINTER(ctypes.c_float),   # codebook_W
            ctypes.c_size_t,                  # B
            ctypes.POINTER(ctypes.c_float),   # x_batch
            ctypes.POINTER(ctypes.c_float),   # y_batch
            ctypes.c_size_t,                  # batch_size
        ]
        lib.cdna_matmul_batch.restype = ctypes.c_int
        _has_batch = True
        print("[cdna_c_backend] Batch processing available")

    # cdna_matmul_batch_parallel - batch + OpenMP
    if hasattr(lib, 'cdna_matmul_batch_parallel'):
        lib.cdna_matmul_batch_parallel.argtypes = [
            ctypes.POINTER(ctypes.c_uint8),  # indices_W
            ctypes.c_size_t,                  # M
            ctypes.c_size_t,                  # K
            ctypes.POINTER(ctypes.c_float),   # codebook_W
            ctypes.c_size_t,                  # B
            ctypes.POINTER(ctypes.c_float),   # x_batch
            ctypes.POINTER(ctypes.c_float),   # y_batch
            ctypes.c_size_t,                  # batch_size
        ]
        lib.cdna_matmul_batch_parallel.restype = ctypes.c_int
        _has_batch_parallel = True
        print("[cdna_c_backend] OpenMP parallel batch available")


# Load on import
_lib = _load_library()
if _lib is not None:
    _setup_functions(_lib)


def is_available() -> bool:
    """Check if C backend is available."""
    return _lib is not None


def has_batch() -> bool:
    """Check if batch processing is available."""
    return _has_batch


def has_parallel() -> bool:
    """Check if OpenMP parallel batch is available."""
    return _has_batch_parallel


def get_backend_info() -> Dict[str, Any]:
    """Get backend capability info."""
    return {
        "available": _lib is not None,
        "library": _lib_path,
        "has_batch": _has_batch,
        "has_parallel": _has_batch_parallel,
    }


def cdna_matmul_c(
    indices: np.ndarray,
    codebook: np.ndarray,
    x: np.ndarray,
) -> np.ndarray:
    """
    Compute y = W @ x using C backend.

    Args:
        indices: [M, K] uint8 codebook indices
        codebook: [B] float32 codebook values
        x: [K] float32 input vector

    Returns:
        y: [M] float32 output vector
    """
    if _lib is None:
        raise RuntimeError("C backend not available")

    # Validate inputs
    assert indices.dtype == np.uint8, f"indices must be uint8, got {indices.dtype}"
    assert codebook.dtype == np.float32, f"codebook must be float32, got {codebook.dtype}"
    assert x.dtype == np.float32, f"x must be float32, got {x.dtype}"

    M, K = indices.shape
    B = len(codebook)

    assert x.shape[0] == K, f"x shape {x.shape} doesn't match indices K={K}"

    # Ensure contiguous
    indices = np.ascontiguousarray(indices)
    codebook = np.ascontiguousarray(codebook)
    x = np.ascontiguousarray(x)

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

    # Get pointers
    indices_ptr = indices.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
    codebook_ptr = codebook.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    x_ptr = x.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    y_ptr = y.ctypes.data_as(ctypes.POINTER(ctypes.c_float))

    # Call C function
    rc = _lib.cdna_matmul_v0(
        indices_ptr, M, K, codebook_ptr, B, x_ptr, y_ptr
    )

    if rc != 0:
        raise RuntimeError(f"cdna_matmul_c failed with error code {rc}")

    return y


def cdna_matmul_c_batch(
    indices: np.ndarray,
    codebook: np.ndarray,
    x_batch: np.ndarray,
    use_parallel: bool = True,
) -> np.ndarray:
    """
    Batch CDNA matmul: y_batch[i] = W @ x_batch[i].

    Uses C-level batch processing when available, with optional OpenMP.

    Args:
        indices: [M, K] uint8
        codebook: [B] float32
        x_batch: [batch, K] float32
        use_parallel: Use OpenMP parallel version if available

    Returns:
        y_batch: [batch, M] float32
    """
    if _lib is None:
        raise RuntimeError("C backend not available")

    batch_size = x_batch.shape[0]
    M, K = indices.shape
    B = len(codebook)

    # Ensure contiguous
    indices = np.ascontiguousarray(indices)
    codebook = np.ascontiguousarray(codebook)
    x_batch = np.ascontiguousarray(x_batch.astype(np.float32))

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

    # Get pointers
    indices_ptr = indices.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))
    codebook_ptr = codebook.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    x_ptr = x_batch.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
    y_ptr = y_batch.ctypes.data_as(ctypes.POINTER(ctypes.c_float))

    # Choose best available function
    if use_parallel and _has_batch_parallel:
        rc = _lib.cdna_matmul_batch_parallel(
            indices_ptr, M, K, codebook_ptr, B, x_ptr, y_ptr, batch_size
        )
    elif _has_batch:
        rc = _lib.cdna_matmul_batch(
            indices_ptr, M, K, codebook_ptr, B, x_ptr, y_ptr, batch_size
        )
    else:
        # Fall back to loop
        for i in range(batch_size):
            y_batch[i] = cdna_matmul_c(indices, codebook, x_batch[i])
        return y_batch

    if rc != 0:
        raise RuntimeError(f"cdna_matmul_batch failed with error code {rc}")

    return y_batch
