# 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 ._emitter import emit_matmul
from ._jit import compile_cpp_and_link
from ._semiring import (
    epilog_impl,
    prolog_a_impl,
    prolog_b_impl,
    prolog_c_impl,
    prolog_epilog_decl,
    semiring_add_impl,
    semiring_atomic_add_impl,
    semiring_mul_impl,
    semiring_ops_decl,
)
from ._utils import COMPLEX_PRECISION, Cache, type_str

# Maps matmul(a,b,c) to pre-defined kernel source code. All kernels should adhere to the
# parameter passing conventions for each of the three UST parameters.
_tensor_matmul_kernels = {
    # (format_a, format_b, format_c, transpose_a, transpose_b)
    ("DIAI", "DenseVector", "DenseVector", False, False): """
extern "C" __global__ void matmul(
    const POS* __restrict__ pos0,
    const CRD* __restrict__ diags,
    const VAL* __restrict__ A,
    unsigned long m,
    unsigned long n,
    unsigned long l0,
    unsigned long l1,
    const VAL *__restrict__ B,
    VAL *__restrict__ C) {
  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 = prolog_c(static_cast<CTP>(C[i]));
    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 = add(acc, mul(prolog_a(static_cast<CTP>(A[di])),
                           prolog_b(static_cast<CTP>(B[j]))));
      }
    }
    C[i] = static_cast<VAL>(epilog(acc));
  }
}
""",
    ("DIAJ", "DenseVector", "DenseVector", False, False): """
extern "C" __global__ void matmul(
    const POS* __restrict__ pos0,
    const CRD* __restrict__ diags,
    const VAL* __restrict__ A,
    unsigned long m,
    unsigned long n,
    unsigned long l0,
    unsigned long l1,
    const VAL *__restrict__ B,
    VAL *__restrict__ C) {
  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 = prolog_c(static_cast<CTP>(C[i]));
    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 = add(acc, mul(static_cast<CTP>(prolog_a(A[di + j])),
                           static_cast<CTP>(prolog_b(B[j]))));
      }
    }
    C[i] = static_cast<VAL>(epilog(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,
        epilog=None,
        semiring=None,
        transpose_a=False,
        transpose_b=False,
        arch=None,
    ):
        """
        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 semiring is None:
            user_mul, user_add, user_atomic_add = None, None, None
        else:
            user_mul, user_add, user_atomic_add = semiring.get("mul"), semiring.get("add"), semiring.get("atomic_add")
        if prologs is None:
            user_prolog_a, user_prolog_b, user_prolog_c = None, None, None
        else:
            user_prolog_a, user_prolog_b, user_prolog_c = prologs.get("a"), prologs.get("b"), prologs.get("c")
        # Try cache lookup.
        key = (
            format_a,
            format_b,
            format_c,
            vtp,
            itp,
            compute_type,
            user_mul,
            user_add,
            user_atomic_add,
            user_prolog_a,
            user_prolog_b,
            user_prolog_c,
            epilog,
            transpose_a,
            transpose_b,
            arch,
        )
        kernel = self._matmul_kernels_cache.get(key)
        if kernel is None:
            # Nothing in cache, try kernel lookup.
            typeless_key = (format_a, format_b, format_c, transpose_a, transpose_b)
            code = _tensor_matmul_kernels.get(typeless_key)

            # Prepare prolog.
            ctp = type_str(compute_type)
            if a.dtype in ["float8_e4m3fn", "float8_e5m2"]:
                prolog = f"#include <cuda_fp8.h>\nusing ATP = {ctp};\nusing CTP = {ctp};\n"  # TODO: ATP=vtp
            elif a.dtype == "float16":
                prolog = f"#include <cuda_fp16.h>\nusing ATP = {vtp};\nusing CTP = {ctp};\n"
            elif a.dtype == "bfloat16":
                prolog = f"#include <cuda_bf16.h>\nusing ATP = {vtp};\nusing CTP = {ctp};\n"
            elif a.dtype in ["complex64", "complex128"]:
                prolog = f"#include <cuda/std/complex>\nusing ATP = {COMPLEX_PRECISION[a.dtype]};\nusing CTP = {ctp};\n"
            else:
                prolog = f"using ATP = {vtp};\nusing 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}\nusing VAL = {vtp};\nusing POS = {itp};\n"
                    f"using CRD = {itp};\n{semiring_ops_decl}{prolog_epilog_decl}{code}"
                )
            else:
                if user_prolog_c is not None or epilog is not None:
                    raise NotImplementedError(
                        "The code generation (emitter) path does not currently support \
prolog for the `c` operand or epilog."
                    )
                # (2) Sparse emitter. Emits proper types.
                source = emit_matmul(a, b, c, ctp, transpose_a, transpose_b)

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

            # User-provided or default semiring mul operation.
            src_code_list, obj_code_list = _set_user_or_default_code(
                user_code=user_mul,
                default_code=semiring_mul_impl,
                default_prolog=prolog,
                src_code_list=src_code_list,
                obj_code_list=obj_code_list,
            )

            # User-provided or default semiring add operation.
            src_code_list, obj_code_list = _set_user_or_default_code(
                user_code=user_add,
                default_code=semiring_add_impl,
                default_prolog=prolog,
                src_code_list=src_code_list,
                obj_code_list=obj_code_list,
            )

            # User-provided or default semiring atomic add operation.
            src_code_list, obj_code_list = _set_user_or_default_code(
                user_code=user_atomic_add,
                default_code=semiring_atomic_add_impl,
                default_prolog=prolog,
                src_code_list=src_code_list,
                obj_code_list=obj_code_list,
            )

            # User-provided or default prolog operation on A.
            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,
            )

            # User-provided or default prolog operation on B.
            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,
            )

            # User-provided or default prolog operation on C.
            src_code_list, obj_code_list = _set_user_or_default_code(
                user_code=user_prolog_c,
                default_code=prolog_c_impl,
                default_prolog=prolog,
                src_code_list=src_code_list,
                obj_code_list=obj_code_list,
            )

            # User-provided or default epilog operation.
            src_code_list, obj_code_list = _set_user_or_default_code(
                user_code=epilog,
                default_code=epilog_impl,
                default_prolog=prolog,
                src_code_list=src_code_list,
                obj_code_list=obj_code_list,
            )

            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 = compile_cpp_and_link(
                src_code=src_code_list,
                object_code=obj_code_list,
                function_name="matmul",
                compiler_options=compiler_options,
                linker_options=linker_options,
            )
            assert kernel is not None, "Internal error."
            self._matmul_kernels_cache[key] = kernel

            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
