"""Testing functions exposed to the user API"""

import functools
import warnings
from collections.abc import Hashable

import numpy as np
import pandas as pd

from xarray.core import duck_array_ops, formatting, utils
from xarray.core.coordinates import Coordinates
from xarray.core.dataarray import DataArray
from xarray.core.dataset import Dataset
from xarray.core.datatree import DataTree
from xarray.core.datatree_mapping import map_over_datasets
from xarray.core.formatting import diff_datatree_repr
from xarray.core.indexes import Index, PandasIndex, PandasMultiIndex, default_indexes
from xarray.core.variable import IndexVariable, Variable


def ensure_warnings(func):
    # sometimes tests elevate warnings to errors
    # -> make sure that does not happen in the assert_* functions
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        __tracebackhide__ = True

        with warnings.catch_warnings():
            # only remove filters that would "error"
            warnings.filters = [f for f in warnings.filters if f[0] != "error"]

            return func(*args, **kwargs)

    return wrapper


def _decode_string_data(data):
    if data.dtype.kind == "S":
        return np.char.decode(data, "utf-8", "replace")
    return data


def _data_allclose_or_equiv(arr1, arr2, rtol=1e-05, atol=1e-08, decode_bytes=True):
    if any(arr.dtype.kind == "S" for arr in [arr1, arr2]) and decode_bytes:
        arr1 = _decode_string_data(arr1)
        arr2 = _decode_string_data(arr2)
    exact_dtypes = ["M", "m", "O", "S", "U"]
    if any(arr.dtype.kind in exact_dtypes for arr in [arr1, arr2]):
        return duck_array_ops.array_equiv(arr1, arr2)
    else:
        return duck_array_ops.allclose_or_equiv(arr1, arr2, rtol=rtol, atol=atol)


@ensure_warnings
def assert_isomorphic(a: DataTree, b: DataTree):
    """
    Two DataTrees are considered isomorphic if the set of paths to their
    descendent nodes are the same.

    Nothing about the data or attrs in each node is checked.

    Isomorphism is a necessary condition for two trees to be used in a nodewise binary operation,
    such as tree1 + tree2.

    Parameters
    ----------
    a : DataTree
        The first object to compare.
    b : DataTree
        The second object to compare.

    See Also
    --------
    DataTree.isomorphic
    assert_equal
    assert_identical
    """
    __tracebackhide__ = True
    assert isinstance(a, type(b))

    if isinstance(a, DataTree):
        assert a.isomorphic(b), diff_datatree_repr(a, b, "isomorphic")
    else:
        raise TypeError(f"{type(a)} not of type DataTree")


def maybe_transpose_dims(a, b, check_dim_order: bool):
    """Helper for assert_equal/allclose/identical"""

    __tracebackhide__ = True

    def _maybe_transpose_dims(a, b):
        if not isinstance(a, Variable | DataArray | Dataset):
            return b
        if set(a.dims) == set(b.dims):
            # Ensure transpose won't fail if a dimension is missing
            # If this is the case, the difference will be caught by the caller
            return b.transpose(*a.dims)
        return b

    if check_dim_order:
        return b

    if isinstance(a, DataTree):
        return map_over_datasets(_maybe_transpose_dims, a, b)

    return _maybe_transpose_dims(a, b)


@ensure_warnings
def assert_equal(a, b, check_dim_order: bool = True):
    """Like :py:func:`numpy.testing.assert_array_equal`, but for xarray
    objects.

    Raises an AssertionError if two objects are not equal. This will match
    data values, dimensions and coordinates, but not names or attributes
    (except for Dataset objects for which the variable names must match).
    Arrays with NaN in the same location are considered equal.

    For DataTree objects, assert_equal is mapped over all Datasets on each node,
    with the DataTrees being equal if both are isomorphic and the corresponding
    Datasets at each node are themselves equal.

    Parameters
    ----------
    a : xarray.Dataset, xarray.DataArray, xarray.Variable, xarray.Coordinates
        or xarray.core.datatree.DataTree. The first object to compare.
    b : xarray.Dataset, xarray.DataArray, xarray.Variable, xarray.Coordinates
        or xarray.core.datatree.DataTree. The second object to compare.
    check_dim_order : bool, optional, default is True
        Whether dimensions must be in the same order.

    See Also
    --------
    assert_identical, assert_allclose, Dataset.equals, DataArray.equals
    numpy.testing.assert_array_equal
    """
    __tracebackhide__ = True
    assert type(a) is type(b) or (
        isinstance(a, Coordinates) and isinstance(b, Coordinates)
    )
    b = maybe_transpose_dims(a, b, check_dim_order)
    if isinstance(a, Variable | DataArray):
        assert a.equals(b), formatting.diff_array_repr(a, b, "equals")
    elif isinstance(a, Dataset):
        assert a.equals(b), formatting.diff_dataset_repr(a, b, "equals")
    elif isinstance(a, Coordinates):
        assert a.equals(b), formatting.diff_coords_repr(a, b, "equals")
    elif isinstance(a, DataTree):
        assert a.equals(b), diff_datatree_repr(a, b, "equals")
    else:
        raise TypeError(f"{type(a)} not supported by assertion comparison")


@ensure_warnings
def assert_identical(a, b):
    """Like :py:func:`xarray.testing.assert_equal`, but also matches the
    objects' names and attributes.

    Raises an AssertionError if two objects are not identical.

    For DataTree objects, assert_identical is mapped over all Datasets on each
    node, with the DataTrees being identical if both are isomorphic and the
    corresponding Datasets at each node are themselves identical.

    Parameters
    ----------
    a : xarray.Dataset, xarray.DataArray, xarray.Variable or xarray.Coordinates
        The first object to compare.
    b : xarray.Dataset, xarray.DataArray, xarray.Variable or xarray.Coordinates
        The second object to compare.

    See Also
    --------
    assert_equal, assert_allclose, Dataset.equals, DataArray.equals
    """
    __tracebackhide__ = True
    assert type(a) is type(b) or (
        isinstance(a, Coordinates) and isinstance(b, Coordinates)
    )
    if isinstance(a, Variable):
        assert a.identical(b), formatting.diff_array_repr(a, b, "identical")
    elif isinstance(a, DataArray):
        assert a.name == b.name, (
            f"DataArray names are different. L: {a.name}, R: {b.name}"
        )
        assert a.identical(b), formatting.diff_array_repr(a, b, "identical")
    elif isinstance(a, Dataset | Variable):
        assert a.identical(b), formatting.diff_dataset_repr(a, b, "identical")
    elif isinstance(a, Coordinates):
        assert a.identical(b), formatting.diff_coords_repr(a, b, "identical")
    elif isinstance(a, DataTree):
        assert a.identical(b), diff_datatree_repr(a, b, "identical")
    else:
        raise TypeError(f"{type(a)} not supported by assertion comparison")


@ensure_warnings
def assert_allclose(
    a, b, rtol=1e-05, atol=1e-08, decode_bytes=True, check_dim_order: bool = True
):
    """Like :py:func:`numpy.testing.assert_allclose`, but for xarray objects.

    Raises an AssertionError if two objects are not equal up to desired
    tolerance.

    Parameters
    ----------
    a : xarray.Dataset, xarray.DataArray or xarray.Variable
        The first object to compare.
    b : xarray.Dataset, xarray.DataArray or xarray.Variable
        The second object to compare.
    rtol : float, optional
        Relative tolerance.
    atol : float, optional
        Absolute tolerance.
    decode_bytes : bool, optional
        Whether byte dtypes should be decoded to strings as UTF-8 or not.
        This is useful for testing serialization methods on Python 3 that
        return saved strings as bytes.
    check_dim_order : bool, optional, default is True
        Whether dimensions must be in the same order.

    See Also
    --------
    assert_identical, assert_equal, numpy.testing.assert_allclose
    """
    __tracebackhide__ = True
    assert type(a) is type(b)
    b = maybe_transpose_dims(a, b, check_dim_order)

    equiv = functools.partial(
        _data_allclose_or_equiv, rtol=rtol, atol=atol, decode_bytes=decode_bytes
    )
    equiv.__name__ = "allclose"  # type: ignore[attr-defined]

    def compat_variable(a, b):
        a = getattr(a, "variable", a)
        b = getattr(b, "variable", b)
        return a.dims == b.dims and (a._data is b._data or equiv(a.data, b.data))

    if isinstance(a, Variable):
        allclose = compat_variable(a, b)
        assert allclose, formatting.diff_array_repr(a, b, compat=equiv)
    elif isinstance(a, DataArray):
        allclose = utils.dict_equiv(
            a.coords, b.coords, compat=compat_variable
        ) and compat_variable(a.variable, b.variable)
        assert allclose, formatting.diff_array_repr(a, b, compat=equiv)
    elif isinstance(a, Dataset):
        allclose = a._coord_names == b._coord_names and utils.dict_equiv(
            a.variables, b.variables, compat=compat_variable
        )
        assert allclose, formatting.diff_dataset_repr(a, b, compat=equiv)
    elif isinstance(a, Coordinates):
        allclose = utils.dict_equiv(a.variables, b.variables, compat=compat_variable)
        assert allclose, formatting.diff_coords_repr(a, b, compat=equiv)
    else:
        raise TypeError(f"{type(a)} not supported by assertion comparison")


def _format_message(x, y, err_msg, verbose):
    diff = x - y
    abs_diff = max(abs(diff))
    rel_diff = "not implemented"

    n_diff = np.count_nonzero(diff)
    n_total = diff.size

    fraction = f"{n_diff} / {n_total}"
    percentage = float(n_diff / n_total * 100)

    parts = [
        "Arrays are not equal",
        err_msg,
        f"Mismatched elements: {fraction} ({percentage:.0f}%)",
        f"Max absolute difference: {abs_diff}",
        f"Max relative difference: {rel_diff}",
    ]
    if verbose:
        parts += [
            f" x: {x!r}",
            f" y: {y!r}",
        ]

    return "\n".join(parts)


@ensure_warnings
def assert_duckarray_allclose(
    actual, desired, rtol=1e-07, atol=0, err_msg="", verbose=True
):
    """Like `np.testing.assert_allclose`, but for duckarrays."""
    __tracebackhide__ = True

    allclose = duck_array_ops.allclose_or_equiv(actual, desired, rtol=rtol, atol=atol)
    assert allclose, _format_message(actual, desired, err_msg=err_msg, verbose=verbose)


@ensure_warnings
def assert_duckarray_equal(x, y, err_msg="", verbose=True):
    """Like `np.testing.assert_array_equal`, but for duckarrays"""
    __tracebackhide__ = True

    if not utils.is_duck_array(x) and not utils.is_scalar(x):
        x = np.asarray(x)

    if not utils.is_duck_array(y) and not utils.is_scalar(y):
        y = np.asarray(y)

    if (utils.is_duck_array(x) and utils.is_scalar(y)) or (
        utils.is_scalar(x) and utils.is_duck_array(y)
    ):
        equiv = duck_array_ops.array_all(x == y)
    else:
        equiv = duck_array_ops.array_equiv(x, y)
    assert equiv, _format_message(x, y, err_msg=err_msg, verbose=verbose)


def assert_chunks_equal(a, b):
    """
    Assert that chunksizes along chunked dimensions are equal.

    Parameters
    ----------
    a : xarray.Dataset or xarray.DataArray
        The first object to compare.
    b : xarray.Dataset or xarray.DataArray
        The second object to compare.
    """

    if isinstance(a, DataArray) != isinstance(b, DataArray):
        raise TypeError("a and b have mismatched types")

    left = a.unify_chunks()
    right = b.unify_chunks()
    assert left.chunks == right.chunks


def _assert_indexes_invariants_checks(
    indexes, possible_coord_variables, dims, check_default=True
):
    assert isinstance(indexes, dict), indexes
    assert all(isinstance(v, Index) for v in indexes.values()), {
        k: type(v) for k, v in indexes.items()
    }

    if check_default:
        index_vars = {
            k
            for k, v in possible_coord_variables.items()
            if isinstance(v, IndexVariable)
        }
        assert indexes.keys() <= index_vars, (set(indexes), index_vars)

    # check pandas index wrappers vs. coordinate data adapters
    for k, index in indexes.items():
        if isinstance(index, PandasIndex):
            pd_index = index.index
            var = possible_coord_variables[k]
            assert (index.dim,) == var.dims, (pd_index, var)
            if k == index.dim:
                # skip multi-index levels here (checked below)
                assert index.coord_dtype == var.dtype, (index.coord_dtype, var.dtype)
            assert isinstance(var._data.array, pd.Index), var._data.array
            # TODO: check identity instead of equality?
            assert pd_index.equals(var._data.array), (pd_index, var)
        if isinstance(index, PandasMultiIndex):
            pd_index = index.index
            for name in index.index.names:
                assert name in possible_coord_variables, (pd_index, index_vars)
                var = possible_coord_variables[name]
                assert (index.dim,) == var.dims, (pd_index, var)
                assert index.level_coords_dtype[name] == var.dtype, (
                    index.level_coords_dtype[name],
                    var.dtype,
                )
                assert isinstance(var._data.array, pd.MultiIndex), var._data.array
                assert pd_index.equals(var._data.array), (pd_index, var)
                # check all all levels are in `indexes`
                assert name in indexes, (name, set(indexes))
                # index identity is used to find unique indexes in `indexes`
                assert index is indexes[name], (pd_index, indexes[name].index)

    if check_default:
        defaults = default_indexes(possible_coord_variables, dims)
        assert indexes.keys() == defaults.keys(), (set(indexes), set(defaults))
        assert all(v.equals(defaults[k]) for k, v in indexes.items()), (
            indexes,
            defaults,
        )


def _assert_variable_invariants(var: Variable, name: Hashable = None):
    if name is None:
        name_or_empty: tuple = ()
    else:
        name_or_empty = (name,)
    assert isinstance(var._dims, tuple), name_or_empty + (var._dims,)
    assert len(var._dims) == len(var._data.shape), name_or_empty + (
        var._dims,
        var._data.shape,
    )
    assert isinstance(var._encoding, type(None) | dict), name_or_empty + (
        var._encoding,
    )
    assert isinstance(var._attrs, type(None) | dict), name_or_empty + (var._attrs,)


def _assert_dataarray_invariants(da: DataArray, check_default_indexes: bool):
    assert isinstance(da._variable, Variable), da._variable
    _assert_variable_invariants(da._variable)

    assert isinstance(da._coords, dict), da._coords
    assert all(isinstance(v, Variable) for v in da._coords.values()), da._coords

    if check_default_indexes:
        assert all(set(v.dims) <= set(da.dims) for v in da._coords.values()), (
            da.dims,
            {k: v.dims for k, v in da._coords.items()},
        )
        assert all(
            isinstance(v, IndexVariable)
            for (k, v) in da._coords.items()
            if v.dims == (k,)
        ), {k: type(v) for k, v in da._coords.items()}

    for k, v in da._coords.items():
        _assert_variable_invariants(v, k)

    if da._indexes is not None:
        _assert_indexes_invariants_checks(
            da._indexes, da._coords, da.dims, check_default=check_default_indexes
        )


def _assert_dataset_invariants(ds: Dataset, check_default_indexes: bool):
    assert isinstance(ds._variables, dict), type(ds._variables)
    assert all(isinstance(v, Variable) for v in ds._variables.values()), ds._variables
    for k, v in ds._variables.items():
        _assert_variable_invariants(v, k)

    assert isinstance(ds._coord_names, set), ds._coord_names
    assert ds._coord_names <= ds._variables.keys(), (
        ds._coord_names,
        set(ds._variables),
    )

    assert type(ds._dims) is dict, ds._dims
    assert all(isinstance(v, int) for v in ds._dims.values()), ds._dims
    var_dims: set[Hashable] = set()
    for v in ds._variables.values():
        var_dims.update(v.dims)
    assert ds._dims.keys() == var_dims, (set(ds._dims), var_dims)
    assert all(
        ds._dims[k] == v.sizes[k] for v in ds._variables.values() for k in v.sizes
    ), (ds._dims, {k: v.sizes for k, v in ds._variables.items()})

    if check_default_indexes:
        assert all(
            isinstance(v, IndexVariable)
            for (k, v) in ds._variables.items()
            if v.dims == (k,)
        ), {k: type(v) for k, v in ds._variables.items() if v.dims == (k,)}

    if ds._indexes is not None:
        _assert_indexes_invariants_checks(
            ds._indexes, ds._variables, ds._dims, check_default=check_default_indexes
        )

    assert isinstance(ds._encoding, type(None) | dict)
    assert isinstance(ds._attrs, type(None) | dict)


def _assert_internal_invariants(
    xarray_obj: DataArray | Dataset | Variable, check_default_indexes: bool
):
    """Validate that an xarray object satisfies its own internal invariants.

    This exists for the benefit of xarray's own test suite, but may be useful
    in external projects if they (ill-advisedly) create objects using xarray's
    private APIs.
    """
    if isinstance(xarray_obj, Variable):
        _assert_variable_invariants(xarray_obj)
    elif isinstance(xarray_obj, DataArray):
        _assert_dataarray_invariants(
            xarray_obj, check_default_indexes=check_default_indexes
        )
    elif isinstance(xarray_obj, Dataset):
        _assert_dataset_invariants(
            xarray_obj, check_default_indexes=check_default_indexes
        )
    elif isinstance(xarray_obj, Coordinates):
        _assert_dataset_invariants(
            xarray_obj.to_dataset(), check_default_indexes=check_default_indexes
        )
    else:
        raise TypeError(
            f"{type(xarray_obj)} is not a supported type for xarray invariant checks"
        )