import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm
from pandas.core.arrays.sparse import SparseArray


@pytest.mark.parametrize(
    "kwargs",
    [
        {},  # Default is check_exact=False
        {"check_exact": False},
        {"check_exact": True},
    ],
)
def test_assert_extension_array_equal_not_exact(kwargs):
    # see gh-23709
    arr1 = SparseArray([-0.17387645482451206, 0.3414148016424936])
    arr2 = SparseArray([-0.17387645482451206, 0.3414148016424937])

    if kwargs.get("check_exact", False):
        msg = """\
ExtensionArray are different

ExtensionArray values are different \\(50\\.0 %\\)
\\[left\\]:  \\[-0\\.173876.*, 0\\.341415.*\\]
\\[right\\]: \\[-0\\.173876.*, 0\\.341415.*\\]"""

        with pytest.raises(AssertionError, match=msg):
            tm.assert_extension_array_equal(arr1, arr2, **kwargs)
    else:
        tm.assert_extension_array_equal(arr1, arr2, **kwargs)


@pytest.mark.parametrize("decimals", range(10))
def test_assert_extension_array_equal_less_precise(decimals):
    rtol = 0.5 * 10**-decimals
    arr1 = SparseArray([0.5, 0.123456])
    arr2 = SparseArray([0.5, 0.123457])

    if decimals >= 5:
        msg = """\
ExtensionArray are different

ExtensionArray values are different \\(50\\.0 %\\)
\\[left\\]:  \\[0\\.5, 0\\.123456\\]
\\[right\\]: \\[0\\.5, 0\\.123457\\]"""

        with pytest.raises(AssertionError, match=msg):
            tm.assert_extension_array_equal(arr1, arr2, rtol=rtol)
    else:
        tm.assert_extension_array_equal(arr1, arr2, rtol=rtol)


def test_assert_extension_array_equal_dtype_mismatch(check_dtype):
    end = 5
    kwargs = {"check_dtype": check_dtype}

    arr1 = SparseArray(np.arange(end, dtype="int64"))
    arr2 = SparseArray(np.arange(end, dtype="int32"))

    if check_dtype:
        msg = """\
ExtensionArray are different

Attribute "dtype" are different
\\[left\\]:  Sparse\\[int64, 0\\]
\\[right\\]: Sparse\\[int32, 0\\]"""

        with pytest.raises(AssertionError, match=msg):
            tm.assert_extension_array_equal(arr1, arr2, **kwargs)
    else:
        tm.assert_extension_array_equal(arr1, arr2, **kwargs)


def test_assert_extension_array_equal_missing_values():
    arr1 = SparseArray([np.nan, 1, 2, np.nan])
    arr2 = SparseArray([np.nan, 1, 2, 3])

    msg = """\
ExtensionArray NA mask are different

ExtensionArray NA mask values are different \\(25\\.0 %\\)
\\[left\\]:  \\[True, False, False, True\\]
\\[right\\]: \\[True, False, False, False\\]"""

    with pytest.raises(AssertionError, match=msg):
        tm.assert_extension_array_equal(arr1, arr2)


@pytest.mark.parametrize("side", ["left", "right"])
def test_assert_extension_array_equal_non_extension_array(side):
    numpy_array = np.arange(5)
    extension_array = SparseArray(numpy_array)

    msg = f"{side} is not an ExtensionArray"
    args = (
        (numpy_array, extension_array)
        if side == "left"
        else (extension_array, numpy_array)
    )

    with pytest.raises(AssertionError, match=msg):
        tm.assert_extension_array_equal(*args)


def test_assert_extension_array_equal_ignore_dtype_mismatch(any_int_dtype):
    # https://github.com/pandas-dev/pandas/issues/35715
    left = pd.array([1, 2, 3], dtype="Int64")
    right = pd.array([1, 2, 3], dtype=any_int_dtype)
    tm.assert_extension_array_equal(left, right, check_dtype=False)


def test_assert_extension_array_equal_time_units():
    # https://github.com/pandas-dev/pandas/issues/55730
    timestamp = pd.Timestamp("2023-11-04T12")
    naive = pd.array([timestamp], dtype="datetime64[ns]")
    utc = pd.array([timestamp], dtype="datetime64[ns, UTC]")

    tm.assert_extension_array_equal(naive, utc, check_dtype=False)
    tm.assert_extension_array_equal(utc, naive, check_dtype=False)


@pytest.mark.parametrize("check_exact", [True, False])
def test_assert_extension_array_equal_zero_dim_duck_array(check_exact):
    # GH#45240 unequal arrays of scalars that define __iter__/__len__ delegating
    #  to a scalar payload (as pint's Quantity does) report the difference
    #  instead of raising TypeError
    class Quantity:
        # mimics pint's Quantity: __iter__/__len__ delegate to the magnitude,
        #  and ndim reports 0 when that magnitude is a scalar
        ndim = 0

        def __init__(self, magnitude):
            self.magnitude = magnitude

        def __iter__(self):
            return iter(self.magnitude)

        def __len__(self):
            return len(self.magnitude)

        def __eq__(self, other):
            return self.magnitude == other.magnitude

    left = pd.array(np.array([Quantity(1), Quantity(2)], dtype=object))
    right = pd.array(np.array([Quantity(1), Quantity(3)], dtype=object))

    tm.assert_extension_array_equal(left, left.copy(), check_exact=check_exact)

    msg = "ExtensionArray are different"
    with pytest.raises(AssertionError, match=msg):
        tm.assert_extension_array_equal(left, right, check_exact=check_exact)


def test_assert_extension_array_equal_large_mixed_integer_float():
    # GH#66699 tolerances must be honored at full integer precision
    left = pd.array([2**60 + 1], dtype="Int64")
    right = pd.array([float(2**60)], dtype="Float64")

    for first, second in [(left, right), (right, left)]:
        with pytest.raises(AssertionError, match="ExtensionArray are different"):
            tm.assert_extension_array_equal(
                first, second, check_dtype=False, check_exact=False, rtol=0, atol=0.5
            )
        tm.assert_extension_array_equal(
            first, second, check_dtype=False, check_exact=False, rtol=0, atol=1
        )