import math

import pytest
import numpy as np
import array
from numpy.testing import assert_equal

from array_api_compat import (  # noqa: F401
    is_numpy_array, is_cupy_array, is_torch_array,
    is_dask_array, is_jax_array, is_pydata_sparse_array,
    is_ndonnx_array,
    is_numpy_namespace, is_cupy_namespace, is_torch_namespace,
    is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
    is_array_api_strict_namespace, is_ndonnx_namespace,
)

from array_api_compat import (
    device, is_array_api_obj, is_lazy_array, is_writeable_array, size, to_device
)
from ._helpers import all_libraries, import_, wrapped_libraries, xfail


is_array_functions = {
    'numpy': 'is_numpy_array',
#    'cupy': 'is_cupy_array',
    'torch': 'is_torch_array',
    'dask.array': 'is_dask_array',
#    'jax.numpy': 'is_jax_array',
#    'sparse': 'is_pydata_sparse_array',
#    'ndonnx': 'is_ndonnx_array',
}

is_namespace_functions = {
    'numpy': 'is_numpy_namespace',
#    'cupy': 'is_cupy_namespace',
    'torch': 'is_torch_namespace',
    'dask.array': 'is_dask_namespace',
#    'jax.numpy': 'is_jax_namespace',
#    'sparse': 'is_pydata_sparse_namespace',
#    'array_api_strict': 'is_array_api_strict_namespace',
#    'ndonnx': 'is_ndonnx_namespace',
}


@pytest.mark.parametrize('library', is_array_functions.keys())
@pytest.mark.parametrize('func', is_array_functions.values())
def test_is_xp_array(library, func):
    lib = import_(library)
    is_func = globals()[func]

    x = lib.asarray([1, 2, 3])

    assert is_func(x) == (func == is_array_functions[library])

    assert is_array_api_obj(x)


@pytest.mark.parametrize('library', is_namespace_functions.keys())
@pytest.mark.parametrize('func', is_namespace_functions.values())
def test_is_xp_namespace(library, func):
    lib = import_(library)
    is_func = globals()[func]

    assert is_func(lib) == (func == is_namespace_functions[library])


@pytest.mark.parametrize('library', all_libraries)
def test_xp_is_array_generics(library):
    """
    Test that scalar selection on a xp.ndarray always returns
    an object that matches with exactly one among the is_*_array
    function of the same library and is_numpy_array.
    """
    lib = import_(library)
    x = lib.asarray([1, 2, 3])
    x0 = x[0]

    matches = []
    for library2, func in is_array_functions.items():
        is_func = globals()[func]
        if is_func(x0):
            matches.append(library2)

    if library == "array_api_strict":
        # There is no is_array_api_strict_array() function
        assert matches == []
    else:
        assert matches in ([library], ["numpy"])


@pytest.mark.parametrize("library", all_libraries)
def test_is_writeable_array(library):
    lib = import_(library)
    x = lib.asarray([1, 2, 3])
    if is_writeable_array(x):
        x[1] = 4
    else:
        with pytest.raises((TypeError, ValueError)):
            x[1] = 4


def test_is_writeable_array_numpy():
    x = np.asarray([1, 2, 3])
    assert is_writeable_array(x)
    x.flags.writeable = False
    assert not is_writeable_array(x)


@pytest.mark.parametrize("library", all_libraries)
def test_size(library):
    xp = import_(library)
    x = xp.asarray([1, 2, 3])
    assert size(x) == 3


@pytest.mark.parametrize("library", all_libraries)
def test_size_none(library):
    if library == "sparse":
        pytest.skip("No arange(); no indexing by sparse arrays")

    xp = import_(library)
    x = xp.arange(10)
    x = x[x < 5]

    # dask.array now has shape=(nan, ) and size=nan
    # ndonnx now has shape=(None, ) and size=None
    # Eager libraries have shape=(5, ) and size=5
    assert size(x) in (None, 5)


@pytest.mark.parametrize("library", all_libraries)
def test_is_lazy_array(library):
    lib = import_(library)
    x = lib.asarray([1, 2, 3])
    assert isinstance(is_lazy_array(x), bool)


@pytest.mark.parametrize("shape", [(math.nan,), (1, math.nan), (None, ), (1, None)])
def test_is_lazy_array_nan_size(shape, monkeypatch):
    """Test is_lazy_array() on an unknown Array API compliant object
    with NaN (like Dask) or None (like ndonnx) in its shape
    """
    xp = import_("array_api_strict")
    x = xp.asarray(1)
    assert not is_lazy_array(x)
    monkeypatch.setattr(type(x), "shape", shape)
    assert is_lazy_array(x)


@pytest.mark.parametrize("exc", [TypeError, AssertionError])
def test_is_lazy_array_bool_raises(exc, monkeypatch):
    """Test is_lazy_array() on an unknown Array API compliant object
    where calling bool() raises:
    - TypeError: e.g. like jitted JAX. This is the proper exception which
      lazy arrays should raise as per the Array API specification
    - something else: e.g. like Dask, where bool() triggers compute()
      which can result in any kind of exception to be raised
    """
    xp = import_("array_api_strict")
    x = xp.asarray(1)
    assert not is_lazy_array(x)

    def __bool__(self):
        raise exc("Hello world")

    monkeypatch.setattr(type(x), "__bool__", __bool__)
    assert is_lazy_array(x)


@pytest.mark.parametrize(
    'func',
    list(is_array_functions.values()) 
    + ["is_array_api_obj", "is_lazy_array", "is_writeable_array"]
)
def test_is_array_any_object(func):
    """Test that is_*_array functions return False and don't raise on non-array objects
    """
    func = globals()[func]

    # These objects are missing attributes such as __name__
    assert not func(object())
    assert not func(None)
    assert not func(1)

    class C:
        pass

    assert not func(C())


@pytest.mark.parametrize("library", all_libraries)
def test_device(library, request):
    if library == "ndonnx":
        xfail(request, reason="Needs ndonnx >=0.9.4")

    xp = import_(library, wrapper=True)

    # We can't test much for device() and to_device() other than that
    # x.to_device(x.device) works.

    x = xp.asarray([1, 2, 3])
    dev = device(x)

    x2 = to_device(x, dev)
    assert device(x2) == device(x)

    x3 = xp.asarray(x, device=dev)
    assert device(x3) == device(x)


@pytest.mark.parametrize("library", wrapped_libraries)
def test_to_device_host(library):
    # different libraries have different semantics
    # for DtoH transfers; ensure that we support a portable
    # shim for common array libs
    # see: https://github.com/scipy/scipy/issues/18286#issuecomment-1527552919
    xp = import_(library, wrapper=True)

    expected = np.array([1, 2, 3])
    x = xp.asarray([1, 2, 3])
    x = to_device(x, "cpu")
    # torch will return a genuine Device object, but
    # the other libs will do something different with
    # a `device(x)` query; however, what's really important
    # here is that we can test portably after calling
    # to_device(x, "cpu") to return to host
    assert_equal(x, expected)


@pytest.mark.parametrize("target_library", is_array_functions.keys())
@pytest.mark.parametrize("source_library", is_array_functions.keys())
def test_asarray_cross_library(source_library, target_library, request):
    if source_library == "dask.array" and target_library == "torch":
        # TODO: remove xfail once
        # https://github.com/dask/dask/issues/8260 is resolved
        xfail(request, reason="Bug in dask raising error on conversion")
    elif (
        source_library == "ndonnx" 
        and target_library not in ("array_api_strict", "ndonnx", "numpy")
    ):
        xfail(request, reason="The truth value of lazy Array Array(dtype=Boolean) is unknown")
    elif source_library == "ndonnx" and target_library == "numpy":
        xfail(request, reason="produces numpy array of ndonnx scalar arrays")
    elif source_library == "jax.numpy" and target_library == "torch":
        xfail(request, reason="casts int to float")
    elif source_library == "cupy" and target_library != "cupy":
        # cupy explicitly disallows implicit conversions to CPU
        pytest.skip(reason="cupy does not support implicit conversion to CPU")
    elif source_library == "sparse" and target_library != "sparse":
        pytest.skip(reason="`sparse` does not allow implicit densification")

    src_lib = import_(source_library, wrapper=True)
    tgt_lib = import_(target_library, wrapper=True)
    is_tgt_type = globals()[is_array_functions[target_library]]

    a = src_lib.asarray([1, 2, 3], dtype=src_lib.int32)
    b = tgt_lib.asarray(a)

    assert is_tgt_type(b), f"Expected {b} to be a {tgt_lib.ndarray}, but was {type(b)}"
    assert b.dtype == tgt_lib.int32



@pytest.mark.parametrize("library", wrapped_libraries)
def test_asarray_copy(library):
    # Note, we have this test here because the test suite currently doesn't
    # test the copy flag to asarray() very rigorously. Once
    # https://github.com/data-apis/array-api-tests/issues/241 is fixed we
    # should be able to delete this.
    xp = import_(library, wrapper=True)
    asarray = xp.asarray
    is_lib_func = globals()[is_array_functions[library]]
    all = xp.all if library != 'dask.array' else lambda x: xp.all(x).compute()

    if library == 'numpy' and xp.__version__[0] < '2' and not hasattr(xp, '_CopyMode') :
        supports_copy_false_other_ns = False
        supports_copy_false_same_ns = False
    elif library == 'cupy':
        supports_copy_false_other_ns = False
        supports_copy_false_same_ns = False
    elif library == 'dask.array':
        supports_copy_false_other_ns = False
        supports_copy_false_same_ns = True
    else:
        supports_copy_false_other_ns = True
        supports_copy_false_same_ns = True

    a = asarray([1])
    b = asarray(a, copy=True)
    assert is_lib_func(b)
    a[0] = 0
    assert all(b[0] == 1)
    assert all(a[0] == 0)

    a = asarray([1])
    if supports_copy_false_same_ns:
        b = asarray(a, copy=False)
        assert is_lib_func(b)
        a[0] = 0
        assert all(b[0] == 0)
    else:
        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False))

    a = asarray([1])
    if supports_copy_false_same_ns:
        pytest.raises(ValueError, lambda: asarray(a, copy=False,
                                                  dtype=xp.float64))
    else:
        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False, dtype=xp.float64))

    a = asarray([1])
    b = asarray(a, copy=None)
    assert is_lib_func(b)
    a[0] = 0
    assert all(b[0] == 0)

    a = asarray([1.0], dtype=xp.float32)
    assert a.dtype == xp.float32
    b = asarray(a, dtype=xp.float64, copy=None)
    assert is_lib_func(b)
    assert b.dtype == xp.float64
    a[0] = 0.0
    assert all(b[0] == 1.0)

    a = asarray([1.0], dtype=xp.float64)
    assert a.dtype == xp.float64
    b = asarray(a, dtype=xp.float64, copy=None)
    assert is_lib_func(b)
    assert b.dtype == xp.float64
    a[0] = 0.0
    assert all(b[0] == 0.0)

    # Python built-in types
    for obj in [True, 0, 0.0, 0j, [0], [[0]]]:
        asarray(obj, copy=True) # No error
        asarray(obj, copy=None) # No error
        if supports_copy_false_other_ns:
            pytest.raises(ValueError, lambda: asarray(obj, copy=False))
        else:
            pytest.raises(NotImplementedError, lambda: asarray(obj, copy=False))

    # Use the standard library array to test the buffer protocol
    a = array.array('f', [1.0])
    b = asarray(a, copy=True)
    assert is_lib_func(b)
    a[0] = 0.0
    assert all(b[0] == 1.0)

    a = array.array('f', [1.0])
    if supports_copy_false_other_ns:
        b = asarray(a, copy=False)
        assert is_lib_func(b)
        a[0] = 0.0
        assert all(b[0] == 0.0)
    else:
        pytest.raises(NotImplementedError, lambda: asarray(a, copy=False))

    a = array.array('f', [1.0])
    b = asarray(a, copy=None)
    assert is_lib_func(b)
    a[0] = 0.0
    if library in ('cupy', 'dask.array'):
        # A copy is required for libraries where the default device is not CPU
        # dask changed behaviour of copy=None in 2024.12 to copy;
        # this wrapper ensures the same behaviour in older versions too.
        # https://github.com/dask/dask/pull/11524/
        assert all(b[0] == 1.0)
    else:
        assert all(b[0] == 0.0)


@pytest.mark.parametrize("library", ["numpy", "cupy", "torch"])
def test_clip_out(library):
    """Test non-standard out= parameter for clip()

    (see "Avoid Restricting Behavior that is Outside the Scope of the Standard"
    in https://data-apis.org/array-api-compat/dev/special-considerations.html)
    """
    xp = import_(library, wrapper=True)
    x = xp.asarray([10, 20, 30])
    out = xp.zeros_like(x)
    xp.clip(x, 15, 25, out=out)
    expect = xp.asarray([15, 20, 25])
    assert xp.all(out == expect)