"""
ctypes bindings for KeyDNN v2 CUDA MaxPool2D kernels.

This module provides low-level Python bindings to the CUDA implementation of
MaxPool2D forward/backward via `ctypes`. It is backend-specific and assumes:

- The CUDA DLL exports C ABI functions:
    - keydnn_cuda_set_device
    - keydnn_cuda_malloc / keydnn_cuda_free
    - keydnn_cuda_memcpy_h2d / keydnn_cuda_memcpy_d2h
    - keydnn_cuda_memset
    - keydnn_cuda_synchronize
    - keydnn_cuda_maxpool2d_forward_f32 / _f64
    - keydnn_cuda_maxpool2d_backward_f32 / _f64

- Device pointers are represented as uintptr_t handles (Python int).
- Tensors are NCHW contiguous.
- Forward expects x_pad already padded (usually with -inf padding for max semantics).
- Backward expects grad_x_pad device buffer to be zero-initialized before calling.

Platform notes
--------------
This wrapper currently targets Windows and loads the CUDA DLL from the build
output path you provided.

Design notes
------------
- Strict dtype and contiguity validation to avoid undefined behavior.
- No implicit device memory caching: caller owns device pointers and frees them.
- These functions are infrastructure-layer utilities; higher-level Tensor/Module
  code should wrap them with safe abstractions.
"""

from __future__ import annotations

import ctypes
from ctypes import (
    c_int,
    c_size_t,
    c_void_p,
    c_uint64,
    c_int64,
    c_float,
    c_double,
)
from pathlib import Path
from typing import Tuple

import numpy as np

from ._native_loader import load_keydnn_cuda_native  # dynamic loading, do not remove


# ---------------------------------------------------------------------
# Low-level CUDA utility bindings
# ---------------------------------------------------------------------

# We represent device pointers as uint64 handles (uintptr_t).
DevPtr = int


class CudaLib:
    """
    Thin binding layer around the KeyDNN CUDA native DLL.

    This class performs one-time `argtypes`/`restype` binding for exported symbols
    and exposes helpers for common CUDA operations and MaxPool2D dispatch.

    Notes
    -----
    - This class does not manage device pointers automatically; callers must
      free device allocations using `cuda_free`.
    - Pointers passed between Python and the DLL are treated as raw addresses
      (`void*`) on the Python side to avoid ctypes pointer-casting pitfalls.
    """

    def __init__(self, lib: ctypes.CDLL) -> None:
        """
        Create a CUDA binding wrapper around an already-loaded DLL.

        Parameters
        ----------
        lib : ctypes.CDLL
            Loaded KeyDNNV2CudaNative.dll handle.
        """
        self.lib = lib
        self._cuda_utils_bound = False
        self._maxpool_bound = False

    def _bind_cuda_utils(self) -> None:
        """Bind argtypes/restype for CUDA utility exports (idempotent)."""
        if self._cuda_utils_bound:
            return

        lib = self.lib
        lib.keydnn_cuda_set_device.argtypes = [c_int]
        lib.keydnn_cuda_set_device.restype = c_int

        lib.keydnn_cuda_malloc.argtypes = [ctypes.POINTER(c_uint64), c_size_t]
        lib.keydnn_cuda_malloc.restype = c_int

        lib.keydnn_cuda_free.argtypes = [c_uint64]
        lib.keydnn_cuda_free.restype = c_int

        lib.keydnn_cuda_memcpy_h2d.argtypes = [c_uint64, c_void_p, c_size_t]
        lib.keydnn_cuda_memcpy_h2d.restype = c_int

        lib.keydnn_cuda_memcpy_d2h.argtypes = [c_void_p, c_uint64, c_size_t]
        lib.keydnn_cuda_memcpy_d2h.restype = c_int

        lib.keydnn_cuda_memset.argtypes = [c_uint64, c_int, c_size_t]
        lib.keydnn_cuda_memset.restype = c_int

        lib.keydnn_cuda_synchronize.argtypes = []
        lib.keydnn_cuda_synchronize.restype = c_int

        lib.keydnn_cuda_memcpy_d2d.argtypes = [c_void_p, c_void_p, c_size_t]
        lib.keydnn_cuda_memcpy_d2d.restype = c_int

        self._cuda_utils_bound = True

    def _bind_maxpool2d(self) -> None:
        """
        Bind argtypes/restype for MaxPool2D CUDA exports (idempotent).

        Important
        ---------
        Although the native exports use typed pointers (float*/double*/int64*),
        on the Python/ctypes side we bind them as `c_void_p` to avoid fragile
        pointer casting and Windows-specific issues. The native side will still
        interpret the raw address correctly.
        """
        if self._maxpool_bound:
            return

        lib = self.lib

        # forward f32
        lib.keydnn_cuda_maxpool2d_forward_f32.argtypes = [
            c_void_p,  # x_pad (device)
            c_void_p,  # y (device)
            c_void_p,  # argmax_idx (device)
            c_int,
            c_int,  # N, C
            c_int,
            c_int,  # H_pad, W_pad
            c_int,
            c_int,  # H_out, W_out
            c_int,
            c_int,  # k_h, k_w
            c_int,
            c_int,  # s_h, s_w
        ]
        lib.keydnn_cuda_maxpool2d_forward_f32.restype = c_int

        # forward f64
        lib.keydnn_cuda_maxpool2d_forward_f64.argtypes = [
            c_void_p,
            c_void_p,
            c_void_p,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
        ]
        lib.keydnn_cuda_maxpool2d_forward_f64.restype = c_int

        # backward f32
        lib.keydnn_cuda_maxpool2d_backward_f32.argtypes = [
            c_void_p,  # grad_out (device)
            c_void_p,  # argmax_idx (device)
            c_void_p,  # grad_x_pad (device)
            c_int,
            c_int,  # N, C
            c_int,
            c_int,  # H_out, W_out
            c_int,
            c_int,  # H_pad, W_pad
        ]
        lib.keydnn_cuda_maxpool2d_backward_f32.restype = c_int

        # backward f64
        lib.keydnn_cuda_maxpool2d_backward_f64.argtypes = [
            c_void_p,
            c_void_p,
            c_void_p,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
            c_int,
        ]
        lib.keydnn_cuda_maxpool2d_backward_f64.restype = c_int

        self._maxpool_bound = True

    @staticmethod
    def _as_dev_ptr(dev_ptr: DevPtr) -> c_void_p:
        """
        Convert a device pointer handle (Python int) into a ctypes void*.

        Parameters
        ----------
        dev_ptr : DevPtr
            Device address stored as an integer (uintptr_t).

        Returns
        -------
        ctypes.c_void_p
            Raw pointer value usable in ctypes calls.
        """
        return c_void_p(int(dev_ptr))

    # ----------------------------
    # CUDA utils
    # ----------------------------

    def cuda_set_device(self, device: int = 0) -> None:
        """
        Set the current CUDA device.

        Parameters
        ----------
        device : int
            CUDA device ordinal.

        Raises
        ------
        RuntimeError
            If the native call returns a non-zero status.
        """
        self._bind_cuda_utils()
        st = self.lib.keydnn_cuda_set_device(int(device))
        if st != 0:
            raise RuntimeError(f"keydnn_cuda_set_device failed with status={st}")

    def cuda_malloc(self, nbytes: int) -> DevPtr:
        """
        Allocate device memory.

        Parameters
        ----------
        nbytes : int
            Number of bytes to allocate.

        Returns
        -------
        DevPtr
            Device pointer handle (uintptr_t as Python int).

        Raises
        ------
        RuntimeError
            If allocation fails.
        """
        self._bind_cuda_utils()
        out = c_uint64(0)
        st = self.lib.keydnn_cuda_malloc(ctypes.byref(out), c_size_t(int(nbytes)))
        if st != 0 or out.value == 0:
            raise RuntimeError(
                f"keydnn_cuda_malloc failed with status={st}, nbytes={nbytes}"
            )
        return int(out.value)

    def cuda_free(self, dev_ptr: DevPtr) -> None:
        """
        Free device memory.

        Parameters
        ----------
        dev_ptr : DevPtr
            Device pointer handle. Safe to pass 0.

        Raises
        ------
        RuntimeError
            If the native free call fails.
        """
        self._bind_cuda_utils()
        st = self.lib.keydnn_cuda_free(c_uint64(int(dev_ptr)))
        if st != 0:
            raise RuntimeError(f"keydnn_cuda_free failed with status={st}")

    def cuda_memcpy_h2d(self, dst_dev: DevPtr, src_host: np.ndarray) -> None:
        """
        Copy a NumPy array from host to device.

        Notes
        -----
        This implementation intentionally routes through the CFUNCTYPE-based
        memcpy wrapper (memcpy_ctypes) to avoid CDLL.argtypes collisions across
        modules (e.g., ops.memcpy_cuda vs maxpool2d_ctypes).
        """
        if not isinstance(src_host, np.ndarray):
            raise TypeError(f"src_host must be np.ndarray, got {type(src_host)!r}")
        if not src_host.flags["C_CONTIGUOUS"]:
            src_host = np.ascontiguousarray(src_host)

        # CFUNCTYPE-based call (immune to lib.symbol.argtypes mutations)
        from src.keydnn.infrastructure.native_cuda.python.ops import memcpy_ctypes as mc

        mc.cuda_memcpy_h2d(self.lib, int(dst_dev), src_host, int(src_host.nbytes))

    def cuda_memcpy_d2h(self, dst_host: np.ndarray, src_dev: DevPtr) -> None:
        """
        Copy from device to a NumPy host array.

        Notes
        -----
        This implementation intentionally routes through the CFUNCTYPE-based
        memcpy wrapper (memcpy_ctypes) to avoid CDLL.argtypes collisions across
        modules (e.g., ops.memcpy_cuda vs maxpool2d_ctypes).
        """
        if not isinstance(dst_host, np.ndarray):
            raise TypeError(f"dst_host must be np.ndarray, got {type(dst_host)!r}")
        if not dst_host.flags["C_CONTIGUOUS"]:
            raise ValueError("dst_host must be C-contiguous")

        # CFUNCTYPE-based call (immune to lib.symbol.argtypes mutations)
        from src.keydnn.infrastructure.native_cuda.python.ops import memcpy_ctypes as mc

        mc.cuda_memcpy_d2h(self.lib, dst_host, int(src_dev), int(dst_host.nbytes))

    def cuda_memset(self, dev_ptr: DevPtr, value: int, nbytes: int) -> None:
        """
        Set a device buffer to a byte value.

        Parameters
        ----------
        dev_ptr : DevPtr
            Device pointer handle.
        value : int
            Byte value [0,255] used by cudaMemset.
        nbytes : int
            Number of bytes to set.

        Raises
        ------
        RuntimeError
            If the native memset call fails.
        """
        self._bind_cuda_utils()
        st = self.lib.keydnn_cuda_memset(
            c_uint64(int(dev_ptr)), c_int(int(value)), c_size_t(int(nbytes))
        )
        if st != 0:
            raise RuntimeError(f"keydnn_cuda_memset failed with status={st}")

    def cuda_synchronize(self) -> None:
        """
        Synchronize the device (blocking).

        Raises
        ------
        RuntimeError
            If device synchronization fails.
        """
        self._bind_cuda_utils()
        st = self.lib.keydnn_cuda_synchronize()
        if st != 0:
            raise RuntimeError(f"keydnn_cuda_synchronize failed with status={st}")

    def cuda_from_host(self, x: np.ndarray) -> DevPtr:
        """
        Convenience helper: allocate a device buffer and copy a host array to GPU.

        Supported dtypes
        ----------------
        - np.float32
        - np.float64
        - np.int64  (needed for argmax indices)

        Parameters
        ----------
        x : np.ndarray
            Host array to upload. If not C-contiguous, it will be made contiguous.

        Returns
        -------
        DevPtr
            Device pointer handle.

        Raises
        ------
        TypeError
            If dtype is unsupported.
        RuntimeError
            If allocation or memcpy fails.
        """
        if x.dtype not in (np.float32, np.float64, np.int64):
            raise TypeError(
                f"cuda_from_host only supports float32/float64/int64, got {x.dtype}"
            )
        if not x.flags["C_CONTIGUOUS"]:
            x = np.ascontiguousarray(x)

        dev = self.cuda_malloc(x.nbytes)
        try:
            self.cuda_memcpy_h2d(dev, x)
        except Exception:
            self.cuda_free(dev)
            raise
        return dev

    # ----------------------------
    # MaxPool2D
    # ----------------------------

    def maxpool2d_forward_cuda(
        self,
        *,
        x_pad_dev: DevPtr,
        y_dev: DevPtr,
        argmax_idx_dev: DevPtr,
        N: int,
        C: int,
        H_pad: int,
        W_pad: int,
        H_out: int,
        W_out: int,
        k_h: int,
        k_w: int,
        s_h: int,
        s_w: int,
        dtype: np.dtype,
        sync: bool = True,
    ) -> None:
        """
        Run CUDA MaxPool2D forward on device buffers.

        Parameters
        ----------
        x_pad_dev, y_dev, argmax_idx_dev : DevPtr
            Device pointer handles (raw addresses).
        N, C : int
            Batch size and channels.
        H_pad, W_pad : int
            Spatial dimensions of the *padded* input.
        H_out, W_out : int
            Output spatial dimensions.
        k_h, k_w : int
            Pooling kernel size.
        s_h, s_w : int
            Strides.
        dtype : np.dtype
            np.float32 or np.float64 selects the kernel variant.
        sync : bool
            If True, call `cuda_synchronize` after kernel launch.

        Raises
        ------
        TypeError
            If dtype is unsupported.
        RuntimeError
            If the native kernel returns a non-zero status.
        """
        self._bind_maxpool2d()

        if dtype == np.float32:
            st = self.lib.keydnn_cuda_maxpool2d_forward_f32(
                self._as_dev_ptr(x_pad_dev),
                self._as_dev_ptr(y_dev),
                self._as_dev_ptr(argmax_idx_dev),
                int(N),
                int(C),
                int(H_pad),
                int(W_pad),
                int(H_out),
                int(W_out),
                int(k_h),
                int(k_w),
                int(s_h),
                int(s_w),
            )
        elif dtype == np.float64:
            st = self.lib.keydnn_cuda_maxpool2d_forward_f64(
                self._as_dev_ptr(x_pad_dev),
                self._as_dev_ptr(y_dev),
                self._as_dev_ptr(argmax_idx_dev),
                int(N),
                int(C),
                int(H_pad),
                int(W_pad),
                int(H_out),
                int(W_out),
                int(k_h),
                int(k_w),
                int(s_h),
                int(s_w),
            )
        else:
            raise TypeError(f"Unsupported dtype for maxpool2d_forward_cuda: {dtype}")

        if st != 0:
            raise RuntimeError(f"keydnn_cuda_maxpool2d_forward failed with status={st}")

        if sync:
            self.cuda_synchronize()

    def maxpool2d_backward_cuda(
        self,
        *,
        grad_out_dev: DevPtr,
        argmax_idx_dev: DevPtr,
        grad_x_pad_dev: DevPtr,
        N: int,
        C: int,
        H_out: int,
        W_out: int,
        H_pad: int,
        W_pad: int,
        dtype: np.dtype,
        sync: bool = True,
    ) -> None:
        """
        Run CUDA MaxPool2D backward on device buffers.

        Parameters
        ----------
        grad_out_dev : DevPtr
            Device pointer to grad_out, shape (N, C, H_out, W_out).
        argmax_idx_dev : DevPtr
            Device pointer to argmax indices from forward, shape (N, C, H_out, W_out),
            dtype int64.
        grad_x_pad_dev : DevPtr
            Device pointer to grad_x_pad, shape (N, C, H_pad, W_pad).
            Must be zero-initialized before calling.
        dtype : np.dtype
            np.float32 or np.float64 selects the kernel variant.
        sync : bool
            If True, call `cuda_synchronize` after kernel launch.

        Notes
        -----
        This backward implementation uses atomic adds to accumulate gradients.
        Callers must ensure `grad_x_pad_dev` is zeroed (e.g., via `cuda_memset`)
        before invoking this function.

        Raises
        ------
        TypeError
            If dtype is unsupported.
        RuntimeError
            If the native kernel returns a non-zero status.
        """
        self._bind_maxpool2d()

        if dtype == np.float32:
            st = self.lib.keydnn_cuda_maxpool2d_backward_f32(
                self._as_dev_ptr(grad_out_dev),
                self._as_dev_ptr(argmax_idx_dev),
                self._as_dev_ptr(grad_x_pad_dev),
                int(N),
                int(C),
                int(H_out),
                int(W_out),
                int(H_pad),
                int(W_pad),
            )
        elif dtype == np.float64:
            st = self.lib.keydnn_cuda_maxpool2d_backward_f64(
                self._as_dev_ptr(grad_out_dev),
                self._as_dev_ptr(argmax_idx_dev),
                self._as_dev_ptr(grad_x_pad_dev),
                int(N),
                int(C),
                int(H_out),
                int(W_out),
                int(H_pad),
                int(W_pad),
            )
        else:
            raise TypeError(f"Unsupported dtype for maxpool2d_backward_cuda: {dtype}")

        if st != 0:
            raise RuntimeError(
                f"keydnn_cuda_maxpool2d_backward failed with status={st}"
            )

        if sync:
            self.cuda_synchronize()

    def cuda_memcpy_d2d(self, dst_dev: DevPtr, src_dev: DevPtr, nbytes: int) -> None:
        self._bind_cuda_utils()
        st = self.lib.keydnn_cuda_memcpy_d2d(
            c_void_p(int(dst_dev)),
            c_void_p(int(src_dev)),
            c_size_t(int(nbytes)),
        )
        if st != 0:
            raise RuntimeError(f"keydnn_cuda_memcpy_d2d failed with status={st}")


# ---------------------------------------------------------------------
# Functional API (keeps backward compatibility with the earlier wrapper)
# ---------------------------------------------------------------------

_cuda_singleton: CudaLib | None = None


def _get_cuda(lib: ctypes.CDLL) -> CudaLib:
    """
    Return a cached `CudaLib` wrapper for a given `ctypes.CDLL`.

    This avoids repeating argtype binding on every call and keeps the public API
    of this module function-based.
    """
    global _cuda_singleton
    if _cuda_singleton is None or _cuda_singleton.lib is not lib:
        _cuda_singleton = CudaLib(lib)
    return _cuda_singleton


def cuda_set_device(lib: ctypes.CDLL, device: int = 0) -> None:
    """Module-level convenience wrapper for `CudaLib.cuda_set_device`."""
    _get_cuda(lib).cuda_set_device(device)


def cuda_malloc(lib: ctypes.CDLL, nbytes: int) -> DevPtr:
    """Module-level convenience wrapper for `CudaLib.cuda_malloc`."""
    return _get_cuda(lib).cuda_malloc(nbytes)


def cuda_free(lib: ctypes.CDLL, dev_ptr: DevPtr) -> None:
    """Module-level convenience wrapper for `CudaLib.cuda_free`."""
    _get_cuda(lib).cuda_free(dev_ptr)


def cuda_memcpy_h2d(lib: ctypes.CDLL, dst_dev: DevPtr, src_host: np.ndarray) -> None:
    """
    Module-level convenience wrapper for host->device memcpy.

    Important
    ---------
    This uses the CFUNCTYPE-based memcpy_ctypes implementation to avoid
    CDLL.argtypes collisions across different wrapper modules.
    """
    if not isinstance(src_host, np.ndarray):
        raise TypeError(f"src_host must be np.ndarray, got {type(src_host)!r}")
    if not src_host.flags["C_CONTIGUOUS"]:
        src_host = np.ascontiguousarray(src_host)

    from .ops import memcpy_ctypes as mc

    mc.cuda_memcpy_h2d(lib, int(dst_dev), src_host, int(src_host.nbytes))


def cuda_memcpy_d2h(lib: ctypes.CDLL, dst_host: np.ndarray, src_dev: DevPtr) -> None:
    """
    Module-level convenience wrapper for device->host memcpy.

    Important
    ---------
    This uses the CFUNCTYPE-based memcpy_ctypes implementation to avoid
    CDLL.argtypes collisions across different wrapper modules.
    """
    if not isinstance(dst_host, np.ndarray):
        raise TypeError(f"dst_host must be np.ndarray, got {type(dst_host)!r}")
    if not dst_host.flags["C_CONTIGUOUS"]:
        raise ValueError("dst_host must be C-contiguous")

    from .ops import memcpy_ctypes as mc

    mc.cuda_memcpy_d2h(lib, dst_host, int(src_dev), int(dst_host.nbytes))


def cuda_memset(lib: ctypes.CDLL, dev_ptr: DevPtr, value: int, nbytes: int) -> None:
    """Module-level convenience wrapper for `CudaLib.cuda_memset`."""
    _get_cuda(lib).cuda_memset(dev_ptr, value, nbytes)


def cuda_synchronize(lib: ctypes.CDLL) -> None:
    """Module-level convenience wrapper for `CudaLib.cuda_synchronize`."""
    _get_cuda(lib).cuda_synchronize()


def cuda_from_host(lib: ctypes.CDLL, x: np.ndarray) -> DevPtr:
    """Module-level convenience wrapper for `CudaLib.cuda_from_host`."""
    return _get_cuda(lib).cuda_from_host(x)


def maxpool2d_forward_cuda(
    lib: ctypes.CDLL,
    *,
    x_pad_dev: DevPtr,
    y_dev: DevPtr,
    argmax_idx_dev: DevPtr,
    N: int,
    C: int,
    H_pad: int,
    W_pad: int,
    H_out: int,
    W_out: int,
    k_h: int,
    k_w: int,
    s_h: int,
    s_w: int,
    dtype: np.dtype,
    sync: bool = True,
) -> None:
    """Module-level convenience wrapper for `CudaLib.maxpool2d_forward_cuda`."""
    _get_cuda(lib).maxpool2d_forward_cuda(
        x_pad_dev=x_pad_dev,
        y_dev=y_dev,
        argmax_idx_dev=argmax_idx_dev,
        N=N,
        C=C,
        H_pad=H_pad,
        W_pad=W_pad,
        H_out=H_out,
        W_out=W_out,
        k_h=k_h,
        k_w=k_w,
        s_h=s_h,
        s_w=s_w,
        dtype=dtype,
        sync=sync,
    )


def maxpool2d_backward_cuda(
    lib: ctypes.CDLL,
    *,
    grad_out_dev: DevPtr,
    argmax_idx_dev: DevPtr,
    grad_x_pad_dev: DevPtr,
    N: int,
    C: int,
    H_out: int,
    W_out: int,
    H_pad: int,
    W_pad: int,
    dtype: np.dtype,
    sync: bool = True,
) -> None:
    """Module-level convenience wrapper for `CudaLib.maxpool2d_backward_cuda`."""
    _get_cuda(lib).maxpool2d_backward_cuda(
        grad_out_dev=grad_out_dev,
        argmax_idx_dev=argmax_idx_dev,
        grad_x_pad_dev=grad_x_pad_dev,
        N=N,
        C=C,
        H_out=H_out,
        W_out=W_out,
        H_pad=H_pad,
        W_pad=W_pad,
        dtype=dtype,
        sync=sync,
    )


# ---------------------------------------------------------------------
# Optional convenience: end-to-end CPU NumPy -> GPU -> CPU for testing
# ---------------------------------------------------------------------


def maxpool2d_forward_cuda_from_numpy(
    lib: ctypes.CDLL,
    *,
    x_pad: np.ndarray,
    N: int,
    C: int,
    H_pad: int,
    W_pad: int,
    H_out: int,
    W_out: int,
    k_h: int,
    k_w: int,
    s_h: int,
    s_w: int,
    device: int = 0,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Convenience wrapper for quick correctness testing:
    - copies x_pad (CPU) -> GPU
    - runs forward
    - copies y and argmax_idx back to CPU
    - frees GPU buffers

    Parameters
    ----------
    x_pad : np.ndarray
        Padded input tensor on CPU, dtype float32/float64, shape (N,C,H_pad,W_pad).
    device : int
        CUDA device ordinal to use.

    Returns
    -------
    (y, argmax_idx) : Tuple[np.ndarray, np.ndarray]
        y : np.ndarray
            Output tensor, dtype float32/float64, shape (N,C,H_out,W_out).
        argmax_idx : np.ndarray
            Argmax indices, dtype int64, shape (N,C,H_out,W_out), flattened index
            into padded plane: h * W_pad + w.

    Notes
    -----
    This function is intended for debug/testing flows and performs allocations
    and device transfers internally.
    """
    if x_pad.dtype not in (np.float32, np.float64):
        raise TypeError(f"x_pad must be float32/float64, got {x_pad.dtype}")
    if not x_pad.flags["C_CONTIGUOUS"]:
        x_pad = np.ascontiguousarray(x_pad)

    cuda = _get_cuda(lib)
    cuda.cuda_set_device(device)

    y = np.empty((N, C, H_out, W_out), dtype=x_pad.dtype)
    argmax_idx = np.empty((N, C, H_out, W_out), dtype=np.int64)

    x_dev = cuda.cuda_from_host(x_pad)
    y_dev = cuda.cuda_malloc(y.nbytes)
    idx_dev = cuda.cuda_malloc(argmax_idx.nbytes)

    try:
        cuda.maxpool2d_forward_cuda(
            x_pad_dev=x_dev,
            y_dev=y_dev,
            argmax_idx_dev=idx_dev,
            N=N,
            C=C,
            H_pad=H_pad,
            W_pad=W_pad,
            H_out=H_out,
            W_out=W_out,
            k_h=k_h,
            k_w=k_w,
            s_h=s_h,
            s_w=s_w,
            dtype=x_pad.dtype,
            sync=True,
        )

        cuda.cuda_memcpy_d2h(y, y_dev)
        cuda.cuda_memcpy_d2h(argmax_idx, idx_dev)

    finally:
        cuda.cuda_free(x_dev)
        cuda.cuda_free(y_dev)
        cuda.cuda_free(idx_dev)

    return y, argmax_idx


def cuda_memcpy_htod(
    lib: ctypes.CDLL,
    dst_dev: DevPtr,
    src_host: np.ndarray,
    nbytes: int | None = None,
) -> None:
    """
    Backward-compatible alias for host->device memcpy.

    Parameters
    ----------
    lib : ctypes.CDLL
        Loaded CUDA DLL handle.
    dst_dev : DevPtr
        Destination device pointer.
    src_host : np.ndarray
        Source host array.
    nbytes : int | None, optional
        Legacy parameter used by older wrappers/tests. If provided, it must match
        `src_host.nbytes` (after contiguity normalization).

    Notes
    -----
    Older code/tests may look for `cuda_memcpy_htod` (or `cudaMemcpyHtoD`) and
    call it as (lib, dst_dev, src_host, nbytes).
    The canonical name in this module is `cuda_memcpy_h2d(lib, dst_dev, src_host)`.
    """
    # Normalize contiguity the same way the canonical wrapper does.
    if not src_host.flags["C_CONTIGUOUS"]:
        src_host = np.ascontiguousarray(src_host)

    if nbytes is not None and int(nbytes) != int(src_host.nbytes):
        raise ValueError(
            f"nbytes mismatch: got {int(nbytes)} but src_host.nbytes is {int(src_host.nbytes)}"
        )

    cuda_memcpy_h2d(lib, dst_dev, src_host)


def cuda_memcpy_dtoh(
    lib: ctypes.CDLL,
    dst_host: np.ndarray,
    src_dev: DevPtr,
    nbytes: int | None = None,
) -> None:
    """
    Backward-compatible alias for device->host memcpy.

    Parameters
    ----------
    lib : ctypes.CDLL
        Loaded CUDA DLL handle.
    dst_host : np.ndarray
        Destination host array (should be C-contiguous).
    src_dev : DevPtr
        Source device pointer.
    nbytes : int | None, optional
        Legacy parameter used by older wrappers/tests. If provided, it must match
        `dst_host.nbytes`.

    Notes
    -----
    The canonical name in this module is `cuda_memcpy_d2h(lib, dst_host, src_dev)`.
    """
    if not dst_host.flags["C_CONTIGUOUS"]:
        raise ValueError("dst_host must be C-contiguous")

    if nbytes is not None and int(nbytes) != int(dst_host.nbytes):
        raise ValueError(
            f"nbytes mismatch: got {int(nbytes)} but dst_host.nbytes is {int(dst_host.nbytes)}"
        )

    cuda_memcpy_d2h(lib, dst_host, src_dev)


# Optional extra aliases if tests probe these exact names
cudaMemcpyHtoD = cuda_memcpy_htod
cudaMemcpyDtoH = cuda_memcpy_dtoh


def cuda_memcpy_d2d(
    lib: ctypes.CDLL, dst_dev: DevPtr, src_dev: DevPtr, nbytes: int
) -> None:
    _get_cuda(lib).cuda_memcpy_d2d(dst_dev, src_dev, nbytes)


cuda_memcpy_dtod = cuda_memcpy_d2d
cudaMemcpyDtoD = cuda_memcpy_d2d
cudaMemcpyD2D = cuda_memcpy_d2d
