# 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.

from __future__ import annotations

from typing import Any, NamedTuple, NoReturn

import numpy.typing as npt
import pytest

try:
    import torch
    import torch.version
except ImportError:
    torch = None  # ty: ignore[invalid-assignment]

import numpy as np
import tvm_ffi


def test_tensor_attributes() -> None:
    data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int16")
    if not hasattr(data, "__dlpack__"):
        return
    x = tvm_ffi.from_dlpack(data)
    assert isinstance(x, tvm_ffi.Tensor)
    assert x.shape == (10, 8, 4, 2)
    assert x.ndim == 4
    assert x.numel() == 640
    assert x.size(0) == 10
    assert x.size(-1) == 2
    assert x.is_contiguous()
    assert x.strides == (64, 8, 2, 1)
    assert x.dtype == tvm_ffi.dtype("int16")
    assert x.device.dlpack_device_type() == tvm_ffi.DLDeviceType.kDLCPU
    assert x.device.index == 0
    x2 = np.from_dlpack(x)
    np.testing.assert_equal(x2, data)


def test_empty_tensor_is_contiguous() -> None:
    # Empty tensors are trivially contiguous regardless of what
    # strides the producer reports (numpy 2.3+ via __dlpack__ now
    # reports (0, 0, 0) for shape (4, 0, 4)). See PR #607 review
    # comment for context.
    data: npt.NDArray[Any] = np.zeros((4, 0, 4), dtype="int16")
    if not hasattr(data, "__dlpack__"):
        return
    x = tvm_ffi.from_dlpack(data)
    assert x.is_contiguous()


def test_non_contiguous_tensor_attributes() -> None:
    data: npt.NDArray[Any] = np.zeros((4, 4, 4), dtype="int16")
    slice = data[1:3, :, 1:3]
    if not hasattr(slice, "__dlpack__"):
        return
    x = tvm_ffi.from_dlpack(slice)
    assert isinstance(x, tvm_ffi.Tensor)
    assert x.shape == (2, 4, 2)
    assert x.numel() == 16
    assert x.size(0) == 2
    assert x.size(-1) == 2
    assert not x.is_contiguous()
    assert x.strides == (16, 4, 1)


def test_shape_object() -> None:
    shape = tvm_ffi.Shape((10, 8, 4, 2))
    assert isinstance(shape, tvm_ffi.Shape)
    assert shape == (10, 8, 4, 2)

    fecho = tvm_ffi.convert(lambda x: x)
    shape2: tvm_ffi.Shape = fecho(shape)
    assert shape2._tvm_ffi_cached_object.same_as(shape._tvm_ffi_cached_object)
    assert isinstance(shape2, tvm_ffi.Shape)
    assert isinstance(shape2, tuple)

    shape3: tvm_ffi.Shape = tvm_ffi.convert(shape)
    assert shape3._tvm_ffi_cached_object.same_as(shape._tvm_ffi_cached_object)
    assert isinstance(shape3, tvm_ffi.Shape)


@pytest.mark.skipif(torch is None, reason="Fast torch dlpack importer is not enabled")
def test_tensor_auto_dlpack() -> None:
    assert torch is not None
    x = torch.arange(128)
    fecho = tvm_ffi.get_global_func("testing.echo")
    y = fecho(x)
    assert isinstance(y, torch.Tensor)
    assert y.data_ptr() == x.data_ptr()
    assert y.dtype == x.dtype
    assert y.shape == x.shape
    assert y.device == x.device
    np.testing.assert_equal(y.numpy(), x.numpy())


@pytest.mark.skipif(torch is None, reason="Fast torch dlpack importer is not enabled")
def test_tensor_auto_dlpack_with_error() -> None:
    assert torch is not None
    x = torch.arange(128)

    def raise_torch_error(x: Any) -> NoReturn:
        raise ValueError("error XYZ")

    f = tvm_ffi.convert(raise_torch_error)
    with pytest.raises(ValueError):
        # pass in torch argment to trigger the error in set allocator path
        f(x)


def test_tensor_class_override() -> None:
    class MyTensor(tvm_ffi.Tensor):
        pass

    old_tensor = tvm_ffi.core._CLASS_TENSOR
    tvm_ffi.core._set_class_tensor(MyTensor)

    data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int16")
    if not hasattr(data, "__dlpack__"):
        return
    x = tvm_ffi.from_dlpack(data)

    fecho = tvm_ffi.get_global_func("testing.echo")
    y = fecho(x)
    assert isinstance(y, MyTensor)
    tvm_ffi.core._set_class_tensor(old_tensor)


def test_tvm_ffi_tensor_compatible() -> None:
    class MyTensor:
        def __init__(self, tensor: tvm_ffi.Tensor) -> None:
            """Initialize the MyTensor."""
            self._tensor = tensor

        def __tvm_ffi_object__(self) -> tvm_ffi.Tensor:
            """Implement __tvm_ffi_object__ protocol."""
            return self._tensor

    data: npt.NDArray[Any] = np.zeros((10, 8, 4, 2), dtype="int32")
    if not hasattr(data, "__dlpack__"):
        return
    x = tvm_ffi.from_dlpack(data)
    y = MyTensor(x)
    fecho = tvm_ffi.get_global_func("testing.echo")
    z = fecho(y)
    assert z.__chandle__() == x.__chandle__()

    class MyNamedTuple(NamedTuple):
        a: MyTensor
        b: int

    args = MyNamedTuple(a=y, b=1)
    z = fecho(args)
    assert z[0].__chandle__() == x.__chandle__()
    assert z[1] == 1

    class MyCustom:
        def __init__(self, a: MyTensor, b: int) -> None:
            self.a = a
            self.b = b

        def __tvm_ffi_value__(self) -> Any:
            """Implement __tvm_ffi_value__ protocol."""
            return (self.a, self.b)

    z = fecho(MyCustom(a=y, b=2))
    assert z[0].__chandle__() == x.__chandle__()
    assert z[1] == 2


@pytest.mark.skipif(
    torch is None or not torch.cuda.is_available() or torch.version.hip is None,
    reason="ROCm is not enabled in PyTorch",
)
def test_tensor_from_pytorch_rocm() -> None:
    assert torch is not None

    @tvm_ffi.register_global_func("testing.check_device", override=True)
    def _check_device(x: tvm_ffi.Tensor) -> str:
        return x.device.type

    # PyTorch uses device name "cuda" to represent ROCm device
    x = torch.randn(128, device="cuda")
    device_type = tvm_ffi.get_global_func("testing.check_device")(x)
    assert device_type == "rocm"


def test_optional_tensor_view() -> None:
    optional_tensor_view_has_value = tvm_ffi.get_global_func(
        "testing.optional_tensor_view_has_value"
    )
    assert not optional_tensor_view_has_value(None)
    x: npt.NDArray[Any] = np.zeros((128,), dtype="float32")
    if not hasattr(x, "__dlpack__"):
        return
    assert optional_tensor_view_has_value(x)


if __name__ == "__main__":
    pytest.main([__file__])