# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import ctypes
from functools import wraps
import gc
import pytest

import pyarrow as pa
from pyarrow.vendored.version import Version

# Marks all of the tests in this module
# Ignore these with pytest ... -m 'not numpy'
pytestmark = pytest.mark.numpy
np = pytest.importorskip("numpy")


def requires_numpy_version(min_version):
    return pytest.mark.skipif(
        Version(np.__version__) < Version(min_version),
        reason=f"Test requires numpy {min_version} or later",
    )


def PyCapsule_IsValid(capsule, name):
    return ctypes.pythonapi.PyCapsule_IsValid(ctypes.py_object(capsule), name) == 1


def check_dlpack_export(arr, expected_arr):
    with pytest.warns(DeprecationWarning, match="unversioned DLPack capsule"):
        DLTensor = arr.__dlpack__()
    assert PyCapsule_IsValid(DLTensor, b"dltensor") is True

    result = np.from_dlpack(arr)
    np.testing.assert_array_equal(result, expected_arr, strict=True)

    assert arr.__dlpack_device__() == (1, 0)


class DLPackForwarder:
    """Forward ``__dlpack__`` to a wrapped object with forced keyword arguments.

    Consumers such as ``np.from_dlpack`` do not expose every ``__dlpack__``
    keyword, so this makes them reachable from a consumer's point of view.
    """

    def __init__(self, obj, **forced):
        self._obj = obj
        self._forced = forced

    def __dlpack__(self, **kwargs):
        return self._obj.__dlpack__(**{**kwargs, **self._forced})

    def __dlpack_device__(self):
        return self._obj.__dlpack_device__()


def check_bytes_allocated(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        gc.collect()
        allocated_bytes = pa.total_allocated_bytes()
        try:
            return f(*args, **kwargs)
        finally:
            assert pa.total_allocated_bytes() == allocated_bytes
    return wrapper


@check_bytes_allocated
@pytest.mark.parametrize(
    ('value_type', 'np_type_str'),
    [
        (pa.uint8(), "uint8"),
        (pa.uint16(), "uint16"),
        (pa.uint32(), "uint32"),
        (pa.uint64(), "uint64"),
        (pa.int8(), "int8"),
        (pa.int16(), "int16"),
        (pa.int32(), "int32"),
        (pa.int64(), "int64"),
        (pa.float16(), "float16"),
        (pa.float32(), "float32"),
        (pa.float64(), "float64"),
    ]
)
def test_dlpack(value_type, np_type_str):
    expected = np.array([1, 2, 3], dtype=np.dtype(np_type_str))
    arr = pa.array(expected, type=value_type)
    check_dlpack_export(arr, expected)

    t = pa.Tensor.from_numpy(expected)
    check_dlpack_export(t, expected)

    arr_sliced = arr.slice(1, 1)
    expected = np.array([2], dtype=np.dtype(np_type_str))
    check_dlpack_export(arr_sliced, expected)

    arr_sliced = arr.slice(0, 1)
    expected = np.array([1], dtype=np.dtype(np_type_str))
    check_dlpack_export(arr_sliced, expected)

    arr_sliced = arr.slice(1)
    expected = np.array([2, 3], dtype=np.dtype(np_type_str))
    check_dlpack_export(arr_sliced, expected)

    arr_zero = pa.array([], type=value_type)
    expected = np.array([], dtype=np.dtype(np_type_str))
    check_dlpack_export(arr_zero, expected)

    t = pa.Tensor.from_numpy(expected)
    check_dlpack_export(t, expected)


@check_bytes_allocated
@pytest.mark.parametrize('np_type',
                         [np.uint8, np.uint16, np.uint32, np.uint64,
                          np.int8, np.int16, np.int32, np.int64,
                          np.float16, np.float32, np.float64,])
def test_tensor_dlpack(np_type):
    arr = np.array([1, 2, 3, 4, 5, 6, 1, 1])
    expected = np.array(arr, dtype=np_type).reshape((2, 2, 2), order='C')
    t = pa.Tensor.from_numpy(expected)
    check_dlpack_export(t, expected)

    expected = np.array(arr, dtype=np_type).reshape((2, 2, 2), order='F')
    t = pa.Tensor.from_numpy(expected)
    check_dlpack_export(t, expected)


def multidim_arrays():
    np_arr = np.arange(12, dtype=np.int32).reshape(3, 2, 2)
    values = pa.array(np_arr.ravel(), type=pa.int32())
    nested_list = pa.FixedSizeListArray.from_arrays(
        pa.FixedSizeListArray.from_arrays(values, 2), 2)
    return [
        pytest.param(nested_list, np_arr, id="nested_fixed_size_list"),
        pytest.param(
            pa.FixedShapeTensorArray.from_numpy_ndarray(np_arr),
            np_arr,
            id="fixed_shape_tensor",
        ),
    ]


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize(('arr', 'expected'), multidim_arrays())
def test_array_to_tensor_dlpack(arr, expected):
    tensor = arr.to_tensor()
    # A Tensor sharing an Array buffer is immutable, so it can only be exported
    # through the versioned DLPack protocol.
    assert not tensor.is_mutable
    result = np.from_dlpack(DLPackForwarder(tensor, max_version=(1, 0)))
    np.testing.assert_array_equal(result, expected, strict=True)
    assert tensor.__dlpack_device__() == (1, 0)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_fixed_shape_tensor_array_dlpack_permuted():
    # A non-trivial permutation makes to_tensor() produce a non-row-major
    # tensor: each row-major [3, 2] block is exposed as a logical [2, 3] cell.
    storage = pa.FixedSizeListArray.from_arrays(
        pa.array(range(24), type=pa.int32()), 6)
    arr = pa.ExtensionArray.from_storage(
        pa.fixed_shape_tensor(pa.int32(), [3, 2], permutation=[1, 0]), storage)

    tensor = arr.to_tensor()
    assert tensor.shape == (4, 2, 3)
    assert not tensor.is_contiguous

    # expected[i, j, k] == i * 6 + k * 2 + j (numpy is only the DLPack consumer)
    expected = np.arange(24, dtype=np.int32).reshape(4, 3, 2).transpose(0, 2, 1)
    result = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 0)))
    np.testing.assert_array_equal(result, expected, strict=True)
    assert arr.__dlpack_device__() == (1, 0)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_fixed_shape_tensor_scalar_dlpack():
    np_arr = np.arange(12, dtype=np.int32).reshape(3, 2, 2)
    arr = pa.FixedShapeTensorArray.from_numpy_ndarray(np_arr)

    scalar = arr[1]
    assert isinstance(scalar, pa.FixedShapeTensorScalar)
    # __dlpack_device__ reads the storage array's device, without building a Tensor.
    assert scalar.__dlpack_device__() == (1, 0)

    result = np.from_dlpack(DLPackForwarder(scalar, max_version=(1, 0)))
    np.testing.assert_array_equal(result, np_arr[1], strict=True)


def multidim_arrays_with_nulls():
    np_arr = np.arange(6, dtype=np.int32).reshape(3, 2)
    # Masked entries keep defined values in the child array, so the tensor
    # contents stay fully predictable.
    nested_list = pa.FixedSizeListArray.from_arrays(
        pa.array(np_arr.ravel(), type=pa.int32()), 2,
        mask=pa.array([False, True, False]))
    return [
        pytest.param(nested_list, np_arr, id="fixed_size_list"),
        pytest.param(
            pa.ExtensionArray.from_storage(
                pa.fixed_shape_tensor(pa.int32(), [2]), nested_list),
            np_arr,
            id="fixed_shape_tensor",
        ),
    ]


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize(('arr', 'expected'), multidim_arrays_with_nulls())
def test_array_to_tensor_dlpack_nulls(arr, expected):
    with pytest.raises(pa.ArrowInvalid, match="Array contains nulls"):
        arr.to_tensor()

    tensor = arr.to_tensor(allow_nulls=True)
    result = np.from_dlpack(DLPackForwarder(tensor, max_version=(1, 0)))
    np.testing.assert_array_equal(result, expected, strict=True)


def dlpack_objects():
    arr = pa.array([1, 2, 3], type=pa.int32())
    return [
        pytest.param(arr, id="array"),
        pytest.param(arr.slice(1), id="sliced_array"),
        pytest.param(pa.array([], type=pa.int32()), id="empty_array"),
        pytest.param(
            pa.Tensor.from_numpy(np.array([[1, 2], [3, 4]], dtype=np.int32)),
            id="tensor",
        ),
    ]


@check_bytes_allocated
@pytest.mark.parametrize('obj', dlpack_objects())
@pytest.mark.parametrize('max_version', [None, (0, 8)])
def test_dlpack_legacy_capsule(obj, max_version):
    with pytest.warns(DeprecationWarning, match="unversioned DLPack capsule"):
        capsule = obj.__dlpack__(max_version=max_version)
    assert PyCapsule_IsValid(capsule, b"dltensor") is True


def immutable_tensor():
    np_arr = np.array([[1, 2], [3, 4]], dtype=np.int32)
    np_arr.flags.writeable = False
    tensor = pa.Tensor.from_numpy(np_arr)
    assert not tensor.is_mutable
    return tensor


@check_bytes_allocated
@pytest.mark.parametrize('max_version', [None, (0, 8)])
def test_dlpack_legacy_capsule_immutable_tensor(max_version):
    tensor = immutable_tensor()
    with pytest.raises(NotImplementedError,
                       match="Legacy DLPack support is not implemented "
                             "for immutable tensors"):
        tensor.__dlpack__(max_version=max_version)


@check_bytes_allocated
@pytest.mark.parametrize('max_version', [(1, 0), (1, 3), (2, 0)])
@pytest.mark.parametrize('copy', [None, False, True])
def test_dlpack_versioned_capsule_immutable_tensor(max_version, copy):
    tensor = immutable_tensor()
    capsule = tensor.__dlpack__(max_version=max_version, copy=copy)
    assert PyCapsule_IsValid(capsule, b"dltensor_versioned") is True


@check_bytes_allocated
@pytest.mark.parametrize('obj', dlpack_objects())
@pytest.mark.parametrize('max_version', [None, (0, 8)])
@pytest.mark.parametrize('copy', [False, True])
def test_dlpack_legacy_capsule_copy_not_supported(obj, max_version, copy):
    with pytest.raises(BufferError, match="copy argument is not supported"):
        obj.__dlpack__(max_version=max_version, copy=copy)


@check_bytes_allocated
@pytest.mark.parametrize('obj', dlpack_objects())
@pytest.mark.parametrize('max_version', [(1, 0), (1, 3), (2, 0)])
@pytest.mark.parametrize('copy', [None, False, True])
def test_dlpack_versioned_capsule(obj, max_version, copy):
    capsule = obj.__dlpack__(max_version=max_version, copy=copy)
    assert PyCapsule_IsValid(capsule, b"dltensor_versioned") is True


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('obj', dlpack_objects())
def test_dlpack_versioned_roundtrip(obj):
    expected = np.from_dlpack(DLPackForwarder(obj, max_version=None))
    for copy in [None, False, True]:
        result = np.from_dlpack(
            DLPackForwarder(obj, max_version=(1, 0), copy=copy))
        np.testing.assert_array_equal(result, expected, strict=True)


@requires_numpy_version("2.2.5")
@check_bytes_allocated
def test_dlpack_copy_is_writeable():
    # NumPy did not set the writeable flag on DLPack imports before 2.2.5.
    arr = pa.array([1, 2, 3], type=pa.int32())

    # Arrow arrays are immutable, so a shared export is read-only
    shared = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 3)))
    assert not shared.flags.writeable

    # A copy is solely owned by the consumer, who may mutate it
    copied = np.from_dlpack(DLPackForwarder(arr, max_version=(1, 3), copy=True))
    assert copied.flags.writeable
    copied[0] = 100
    assert arr.to_pylist() == [1, 2, 3]


def test_dlpack_not_supported():
    arr = pa.array([1, None, 3])
    with pytest.raises(TypeError, match="Can only use DLPack "
                       "on arrays with no nulls."):
        np.from_dlpack(arr)

    arr = pa.array(
        [[0, 1], [3, 4]],
        type=pa.list_(pa.int32())
    )
    with pytest.raises(TypeError, match="DataType is not compatible with DLPack spec"):
        np.from_dlpack(arr)

    arr = pa.array([])
    with pytest.raises(TypeError, match="DataType is not compatible with DLPack spec"):
        np.from_dlpack(arr)

    # DLPack doesn't support bit-packed boolean values
    arr = pa.array([True, False, True])
    with pytest.raises(TypeError, match="Bit-packed boolean data type "
                       "not supported by DLPack."):
        np.from_dlpack(arr)


def test_dlpack_cuda_not_supported():
    cuda = pytest.importorskip("pyarrow.cuda", exc_type=ImportError)

    schema = pa.schema([pa.field('f0', pa.int16())])
    a0 = pa.array([1, 2, 3], type=pa.int16())
    batch = pa.record_batch([a0], schema=schema)

    cbuf = cuda.serialize_record_batch(batch, cuda.Context(0))
    cbatch = cuda.read_record_batch(cbuf, batch.schema)
    carr = cbatch["f0"]

    # CudaBuffers not yet supported
    with pytest.raises(NotImplementedError, match="DLPack support is implemented "
                       "only for buffers on CPU device."):
        np.from_dlpack(carr)

    with pytest.raises(NotImplementedError, match="DLPack support is implemented "
                       "only for buffers on CPU device."):
        carr.__dlpack_device__()


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('np_type',
                         [np.uint8, np.uint16, np.uint32, np.uint64,
                          np.int8, np.int16, np.int32, np.int64,
                          np.float16, np.float32, np.float64])
def test_tensor_from_dlpack(np_type):
    def make_array():
        base = np.arange(24, dtype=np_type).reshape((4, 6))
        array = base[::2, 1::2]
        assert not array.flags['C_CONTIGUOUS']
        return array

    # Non-contiguous, strided slice: DLPack carries explicit strides, so this
    # should not need a copy on export.
    tensor = pa.Tensor.from_dlpack(make_array())
    assert isinstance(tensor, pa.Tensor)
    gc.collect()  # Attempts to free input array memory
    np.testing.assert_array_equal(tensor.to_numpy(), make_array(), strict=True)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('np_type',
                         [np.uint8, np.uint16, np.uint32, np.uint64,
                          np.int8, np.int16, np.int32, np.int64,
                          np.float16, np.float32, np.float64])
def test_array_from_dlpack(np_type):
    expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
    arr = pa.Array.from_dlpack(expected)
    arr.validate(full=True)
    assert isinstance(arr, pa.Array)
    np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('np_type',
                         [np.uint8, np.uint16, np.uint32, np.uint64,
                          np.int8, np.int16, np.int32, np.int64,
                          np.float16, np.float32, np.float64])
def test_fixed_shape_tensor_array_from_dlpack(np_type):
    source = np.arange(12, dtype=np_type).reshape((3, 2, 2))
    arr = pa.FixedShapeTensorArray.from_dlpack(source)
    arr.validate(full=True)
    assert arr.type == pa.fixed_shape_tensor(pa.from_numpy_dtype(np_type), [2, 2])
    assert arr.to_pylist() == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]

    # Zero-copy import: mutating the source is visible through the array.
    source[0, 0, 0] = 100
    assert arr.to_pylist()[0] == [100, 1, 2, 3]

    copied = pa.FixedShapeTensorArray.from_dlpack(source, copy=True)
    source[0, 0, 0] = 0
    assert copied.to_pylist()[0] == [100, 1, 2, 3]


@requires_numpy_version("2.1.0")
@check_bytes_allocated
@pytest.mark.parametrize('np_type', [np.uint8, np.int32, np.float64])
def test_fixed_shape_tensor_array_from_dlpack_transposed(np_type):
    source = np.arange(12, dtype=np_type).reshape((3, 2, 2)).transpose(0, 2, 1)
    arr = pa.FixedShapeTensorArray.from_dlpack(source)
    arr.validate(full=True)
    assert arr.type == pa.fixed_shape_tensor(
        pa.from_numpy_dtype(np_type), [2, 2], permutation=[1, 0])
    np.testing.assert_array_equal(arr.to_numpy_ndarray(), source)

    # Zero-copy import: mutating the source is visible through the array.
    source[0, 0, 0] = 100
    np.testing.assert_array_equal(arr.to_numpy_ndarray(), source)

    copied = pa.FixedShapeTensorArray.from_dlpack(source, copy=True)
    expected = source.copy()
    source[0, 0, 0] = 0
    np.testing.assert_array_equal(copied.to_numpy_ndarray(), expected)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_fixed_shape_tensor_array_from_dlpack_not_first_major():
    # The outermost dimension indexes the tensor elements, so it must remain
    # the major one.
    source = np.arange(12, dtype=np.int32).reshape((3, 2, 2)).transpose(1, 0, 2)
    with pytest.raises(pa.ArrowInvalid,
                       match="Only first-major tensors can be zero-copy"):
        pa.FixedShapeTensorArray.from_dlpack(source)


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_from_dlpack_zero_copy():
    expected = np.array([1, 2, 3], dtype=np.int64)
    tensor = pa.Tensor.from_dlpack(expected)
    result = tensor.to_numpy()
    # Zero-copy import: mutating the source is visible through the tensor.
    expected[0] = 100
    assert result[0] == 100
    # Same for mutating the result
    result[1] = 42
    assert expected[1] == 42


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_from_dlpack_explicit_copy():
    expected = np.array([1, 2, 3], dtype=np.int64)
    tensor = pa.Tensor.from_dlpack(expected, copy=True)
    result = tensor.to_numpy()
    expected[0] = 100
    # The data was copied, so mutating the source is not visible.
    assert result[0] == 1


def test_from_dlpack_no_dlpack_method():
    with pytest.raises(AttributeError):
        pa.Tensor.from_dlpack(object())


@requires_numpy_version("2.1.0")
@check_bytes_allocated
def test_array_from_dlpack_multi_dim_not_supported():
    expected = np.arange(6, dtype=np.int32).reshape((2, 3))
    with pytest.raises(
        pa.ArrowInvalid,
        match="Only contiguous one dimensional tensor can be imported as arrays",
    ):
        pa.Array.from_dlpack(expected)