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

from collections.abc import Sequence

import numpy as np

from nvmath.bindings import cudss
from nvmath.internal.ndbuffer import NDBuffer
from nvmath.internal.tensor_ifc_ndbuffer import NDBufferTensor
from nvmath.internal.tensor_wrapper import wrap_operand
from nvmath.internal.utils import device_ctx

COMPATIBLE_LHS_RHS_PACKAGES = {("scipy", "numpy"), ("cupyx", "cupy"), ("torch", "torch")}


class _BatchPtrData:
    """
    In batched cuDSS operations the library needs a device array of pointers (one
    per sample in the batch). When operands are reset between solves, because of
    the reset_operands()/reset_operands_unchecked() requirements, only the
    pointer values can change and the device array that holds them is/can be reused.
    A naive update path would construct transient numpy arrays and NDBuffer wrappers
    on every call (via wrap_operand -> NumpyTensor -> asndbuffer), which adds
    substantial Python-side overhead.
    _BatchPtrData mitigates this overhead by pre-allocating, at creation time, an
    NDBuffer wrapper backed by a numpy array on the host.
    When DirectSolver calls the update_cudss_*_batch_ptr functions below, it writes
    new ptr values directly into ``_host_ndbuf.data`` and calls ``_device_ptrs.copy_``
    with the already-existing ``_host_ndbuf``, skipping all
    intermediate object construction.

    Note: this wrapper is intentionally private because it is only used in this file
    and is transparent to the rest of the code. The DirectSolver class in fact
    stores _BatchPtrData instances in its ``resources_*_ptrs`` lists and passes
    them back to the update functions without knowing what these objects are.
    Since, for now, the only attributes of the underlying NDBufferTensor that the update
    functions need are ``.tensor`` and ``.device_id``, these are exposed
    as explicit properties.
    """

    __slots__ = ("_device_ptrs", "_host_ndbuf")

    def __init__(self, ptr: NDBufferTensor):
        assert isinstance(ptr, NDBufferTensor), "Internal error."
        self._device_ptrs = ptr
        # ptr is a 1-D device array with one element per batch sample,
        # so ptr.shape[0] is the batch count.
        self._host_ndbuf = NDBuffer.from_numpy(np.empty(ptr.shape[0], dtype=np.uint64))

    @property
    def tensor(self):
        return self._device_ptrs.tensor

    @property
    def device_id(self):
        return self._device_ptrs.device_id

    def update_from_host_ptrs(self, ptrs, stream_holder):
        """
        Write new pointer values into the host buffer and copy to device.
        Must be called inside a ``device_ctx`` for the target device.
        """
        self._host_ndbuf.data[:] = ptrs
        self._device_ptrs.tensor.copy_(self._host_ndbuf, stream=stream_holder, blocking=False)


# TODO: Add debug log messages.
def wrap_dense_tensors(tensor_type, native_operands):
    # Unchecked faster version of tensor_wrapper.wrap_operands().
    return tuple(tensor_type(o) for o in native_operands)


def wrap_sparse_tensors(tensor_type, attr_name_map, native_operands):
    # A more direct version of nvmath.sparse._internal.common_utils.wrap_sparse_operands().
    return tuple(tensor_type.create_from_tensor(o, attr_name_map=attr_name_map) for o in native_operands)


def create_cudss_dense_matrix(cuda_value_type, rhs):
    # The rhs should be a wrapped (TensorHolder object) matrix or vector in a memory space
    # consistent with the execution mode.
    shape = rhs.shape
    assert len(shape) <= 2, "Internal error."

    m = shape[0]
    n = shape[1] if len(shape) == 2 else 1

    # Only column-major is supported.
    ld = m

    rhs_ptr = cudss.matrix_create_dn(m, n, ld, rhs.data_ptr, cuda_value_type, cudss.Layout.COL_MAJOR)
    return [], [], rhs_ptr


def create_cudss_dense_explicit_batch(cuda_index_type, cuda_value_type, index_type, rhs, stream_holder):
    # The rhs should be a sequence of wrapped (TensorHolder object) matrix or vector in a
    # memory space consistent with the execution mode.
    assert isinstance(rhs, Sequence), "Internal error."

    batch_count = len(rhs)

    m = wrap_operand(np.array([r.shape[0] for r in rhs], dtype=index_type))
    n = wrap_operand(np.array([r.shape[1] if len(r.shape) == 2 else 1 for r in rhs], dtype=index_type))

    # Take device ID of the first matrix, since it's the same for all.
    device_id = rhs[0].device_id

    # The pointers must be a device array. We'll first create it on the host, wrap and
    # copy it to the device.
    p = wrap_operand(np.array([r.data_ptr for r in rhs], dtype=np.uint64))
    p = p.to(device_id=device_id, stream_holder=stream_holder)

    # LD is the same as M, since only column-major is supported.
    rhs_ptr = cudss.matrix_create_batch_dn(
        batch_count, m.data_ptr, n.data_ptr, m.data_ptr, p.data_ptr, cuda_index_type, cuda_value_type, cudss.Layout.COL_MAJOR
    )

    return [m, n], [_BatchPtrData(p)], rhs_ptr


def create_cudss_dense_implicit_batch(cuda_index_type, cuda_value_type, index_type, batch_indices, rhs, stream_holder):
    # The rhs should be a wrapped (TensorHolder object) NDarray 3D or greater dimension.
    assert len(rhs.shape) > 2, "Internal error."

    # Create a sequence of samples in the batch.
    # Explicitly unpack indices to avoid tensor[*b, ...] (support 3.10).
    unpacked_indices = [tuple(b) + (...,) for b in batch_indices]
    rhs_sequence = [rhs.tensor[u] for u in unpacked_indices]

    # Wrap the samples using the lightweight interface, since we've already checked.
    rhs_sequence = wrap_dense_tensors(rhs.__class__, rhs_sequence)

    return create_cudss_dense_explicit_batch(cuda_index_type, cuda_value_type, index_type, rhs_sequence, stream_holder)


def create_cudss_dense_wrapper(cuda_index_type, cuda_value_type, index_type, batch_indices, rhs, stream_holder):
    # A convenience function to forward to the right implementation.

    if isinstance(rhs, Sequence):
        return create_cudss_dense_explicit_batch(cuda_index_type, cuda_value_type, index_type, rhs, stream_holder)

    if len(rhs.shape) > 2:
        return create_cudss_dense_implicit_batch(
            cuda_index_type, cuda_value_type, index_type, batch_indices, rhs, stream_holder
        )

    return create_cudss_dense_matrix(cuda_value_type, rhs)


def update_cudss_dense_matrix_ptr(rhs_ptr, new_rhs):
    # The new rhs should be a wrapped (TensorHolder object) matrix or vector in a memory
    #   space consistent with the execution mode.
    # rhs_ptr is a cuDSS matrix ptr.
    assert len(new_rhs.shape) <= 2, "Internal error."
    cudss.matrix_set_values(rhs_ptr, new_rhs.data_ptr)

    return []


def update_cudss_dense_explicit_batch_ptr(rhs_ptr, new_rhs, stream_holder, existing_ptrs):
    # The new rhs should be a sequence of wrapped (TensorHolder object) matrix or vector
    # in a memory space consistent with the execution mode.
    assert isinstance(new_rhs, Sequence), "Internal error."

    with device_ctx(existing_ptrs[0].device_id):
        existing_ptrs[0].update_from_host_ptrs([r.data_ptr for r in new_rhs], stream_holder)


def update_cudss_dense_implicit_batch_ptr(rhs_ptr, batch_indices, new_rhs, stream_holder, existing_ptrs):
    # The new rhs should be a wrapped (TensorHolder object) NDarray 3D or greater dimension.
    assert len(new_rhs.shape) > 2, "Internal error."

    # Create a sequence of samples in the batch.
    # Explicitly unpack indices to avoid tensor[*b, ...] (support 3.10).
    unpacked_indices = [tuple(b) + (...,) for b in batch_indices]
    rhs_sequence = [new_rhs.tensor[u] for u in unpacked_indices]

    # Wrap the samples using the lightweight interface, since we've already checked.
    rhs_sequence = wrap_dense_tensors(new_rhs.__class__, rhs_sequence)

    update_cudss_dense_explicit_batch_ptr(rhs_ptr, rhs_sequence, stream_holder, existing_ptrs)


def update_cudss_dense_ptr_wrapper(*, existing_ptrs, rhs_ptr, batch_indices=None, new_rhs=None, stream_holder=None):
    # A convenience function to forward to the right update implementation.
    # For batched cases, existing_ptrs must be the list originally returned by the
    # corresponding create function; pointer values are copied into the existing
    # device buffers in place.  For non-batched cases existing_ptrs is unused.
    assert new_rhs is not None, "Internal error."

    if isinstance(new_rhs, Sequence):
        assert stream_holder is not None and existing_ptrs, "Internal error."
        update_cudss_dense_explicit_batch_ptr(rhs_ptr, new_rhs, stream_holder, existing_ptrs)
        return

    if len(new_rhs.shape) > 2:
        assert batch_indices is not None and stream_holder is not None and existing_ptrs, "Internal error."
        update_cudss_dense_implicit_batch_ptr(rhs_ptr, batch_indices, new_rhs, stream_holder, existing_ptrs)
        return

    update_cudss_dense_matrix_ptr(rhs_ptr, new_rhs)


def create_cudss_csr_matrix(cuda_index_type, cuda_value_type, matrix_type, matrix_view_type, lhs):
    # The lhs should be a wrapped (CSRTensorHolder object) matrix in a memory space
    # consistent with the execution mode.
    assert lhs.format_name == "CSR", "Internal error."

    shape = lhs.shape
    assert len(shape) == 2, "Internal error."

    m, n = shape
    nnz = lhs.values.size

    lhs_ptr = cudss.matrix_create_csr(
        m,
        n,
        nnz,
        lhs.crow_indices.data_ptr,
        0,
        lhs.col_indices.data_ptr,
        lhs.values.data_ptr,
        cuda_index_type,
        cuda_index_type,
        cuda_value_type,
        matrix_type,
        matrix_view_type,
        cudss.IndexBase.ZERO,
    )
    return [], [], lhs_ptr


def create_cudss_csr_explicit_batch(
    cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, lhs, stream_holder
):
    # The lhs should be a sequence of wrapped (CSRTensorHolder object) matrix in a memory
    # space consistent with the execution mode.
    assert isinstance(lhs, Sequence), "Internal error."

    batch_count = len(lhs)

    m = wrap_operand(np.array([o.shape[0] for o in lhs], dtype=index_type))
    n = wrap_operand(np.array([o.shape[1] for o in lhs], dtype=index_type))
    nnz = wrap_operand(np.array([o.values.size for o in lhs], dtype=index_type))

    # Take device ID of the first operator, since it's the same for all.
    device_id = lhs[0].device_id

    # The pointer arrays must be on device. We'll first create it on the host, wrap and
    # copy it to the device.
    crow_indices = wrap_operand(np.array([o.crow_indices.data_ptr for o in lhs], dtype=np.uint64))
    crow_indices = crow_indices.to(device_id=device_id, stream_holder=stream_holder)

    col_indices = wrap_operand(np.array([o.col_indices.data_ptr for o in lhs], dtype=np.uint64))
    col_indices = col_indices.to(device_id=device_id, stream_holder=stream_holder)

    values = wrap_operand(np.array([o.values.data_ptr for o in lhs], dtype=np.uint64))
    values = values.to(device_id=device_id, stream_holder=stream_holder)

    lhs_ptr = cudss.matrix_create_batch_csr(
        batch_count,
        m.data_ptr,
        n.data_ptr,
        nnz.data_ptr,
        crow_indices.data_ptr,
        0,
        col_indices.data_ptr,
        values.data_ptr,
        cuda_index_type,
        cuda_index_type,
        cuda_value_type,
        matrix_type,
        matrix_view_type,
        cudss.IndexBase.ZERO,
    )
    return (
        [m, n, nnz],
        [
            _BatchPtrData(crow_indices),
            _BatchPtrData(col_indices),
            _BatchPtrData(values),
        ],
        lhs_ptr,
    )


def create_cudss_csr_implicit_batch(
    cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, batch_indices, lhs, stream_holder
):
    # The lhs should be a wrapped CSRTensorHolder object in a memory space consistent with
    # the execution mode.
    assert lhs.num_dimensions > 2, "Internal error."

    # We can't use an elegant implementation like for the dense implicit batch, where we
    # index the samples in the tensor to form an explicit sequence and forward to the dense
    # explicit batch implementation. This is because we can't rely on having the native
    # LHS tensor after copying across memory spaces.

    # For this function, we assume that the batch_indices are consistent with the
    # specified lhs.
    batch_count = len(batch_indices)

    # The shape is constant for each sample.
    shape = lhs.shape[-2], lhs.shape[-1]
    m = wrap_operand(np.full(batch_count, shape[0], dtype=index_type))
    n = wrap_operand(np.full(batch_count, shape[1], dtype=index_type))

    # Explicitly unpack indices to avoid tensor[*b, ...] (support 3.10).
    unpacked_indices = [tuple(b) + (...,) for b in batch_indices]

    # Create a sequence of wrapped samples in the batch for each of the constituent dense
    # arrays.
    crow_indices_sequence = [lhs.crow_indices.__class__(lhs.crow_indices.tensor[u]) for u in unpacked_indices]
    col_indices_sequence = [lhs.col_indices.__class__(lhs.col_indices.tensor[u]) for u in unpacked_indices]
    values_sequence = [lhs.values.__class__(lhs.values.tensor[u]) for u in unpacked_indices]

    # The pointer arrays must be on device. We'll first create it on the host, wrap and
    # copy it to the device.
    crow_indices = wrap_operand(np.array([o.data_ptr for o in crow_indices_sequence], dtype=np.uint64))
    crow_indices = crow_indices.to(device_id=lhs.device_id, stream_holder=stream_holder)

    col_indices = wrap_operand(np.array([o.data_ptr for o in col_indices_sequence], dtype=np.uint64))
    col_indices = col_indices.to(device_id=lhs.device_id, stream_holder=stream_holder)

    values = wrap_operand(np.array([o.data_ptr for o in values_sequence], dtype=np.uint64))
    values = values.to(device_id=lhs.device_id, stream_holder=stream_holder)

    # The nnz should be the same for each sample, but we'll iterate anyway.
    nnz = wrap_operand(np.array([o.size for o in values_sequence], dtype=index_type))

    lhs_ptr = cudss.matrix_create_batch_csr(
        batch_count,
        m.data_ptr,
        n.data_ptr,
        nnz.data_ptr,
        crow_indices.data_ptr,
        0,
        col_indices.data_ptr,
        values.data_ptr,
        cuda_index_type,
        cuda_index_type,
        cuda_value_type,
        matrix_type,
        matrix_view_type,
        cudss.IndexBase.ZERO,
    )
    return (
        [m, n, nnz],
        [
            _BatchPtrData(crow_indices),
            _BatchPtrData(col_indices),
            _BatchPtrData(values),
        ],
        lhs_ptr,
    )


def create_cudss_batchedcsr_implicit_batch(
    cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, batch_indices, lhs, stream_holder
):
    # The lhs should be a wrapped CSRTensorHolder object containing an 3D BatchedCSR UST.
    assert lhs.num_dimensions == 3, "Internal error."

    # For this function, we assume that the batch_indices are consistent with the
    # specified lhs.
    batch_count = len(batch_indices)
    assert batch_count == lhs.shape[0], "Internal error."

    # The shape is constant for each sample.
    shape = lhs.shape[-2], lhs.shape[-1]
    m = wrap_operand(np.full(batch_count, shape[0], dtype=index_type))
    n = wrap_operand(np.full(batch_count, shape[1], dtype=index_type))

    # Create a sequence of wrapped samples in the batch for each of the constituent dense
    # arrays.
    crow_indices_sequence = []
    col_indices_sequence = []
    values_sequence = []
    for b in range(batch_count):
        s = shape[0]
        # Take a copy, since we are offsetting crow_indices to start at 0.
        crow_b = lhs.crow_indices.tensor[b * s : (b + 1) * s + 1].copy()
        crow_indices_sequence.append(lhs.crow_indices.__class__(crow_b))
        col_b = lhs.col_indices.tensor[crow_b[0] : crow_b[-1]]
        col_indices_sequence.append(lhs.col_indices.__class__(col_b))
        values_b = lhs.values.tensor[crow_b[0] : crow_b[-1]]
        values_sequence.append(lhs.values.__class__(values_b))
        # Offset the crow_indices after col_indices and values are sliced.
        crow_b -= crow_b[0]

    # The pointer arrays must be on device. We'll first create it on the host, wrap and
    # copy it to the device.
    crow_indices = wrap_operand(np.array([o.data_ptr for o in crow_indices_sequence], dtype=np.uint64))
    crow_indices = crow_indices.to(device_id=lhs.device_id, stream_holder=stream_holder)

    col_indices = wrap_operand(np.array([o.data_ptr for o in col_indices_sequence], dtype=np.uint64))
    col_indices = col_indices.to(device_id=lhs.device_id, stream_holder=stream_holder)

    values = wrap_operand(np.array([o.data_ptr for o in values_sequence], dtype=np.uint64))
    values = values.to(device_id=lhs.device_id, stream_holder=stream_holder)

    # The nnz can vary for each sample in UST BatchedCSR, it is important to iterate.
    nnz = wrap_operand(np.array([o.size for o in values_sequence], dtype=index_type))

    lhs_ptr = cudss.matrix_create_batch_csr(
        batch_count,
        m.data_ptr,
        n.data_ptr,
        nnz.data_ptr,
        crow_indices.data_ptr,
        0,
        col_indices.data_ptr,
        values.data_ptr,
        cuda_index_type,
        cuda_value_type,
        matrix_type,
        matrix_view_type,
        cudss.IndexBase.ZERO,
    )
    # Since we have taken a copy to offset crow_indices, we need to hold on to it.
    references = [m, n, nnz, crow_indices_sequence, col_indices_sequence, values_sequence, crow_indices, col_indices, values]
    return references, lhs_ptr


def create_cudss_csr_wrapper(
    cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, batch_indices, lhs, stream_holder
):
    if isinstance(lhs, Sequence):
        return create_cudss_csr_explicit_batch(
            cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, lhs, stream_holder
        )

    if len(lhs.shape) > 2:
        return create_cudss_csr_implicit_batch(
            cuda_index_type, cuda_value_type, index_type, matrix_type, matrix_view_type, batch_indices, lhs, stream_holder
        )

    return create_cudss_csr_matrix(cuda_index_type, cuda_value_type, matrix_type, matrix_view_type, lhs)


def update_cudss_csr_matrix_ptr(lhs_ptr, new_lhs):
    # The new lhs should be a wrapped (CSRTensorHolder object) matrix in a memory space
    # consistent with the execution mode.
    # lhs_ptr is a cuDSS CSR matrix ptr.
    assert len(new_lhs.shape) == 2, "Internal error."
    cudss.matrix_set_csr_pointers(
        lhs_ptr, new_lhs.crow_indices.data_ptr, 0, new_lhs.col_indices.data_ptr, new_lhs.values.data_ptr
    )
    return []


def update_cudss_csr_explicit_batch_ptr(lhs_ptr, new_lhs, stream_holder, existing_ptrs):
    # The new lhs should be a sequence of wrapped (CSRTensorHolder object) matrix in a
    # memory space consistent with the execution mode.
    assert isinstance(new_lhs, Sequence), "Internal error."

    with device_ctx(existing_ptrs[0].device_id):
        existing_ptrs[0].update_from_host_ptrs([o.crow_indices.data_ptr for o in new_lhs], stream_holder)
        existing_ptrs[1].update_from_host_ptrs([o.col_indices.data_ptr for o in new_lhs], stream_holder)
        existing_ptrs[2].update_from_host_ptrs([o.values.data_ptr for o in new_lhs], stream_holder)


def update_cudss_csr_implicit_batch_ptr(lhs_ptr, batch_indices, new_lhs, stream_holder, existing_ptrs):
    # The new lhs should be a wrapped (CSRTensorHolder object) tensor 3D or greater
    # dimension in a memory space consistent with the execution mode.
    assert new_lhs.num_dimensions > 2, "Internal error."

    # We can't use an elegant implementation like for the dense implicit batch, where we
    # index the samples in the tensor to form an explicit sequence and forward to the dense
    # explicit batch implementation. This is because we can't rely on having the native
    # LHS tensor after copying across memory spaces.

    # Explicitly unpack indices to avoid tensor[*b, ...] (support 3.10).
    unpacked_indices = [tuple(b) + (...,) for b in batch_indices]

    # Create a sequence of wrapped samples in the batch for each of the constituent dense
    # arrays.
    crow_indices_sequence = [new_lhs.crow_indices.__class__(new_lhs.crow_indices.tensor[u]) for u in unpacked_indices]
    col_indices_sequence = [new_lhs.col_indices.__class__(new_lhs.col_indices.tensor[u]) for u in unpacked_indices]
    values_sequence = [new_lhs.values.__class__(new_lhs.values.tensor[u]) for u in unpacked_indices]

    with device_ctx(existing_ptrs[0].device_id):
        existing_ptrs[0].update_from_host_ptrs([o.data_ptr for o in crow_indices_sequence], stream_holder)
        existing_ptrs[1].update_from_host_ptrs([o.data_ptr for o in col_indices_sequence], stream_holder)
        existing_ptrs[2].update_from_host_ptrs([o.data_ptr for o in values_sequence], stream_holder)


def update_cudss_csr_ptr_wrapper(*, existing_ptrs, lhs_ptr, batch_indices=None, new_lhs=None, stream_holder=None):
    # A convenience function to forward to the right update implementation.
    # For batched cases, existing_ptrs must be the list originally returned by the
    # corresponding create function; pointer values are copied into the existing
    # device buffers in place.  For non-batched cases existing_ptrs is unused.
    assert new_lhs is not None, "Internal error."

    if isinstance(new_lhs, Sequence):
        assert stream_holder is not None and existing_ptrs, "Internal error."
        update_cudss_csr_explicit_batch_ptr(lhs_ptr, new_lhs, stream_holder, existing_ptrs)
        return

    if len(new_lhs.shape) > 2:
        assert batch_indices is not None and stream_holder is not None and existing_ptrs, "Internal error."
        update_cudss_csr_implicit_batch_ptr(lhs_ptr, batch_indices, new_lhs, stream_holder, existing_ptrs)
        return

    update_cudss_csr_matrix_ptr(lhs_ptr, new_lhs)
