#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
#  See https://llvm.org/LICENSE.txt for license information.
#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# This file contains the utilities to process sparse tensor outputs.

from typing import Callable, Dict, Sequence, Tuple
import ctypes
import functools
import numpy as np
import os

# Import MLIR related modules.
from mlir import execution_engine
from mlir import ir
from mlir import runtime
from mlir.dialects import sparse_tensor

from . import mlir_sparse_compiler

# Type aliases for type annotation.
_SupportFunc = Callable[..., None]
_SupportFuncLocator = Callable[[np.dtype], Tuple[_SupportFunc, _SupportFunc]]

# The name for the environment variable that provides the full path for the
# supporting library.
_SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
# The default supporting library if the environment variable is not provided.
_DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"

# The JIT compiler optimization level.
_OPT_LEVEL = 2
# The entry point to the JIT compiled program.
_ENTRY_NAME = "main"


@functools.lru_cache()
def _get_support_lib_name() -> str:
    """Gets the string name for the supporting C shared library."""
    return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)


@functools.lru_cache()
def _get_sparse_compiler() -> mlir_sparse_compiler.SparseCompiler:
    """Gets the MLIR sparse compiler with default setting."""
    return mlir_sparse_compiler.SparseCompiler(
        options="", opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()]
    )


def _record_support_funcs(
    ty: np.dtype,
    to_func: _SupportFunc,
    from_func: _SupportFunc,
    ty_to_funcs: Dict[np.dtype, Tuple[_SupportFunc, _SupportFunc]],
) -> None:
    """Records the two supporting functions for a given data type."""
    to_func.restype = ctypes.c_void_p
    from_func.restype = ctypes.c_void_p
    ty_to_funcs[ty] = (to_func, from_func)


@functools.lru_cache()
def _get_support_func_locator() -> _SupportFuncLocator:
    """Constructs a function to locate the supporting functions for a data type.

    Loads the supporting C shared library with the needed routines. Constructs a
    dictionary from the supported data types to the routines for the data types,
    and then a function to look up the dictionary for a given data type.

    The name of the supporting C shared library is either provided by an
    an environment variable or a default value.

    Returns:
      The function to look up the supporting functions for a given data type.

    Raises:
      OSError: If there is any problem in loading the shared library.
      ValueError: If the shared library doesn't contain the needed routines.
    """
    # This raises OSError exception if there is any problem in loading the shared
    # library.
    c_lib = ctypes.CDLL(_get_support_lib_name())

    type_to_funcs = {}
    try:
        support_types = [
            (
                np.int8,
                c_lib.convertToMLIRSparseTensorI8,
                c_lib.convertFromMLIRSparseTensorI8,
            ),
            (
                np.int16,
                c_lib.convertToMLIRSparseTensorI16,
                c_lib.convertFromMLIRSparseTensorI16,
            ),
            (
                np.int32,
                c_lib.convertToMLIRSparseTensorI32,
                c_lib.convertFromMLIRSparseTensorI32,
            ),
            (
                np.int64,
                c_lib.convertToMLIRSparseTensorI64,
                c_lib.convertFromMLIRSparseTensorI64,
            ),
            (
                np.float16,
                c_lib.convertToMLIRSparseTensorF16,
                c_lib.convertFromMLIRSparseTensorF16,
            ),
            (
                np.float32,
                c_lib.convertToMLIRSparseTensorF32,
                c_lib.convertFromMLIRSparseTensorF32,
            ),
            (
                np.float64,
                c_lib.convertToMLIRSparseTensorF64,
                c_lib.convertFromMLIRSparseTensorF64,
            ),
            (
                np.complex64,
                c_lib.convertToMLIRSparseTensorC32,
                c_lib.convertFromMLIRSparseTensorC32,
            ),
            (
                np.complex128,
                c_lib.convertToMLIRSparseTensorC64,
                c_lib.convertFromMLIRSparseTensorC64,
            ),
        ]
    except Exception as e:
        raise ValueError(f"Missing supporting function: {e}") from e
    for i, info in enumerate(support_types):
        _record_support_funcs(info[0], info[1], info[2], type_to_funcs)

    def get_support_funcs(ty: np.dtype):
        funcs = type_to_funcs[ty]
        assert funcs is not None
        return funcs

    return get_support_funcs


def sparse_tensor_to_coo_tensor(
    sparse_tensor: ctypes.c_void_p,
    dtype: np.dtype,
) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
    """Converts an MLIR sparse tensor to a COO-flavored format tensor.

    Args:
       sparse_tensor: A ctypes.c_void_p to the MLIR sparse tensor descriptor.
       dtype: The numpy data type for the tensor elements.

    Returns:
      A tuple that contains the following values for the COO-flavored format
      tensor:
      rank: An integer for the rank of the tensor.
      nse: An integer for the number of non-zero values in the tensor.
      shape: A 1D numpy array of integers, for the shape of the tensor.
      values: A 1D numpy array, for the non-zero values in the tensor.
      indices: A 2D numpy array of integers, representing the indices for the
        non-zero values in the tensor.

    Raises:
      OSError: If there is any problem in loading the shared library.
      ValueError: If the shared library doesn't contain the needed routines.
    """
    convert_from = _get_support_func_locator()(dtype)[1]
    rank = ctypes.c_ulonglong(0)
    nse = ctypes.c_ulonglong(0)
    shape = ctypes.POINTER(ctypes.c_ulonglong)()

    values = ctypes.POINTER(runtime.as_ctype(np.dtype(dtype)))()
    indices = ctypes.POINTER(ctypes.c_ulonglong)()
    convert_from(
        sparse_tensor,
        ctypes.byref(rank),
        ctypes.byref(nse),
        ctypes.byref(shape),
        ctypes.byref(values),
        ctypes.byref(indices),
    )

    # Convert the returned values to the corresponding numpy types.
    shape = np.ctypeslib.as_array(shape, shape=[rank.value])
    values = runtime.to_numpy(np.ctypeslib.as_array(values, shape=[nse.value]))
    indices = np.ctypeslib.as_array(indices, shape=[nse.value, rank.value])
    return rank.value, nse.value, shape, values, indices


def coo_tensor_to_sparse_tensor(
    np_shape: np.ndarray,
    np_values: np.ndarray,
    np_indices: np.ndarray,
    np_perm: np.ndarray,
    np_sparse: np.ndarray,
) -> int:
    """Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.

    Args:
       np_shape: A 1D numpy array of integers, for the shape of the tensor.
       np_values: A 1D numpy array, for the non-zero values in the tensor.
       np_indices: A 2D numpy array of integers, representing the indices for the
         non-zero values in the tensor.
       np_perm: A 1D numpy array of integers, representing the storage ordering
         for the dimensions.
       np_sparse: A 1D numpy array of uint8, representing the sparsity values
         for the dimensions.

    Returns:
       An integer for the non-null ctypes.c_void_p to the MLIR sparse tensor
       descriptor.

    Raises:
      OSError: If there is any problem in loading the shared library.
      ValueError: If the shared library doesn't contain the needed routines.
    """

    r = len(np_shape)
    rank = ctypes.c_ulonglong(r)
    nse = ctypes.c_ulonglong(len(np_values))
    shape = np_shape.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
    values = np_values.ctypes.data_as(
        ctypes.POINTER(runtime.as_ctype(np.dtype(np_values.dtype)))
    )
    indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))

    perm = np_perm.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
    sparse = np_sparse.ctypes.data_as(ctypes.POINTER(ctypes.c_uint8))

    convert_to = _get_support_func_locator()(np_values.dtype.type)[0]
    ptr = convert_to(rank, nse, shape, values, indices, perm, sparse)
    assert ptr is not None, "Problem with calling convertToMLIRSparseTensorF64"
    return ptr


def compile_and_build_engine(module: ir.Module) -> execution_engine.ExecutionEngine:
    """Compiles an MLIR module and builds a JIT execution engine.

    Args:
      module: The MLIR module.

    Returns:
      A JIT execution engine for the MLIR module.

    """
    return _get_sparse_compiler().compile_and_jit(module)


class _SparseTensorDescriptor(ctypes.Structure):
    """A C structure for an MLIR sparse tensor."""

    _fields_ = [
        # A pointer for the MLIR sparse tensor storage.
        ("storage", ctypes.POINTER(ctypes.c_ulonglong)),
        # An MLIR MemRef descriptor for the shape of the sparse tensor.
        ("shape", runtime.make_nd_memref_descriptor(1, ctypes.c_ulonglong)),
    ]


def _output_one_dim(dim: int, rank: int, shape: str, type: str) -> str:
    """Produces the MLIR text code to output the size for the given dimension."""
    return f"""
  %c{dim} = arith.constant {dim} : index
  %d{dim} = tensor.dim %t, %c{dim} : tensor<{shape}x{type}, #enc>
  memref.store %d{dim}, %b[%c{dim}] : memref<{rank}xindex>
"""


# TODO: With better support from MLIR, we may improve the current implementation
# by doing the following:
# (1) Use Python code to generate the kernel instead of doing MLIR text code
#     stitching.
# (2) Use scf.for instead of an unrolled loop to write out the dimension sizes
#     when tensor.dim supports non-constant dimension value.
def _get_create_sparse_tensor_kernel(
    sparsity_codes: Sequence[sparse_tensor.DimLevelType], type: str
) -> str:
    """Creates an MLIR text kernel to contruct a sparse tensor from a file.

    The kernel returns a _SparseTensorDescriptor structure.
    """
    rank = len(sparsity_codes)

    # Use ? to represent a dimension in the dynamic shape string representation.
    shape = "x".join(map(lambda d: "?", range(rank)))

    # Convert the encoded sparsity values to a string representation.
    sparsity = ", ".join(
        map(lambda s: '"compressed"' if s.value else '"dense"', sparsity_codes)
    )

    # Get the MLIR text code to write the dimension sizes to the output buffer.
    output_dims = "\n".join(
        map(lambda d: _output_one_dim(d, rank, shape, type), range(rank))
    )

    # Return the MLIR text kernel.
    return f"""
!Ptr = !llvm.ptr<i8>
#enc = #sparse_tensor.encoding<{{
  lvlTypes = [ {sparsity} ]
}}>
func.func @{_ENTRY_NAME}(%filename: !Ptr) -> (tensor<{shape}x{type}, #enc>, memref<{rank}xindex>)
attributes {{ llvm.emit_c_interface }} {{
  %t = sparse_tensor.new %filename : !Ptr to tensor<{shape}x{type}, #enc>
  %b = memref.alloc() : memref<{rank}xindex>
  {output_dims}
  return %t, %b : tensor<{shape}x{type}, #enc>, memref<{rank}xindex>
}}"""


def create_sparse_tensor(
    filename: str, sparsity: Sequence[sparse_tensor.DimLevelType], type: str
) -> Tuple[ctypes.c_void_p, np.ndarray]:
    """Creates an MLIR sparse tensor from the input file.

    Args:
      filename: A string for the name of the file that contains the tensor data in
        a COO-flavored format.
      sparsity: A sequence of DimLevelType values, one for each dimension of the
        tensor.

    Returns:
      A Tuple containing the following values:
      storage: A ctypes.c_void_p for the MLIR sparse tensor storage.
      shape: A 1D numpy array of integers, for the shape of the tensor.

    Raises:
      OSError: If there is any problem in loading the supporting C shared library.
      ValueError:  If the shared library doesn't contain the needed routine.
    """
    with ir.Context() as ctx, ir.Location.unknown():
        module = _get_create_sparse_tensor_kernel(sparsity, type)
        module = ir.Module.parse(module)
        engine = compile_and_build_engine(module)

    # A sparse tensor descriptor to receive the kernel result.
    c_tensor_desc = _SparseTensorDescriptor()
    # Convert the filename to a byte stream.
    c_filename = ctypes.c_char_p(bytes(filename, "utf-8"))

    arg_pointers = [
        ctypes.byref(ctypes.pointer(c_tensor_desc)),
        ctypes.byref(c_filename),
    ]

    # Invoke the execution engine to run the module and return the result.
    engine.invoke(_ENTRY_NAME, *arg_pointers)
    shape = runtime.ranked_memref_to_numpy(ctypes.pointer(c_tensor_desc.shape))
    return c_tensor_desc.storage, shape


# TODO: With better support from MLIR, we may improve the current implementation
# by using Python code to generate the kernel instead of doing MLIR text code
# stitching.
def _get_output_sparse_tensor_kernel(
    sparsity_codes: Sequence[sparse_tensor.DimLevelType], type: str
) -> str:
    """Creates an MLIR text kernel to output a sparse tensor to a file.

    The kernel returns void.
    """
    rank = len(sparsity_codes)

    # Use ? to represent a dimension in the dynamic shape string representation.
    shape = "x".join(map(lambda d: "?", range(rank)))

    # Convert the encoded sparsity values to a string representation.
    sparsity = ", ".join(
        map(lambda s: '"compressed"' if s.value else '"dense"', sparsity_codes)
    )

    # Return the MLIR text kernel.
    return f"""
!Ptr = !llvm.ptr<i8>
#enc = #sparse_tensor.encoding<{{
  lvlTypes = [ {sparsity} ]
}}>
func.func @{_ENTRY_NAME}(%t: tensor<{shape}x{type}, #enc>, %filename: !Ptr)
attributes {{ llvm.emit_c_interface }} {{
  sparse_tensor.out %t, %filename : tensor<{shape}x{type}, #enc>, !Ptr
  func.return
}}"""


def output_sparse_tensor(
    tensor: ctypes.c_void_p,
    filename: str,
    sparsity: Sequence[sparse_tensor.DimLevelType],
    type: str,
) -> None:
    """Outputs an MLIR sparse tensor to the given file.

    Args:
      tensor: A C pointer to the MLIR sparse tensor.
      filename: A string for the name of the file that contains the tensor data in
        a COO-flavored format.
      sparsity: A sequence of DimLevelType values, one for each dimension of the
        tensor.
      type: The MLIR string for the data type.

    Raises:
      OSError: If there is any problem in loading the supporting C shared library.
      ValueError:  If the shared library doesn't contain the needed routine.
    """
    with ir.Context() as ctx, ir.Location.unknown():
        module = _get_output_sparse_tensor_kernel(sparsity, type)
        module = ir.Module.parse(module)
        engine = compile_and_build_engine(module)

    # Convert the filename to a byte stream.
    c_filename = ctypes.c_char_p(bytes(filename, "utf-8"))

    arg_pointers = [
        ctypes.byref(ctypes.cast(tensor, ctypes.c_void_p)),
        ctypes.byref(c_filename),
    ]

    # Invoke the execution engine to run the module and return the result.
    engine.invoke(_ENTRY_NAME, *arg_pointers)