# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0

"""
This module defines a kernel generator for various sparse tensor operations.
"""

__all__ = []

import time

from cuda import pathfinder

from ._cpp import prolog_a_impl, prolog_b_impl, prolog_decl
from ._emitter import Index, emit_matmul_kernel
from ._jit import compile_cpp_and_link
from ._utils import Cache, type_str

# Maps matmul(a,b,c) to pre-defined kernel source code.
_tensor_matmul_kernels = {
    # (format_a, format_b, format_c, transpose_a, transpose_b)
    ("DIAI", "DenseVector", "DenseVector", False, False): """
extern "C" __global__ void matmul_kernel(
    const POS* __restrict__ pos0,
    const CRD* __restrict__ diags,
    const VAL* __restrict__ A,
    const VAL *__restrict__ B,
    VAL *__restrict__ C,
    unsigned long m,
    unsigned long n,
    unsigned long l0,
    unsigned long l1) {
  const unsigned long tidx = threadIdx.x;
  const unsigned long bidx = blockIdx.x;
  const unsigned long i = bidx * blockDim.x + tidx;
  if (i < m) {
    CTP acc = 0;
    const POS numDiags = pos0[1];
    for (unsigned long d = 0, di = i; d < numDiags; d++, di += m) {
      const long j = i + diags[d];
      if (0 <= j && j < n) {
        acc += prolog_a(static_cast<CTP>(A[di])) *
               prolog_b(static_cast<CTP>(B[j]));
      }
    }
    C[i] += static_cast<VAL>(acc);
  }
}
""",
    ("DIAJ", "DenseVector", "DenseVector", False, False): """
extern "C" __global__ void matmul_kernel(
    const POS* __restrict__ pos0,
    const CRD* __restrict__ diags,
    const VAL* __restrict__ A,
    const VAL *__restrict__ B,
    VAL *__restrict__ C,
    unsigned long m,
    unsigned long n,
    unsigned long l0,
    unsigned long l1) {
  const unsigned long tidx = threadIdx.x;
  const unsigned long bidx = blockIdx.x;
  const unsigned long i = bidx * blockDim.x + tidx;
  if (i < m) {
    CTP acc = 0;
    const POS numDiags = pos0[1];
    for (unsigned long d = 0, di = 0; d < numDiags; d++, di += n) {
      const long j = i + diags[d];
      if (0 <= j && j < n) {
        acc += static_cast<CTP>(prolog_a(A[di + j])) *
               static_cast<CTP>(prolog_b(B[j]));
      }
    }
    C[i] += static_cast<VAL>(acc);
  }
}
""",
}


def _set_user_or_default_code(*, user_code, default_code, default_prolog, src_code_list, obj_code_list):
    """
    Args:
        user_code: The user-provided LTO-IR as a bytes object or CUDA C++ source code
            string.
        default_code: The CUDA C++ source code that will be used if the user-provided
            code is None.
        default_prolog: The prolog to the CUDA C++ source code that will be used if the
            user-provided code is None.
        src_code_list: The list to which the source code will be appended to.
        obj_code_list: The list to which the LTO-IR code will be appended to.

    Returns:
        The pair (src_code_list, obj_code_list)
    """

    if user_code is not None:
        if isinstance(user_code, str):
            target = src_code_list
        elif isinstance(user_code, bytes):
            target = obj_code_list
        else:
            raise AssertionError("Internal error.")
        target.append(user_code)
        return src_code_list, obj_code_list

    assert default_code is not None and isinstance(default_code, str), "Internal error."

    source_code = f"{default_prolog} {default_code}"
    src_code_list.append(source_code)

    return src_code_list, obj_code_list


class KernelGen:
    """
    This class defines a kernel generator for various sparse tensor operations.
    The generated kernels can be JITted on-demand. Caching avoids recompilation
    on subsequent uses.

    The generation uses:
    (1) kernel lookup from a collection of pre-defined kernels
    (2) kernel generated by a sparse emitter

    Clients can easily add new pre-defined kernels in the table above (for example,
    by hand optimizing the automatically generated cuda code from the sparse emitter).
    """

    def __init__(self, debug=False):
        self._matmul_kernels_cache = Cache()
        self._debug = debug

    def generate_matmul(
        self,
        a,
        b,
        c,
        compute_type,
        prologs=None,
        transpose_a=False,
        transpose_b=False,
        arch=None,
        kernel=0,
    ):
        """
        Matmul semantics:

        Dot Product: If both inputs are 1-dimensional, it computes the dot
        product (inner product).

        Matrix-Vector Product: If the first input is a 2D tensor (matrix) and
        the second is a 1D tensor (vector), it performs a matrix-vector
        multiplication.

        Vector-Matrix Product: If the first input is a 1D tensor and the second
        is a 2D tensor, it performs a vector-matrix multiplication.

        Matrix-Matrix Product: If both inputs are 2D tensors, it performs
        standard matrix multiplication, requiring the number of columns in
        the first matrix to equal the number of rows in the second.

        Batched Matrix Product: If either input has more than two dimensions,
        the extra leading dimensions are treated as "batch dimensions" and it
        performs a batched matrix multiplication over the last two dimensions.
        """

        format_a = a.tensor_format.name
        format_b = b.tensor_format.name
        format_c = c.tensor_format.name
        vtp = type_str(a.dtype)
        itp = type_str(a.index_type)

        # Get user code.
        if prologs is None:
            user_prolog_a, user_prolog_b = None, None
        else:
            user_prolog_a, user_prolog_b, user_prolog_c = prologs.get("a"), prologs.get("b"), prologs.get("c")
            if user_prolog_c is not None:
                raise NotImplementedError("The code generation currently does not support prolog for C.")

        # Try cache lookup.
        key = (
            format_a,
            format_b,
            format_c,
            vtp,
            itp,
            compute_type,
            user_prolog_a,
            user_prolog_b,
            transpose_a,
            transpose_b,
            arch,
            kernel,
        )
        entry = self._matmul_kernels_cache.get(key)
        if entry is not None:
            return entry

        # Nothing in cache, try kernel lookup. Note that kernel==0 overrides codegen
        # with this lookup mechanism, all subsequent choices go to actual codegen.
        typeless_key = (format_a, format_b, format_c, transpose_a, transpose_b)
        code = _tensor_matmul_kernels.get(typeless_key) if kernel == 0 else None

        # Prepare prolog.
        ctp = type_str(compute_type)
        if a.dtype in ["float8_e4m3fn", "float8_e5m2"]:
            prolog = f"#include <cuda_fp8.h>\n\nusing CTP = {ctp};\n"
        elif a.dtype == "float16":
            prolog = f"#include <cuda_fp16.h>\n\nusing CTP = {ctp};\n"
        elif a.dtype == "bfloat16":
            prolog = f"#include <cuda_bf16.h>\n\nusing CTP = {ctp};\n"
        elif a.dtype in ["complex64", "complex128"]:
            prolog = f"#include <cuda/std/complex>\n\nusing CTP = {ctp};\n"
        else:
            prolog = f"using CTP = {ctp};\n"

        # Use kernel lookup or emitter.
        if code is not None:
            # (1) Kernel lookup successful. Inject types into the typeless kernel.
            source = f"{prolog}using VAL = {vtp};\nusing POS = {itp};\nusing CRD = {itp};\n{prolog_decl}{code}"
            grid_iter = [(Index.FREE, 0)], True
        else:
            # (2) Sparse emitter. Emits proper types.
            source, grid_iter = emit_matmul_kernel(
                a, b, c, ctp, transpose_a=transpose_a, transpose_b=transpose_b, kernel=kernel
            )

        # The source and object code lists.
        src_code_list, obj_code_list = [source], []

        # User-provided or default prolog operation on A and B.
        src_code_list, obj_code_list = _set_user_or_default_code(
            user_code=user_prolog_a,
            default_code=prolog_a_impl,
            default_prolog=prolog,
            src_code_list=src_code_list,
            obj_code_list=obj_code_list,
        )
        src_code_list, obj_code_list = _set_user_or_default_code(
            user_code=user_prolog_b,
            default_code=prolog_b_impl,
            default_prolog=prolog,
            src_code_list=src_code_list,
            obj_code_list=obj_code_list,
        )

        # TODO: this should use logger.debug(...)
        if self._debug:
            print("\n==> Compiling matmul\n")
            for s in src_code_list:
                print(s)
                print("/////////////////////\n")
            t1 = time.perf_counter_ns()

        # Compile and link matmul function. Cache the result.
        complex_path = pathfinder.find_nvidia_header_directory("cccl")
        narrow_prec_path = pathfinder.find_nvidia_header_directory("cudart")
        compiler_options = {
            "std": "c++17",
            "link_time_optimization": True,
            "arch": arch,
            "include_path": [complex_path, narrow_prec_path],
        }
        linker_options = {"link_time_optimization": True}
        kernel_obj = compile_cpp_and_link(
            src_code=src_code_list,
            object_code=obj_code_list,
            function_name="matmul_kernel",
            compiler_options=compiler_options,
            linker_options=linker_options,
        )
        assert kernel_obj is not None, "Internal error."
        self._matmul_kernels_cache[key] = kernel_obj, grid_iter

        if self._debug:
            t2 = time.perf_counter_ns()
            tt = (t2 - t1) / 1000.0  # ns -> us
            print(f"<== Compiled matmul in {tt:10.2f} us.\n")

        return kernel_obj, grid_iter
