"""Helper functions used by `array_api_extra/_funcs.py`."""

from __future__ import annotations

import io
import math
import pickle
import types
from collections.abc import Callable, Generator, Iterable, Iterator
from functools import wraps
from types import ModuleType
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Generic,
    Literal,
    ParamSpec,
    TypeAlias,
    TypeVar,
    cast,
)

from . import _compat
from ._compat import (
    array_namespace,
    is_array_api_obj,
    is_dask_namespace,
    is_jax_namespace,
    is_numpy_array,
    is_pydata_sparse_namespace,
    is_torch_namespace,
)
from ._typing import Array, Device

if TYPE_CHECKING:  # pragma: no cover
    # TODO import from typing (requires Python >=3.12 and >=3.13)
    from typing_extensions import TypeIs, override
else:

    def override(func):
        return func


P = ParamSpec("P")
T = TypeVar("T")


__all__ = [
    "asarrays",
    "capabilities",
    "eager_shape",
    "in1d",
    "is_python_scalar",
    "jax_autojit",
    "mean",
    "meta_namespace",
    "pickle_flatten",
    "pickle_unflatten",
]


def in1d(
    x1: Array,
    x2: Array,
    /,
    *,
    assume_unique: bool = False,
    invert: bool = False,
    xp: ModuleType | None = None,
) -> Array:  # numpydoc ignore=PR01,RT01
    """
    Check whether each element of an array is also present in a second array.

    Returns a boolean array the same length as `x1` that is True
    where an element of `x1` is in `x2` and False otherwise.

    This function has been adapted using the original implementation
    present in numpy:
    https://github.com/numpy/numpy/blob/v1.26.0/numpy/lib/arraysetops.py#L524-L758
    """
    if xp is None:
        xp = array_namespace(x1, x2)

    x1_shape = eager_shape(x1)
    x2_shape = eager_shape(x2)

    # This code is run to make the code significantly faster
    if x2_shape[0] < 10 * x1_shape[0] ** 0.145 and isinstance(x2, Iterable):
        if invert:
            mask = xp.ones(x1_shape[0], dtype=xp.bool, device=_compat.device(x1))
            for a in x2:
                mask &= x1 != a
        else:
            mask = xp.zeros(x1_shape[0], dtype=xp.bool, device=_compat.device(x1))
            for a in x2:
                mask |= x1 == a
        return mask

    rev_idx = xp.empty(0)  # placeholder
    if not assume_unique:
        x1, rev_idx = xp.unique_inverse(x1)
        x2 = xp.unique_values(x2)

    ar = xp.concat((x1, x2))
    device_ = _compat.device(ar)
    # We need this to be a stable sort.
    order = xp.argsort(ar, stable=True)
    reverse_order = xp.argsort(order, stable=True)
    sar = xp.take(ar, order, axis=0)
    ar_size = _compat.size(sar)
    assert ar_size is not None, "xp.unique*() on lazy backends raises"
    if ar_size >= 1:
        bool_ar = sar[1:] != sar[:-1] if invert else sar[1:] == sar[:-1]
    else:
        bool_ar = xp.asarray([False]) if invert else xp.asarray([True])
    flag = xp.concat((bool_ar, xp.asarray([invert], device=device_)))
    ret = xp.take(flag, reverse_order, axis=0)

    if assume_unique:
        return ret[: x1.shape[0]]
    return xp.take(ret, rev_idx, axis=0)


def mean(
    x: Array,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    keepdims: bool = False,
    xp: ModuleType | None = None,
) -> Array:  # numpydoc ignore=PR01,RT01
    """
    Complex mean, https://github.com/data-apis/array-api/issues/846.
    """
    if xp is None:
        xp = array_namespace(x)

    if xp.isdtype(x.dtype, "complex floating"):
        x_real = xp.real(x)
        x_imag = xp.imag(x)
        mean_real = xp.mean(x_real, axis=axis, keepdims=keepdims)
        mean_imag = xp.mean(x_imag, axis=axis, keepdims=keepdims)
        return mean_real + (mean_imag * xp.asarray(1j))
    return xp.mean(x, axis=axis, keepdims=keepdims)


def is_python_scalar(x: object) -> TypeIs[complex]:  # numpydoc ignore=PR01,RT01
    """Return True if `x` is a Python scalar, False otherwise."""
    # isinstance(x, float) returns True for np.float64
    # isinstance(x, complex) returns True for np.complex128
    # bool is a subclass of int
    return isinstance(x, int | float | complex) and not is_numpy_array(x)


def asarrays(
    a: Array | complex,
    b: Array | complex,
    xp: ModuleType,
) -> tuple[Array, Array]:
    """
    Ensure both `a` and `b` are arrays.

    If `b` is a python scalar, it is converted to the same dtype as `a`, and vice versa.

    Behavior is not specified when mixing a Python ``float`` and an array with an
    integer data type; this may give ``float32``, ``float64``, or raise an exception.
    Behavior is implementation-specific.

    Similarly, behavior is not specified when mixing a Python ``complex`` and an array
    with a real-valued data type; this may give ``complex64``, ``complex128``, or raise
    an exception. Behavior is implementation-specific.

    Parameters
    ----------
    a, b : Array | int | float | complex | bool
        Input arrays or scalars. At least one must be an array.
    xp : array_namespace, optional
        The standard-compatible namespace for `x`. Default: infer.

    Returns
    -------
    Array, Array
        The input arrays, possibly converted to arrays if they were scalars.

    See Also
    --------
    mixing-arrays-with-python-scalars : Array API specification for the behavior.
    """
    a_scalar = is_python_scalar(a)
    b_scalar = is_python_scalar(b)
    if not a_scalar and not b_scalar:
        # This includes misc. malformed input e.g. str
        return a, b  # type: ignore[return-value]

    swap = False
    if a_scalar:
        swap = True
        b, a = a, b

    if is_array_api_obj(a):
        # a is an Array API object
        # b is a int | float | complex | bool
        xa = a

        # https://data-apis.org/array-api/draft/API_specification/type_promotion.html#mixing-arrays-with-python-scalars
        same_dtype = {
            bool: "bool",
            int: ("integral", "real floating", "complex floating"),
            float: ("real floating", "complex floating"),
            complex: "complex floating",
        }
        kind = same_dtype[type(cast(complex, b))]
        if xp.isdtype(a.dtype, kind):
            xb = xp.asarray(b, dtype=a.dtype, device=_compat.device(a))
        else:
            # Undefined behaviour. Let the function deal with it, if it can.
            xb = xp.asarray(b, device=_compat.device(a))

    else:
        # Neither a nor b are Array API objects.
        # Note: we can only reach this point when one explicitly passes
        # xp=xp to the calling function; otherwise we fail earlier on
        # array_namespace(a, b).
        xa, xb = xp.asarray(a), xp.asarray(b)

    return (xb, xa) if swap else (xa, xb)


def ndindex(*x: int) -> Generator[tuple[int, ...]]:
    """
    Generate all N-dimensional indices for a given array shape.

    Given the shape of an array, an ndindex instance iterates over the N-dimensional
    index of the array. At each iteration a tuple of indices is returned, the last
    dimension is iterated over first.

    This has an identical API to numpy.ndindex.

    Parameters
    ----------
    *x : int
        The shape of the array.
    """
    if not x:
        yield ()
        return
    for i in ndindex(*x[:-1]):
        for j in range(x[-1]):
            yield *i, j


def eager_shape(x: Array, /, axis: int | None = None) -> tuple[int, ...]:
    """
    Return shape of an array. Raise if shape is not fully defined.

    Parameters
    ----------
    x : Array
        Input array.
    axis : int, optional
        If provided, only returns the tuple (shape[axis],).

    Returns
    -------
    tuple[int, ...]
        Shape of the array.
    """
    shape = x.shape
    if axis is not None:
        s = shape[axis]
        # Dask arrays uses non-standard NaN instead of None
        if s is None or math.isnan(s):
            msg = f"Unsupported lazy shape for axis {axis}"
            raise TypeError(msg)
        return (s,)

    if any(s is None or math.isnan(s) for s in shape):
        msg = "Unsupported lazy shape"
        raise TypeError(msg)
    return cast(tuple[int, ...], shape)


def meta_namespace(
    *arrays: Array | complex | None, xp: ModuleType | None = None
) -> ModuleType:
    """
    Get the namespace of Dask chunks.

    On all other backends, just return the namespace of the arrays.

    Parameters
    ----------
    *arrays : Array | int | float | complex | bool | None
        Input arrays.
    xp : array_namespace, optional
        The standard-compatible namespace for the input arrays. Default: infer.

    Returns
    -------
    array_namespace
        If xp is Dask, the namespace of the Dask chunks;
        otherwise, the namespace of the arrays.
    """
    xp = array_namespace(*arrays) if xp is None else xp
    if not is_dask_namespace(xp):
        return xp
    # Quietly skip scalars and None's
    metas = [cast(Array | None, getattr(a, "_meta", None)) for a in arrays]
    return array_namespace(*metas)


def capabilities(
    xp: ModuleType, *, device: Device | None = None
) -> dict[str, int | None]:
    """
    Return patched ``xp.__array_namespace_info__().capabilities()``.

    TODO this helper should be eventually removed once all the special cases
    it handles are fixed in the respective backends.

    Parameters
    ----------
    xp : array_namespace
        The standard-compatible namespace.
    device : Device, optional
        The device to use.

    Returns
    -------
    dict
        Capabilities of the namespace.
    """
    out = xp.__array_namespace_info__().capabilities()
    if is_pydata_sparse_namespace(xp):
        if out["boolean indexing"]:
            # FIXME https://github.com/pydata/sparse/issues/876
            # boolean indexing is supported, but not when the index is a sparse array.
            # boolean indexing by list or numpy array is not part of the Array API.
            out = out.copy()
            out["boolean indexing"] = False
    elif is_jax_namespace(xp):
        if out["boolean indexing"]:  # pragma: no cover
            # Backwards compatibility with jax <0.6.0
            # https://github.com/jax-ml/jax/issues/27418
            out = out.copy()
            out["boolean indexing"] = False
    elif is_torch_namespace(xp):
        # FIXME https://github.com/data-apis/array-api/issues/945
        device = xp.get_default_device() if device is None else xp.device(device)
        if device.type == "meta":  # type: ignore[union-attr]  # pyright: ignore[reportAttributeAccessIssue,reportOptionalMemberAccess]
            out = out.copy()
            out["boolean indexing"] = False
            out["data-dependent shapes"] = False

    return out


_BASIC_PICKLED_TYPES = frozenset((
    bool, int, float, complex, str, bytes, bytearray,
    list, tuple, dict, set, frozenset, range, slice,
    types.NoneType, types.EllipsisType,
))  # fmt: skip
_BASIC_REST_TYPES = frozenset((
    type, types.BuiltinFunctionType, types.FunctionType, types.ModuleType
))  # fmt: skip

FlattenRest: TypeAlias = tuple[object, ...]


def pickle_flatten(
    obj: object, cls: type[T] | tuple[type[T], ...]
) -> tuple[list[T], FlattenRest]:
    """
    Use the pickle machinery to extract objects out of an arbitrary container.

    Unlike regular ``pickle.dumps``, this function always succeeds.

    Parameters
    ----------
    obj : object
        The object to pickle.
    cls : type | tuple[type, ...]
        One or multiple classes to extract from the object.
        The instances of these classes inside ``obj`` will not be pickled.

    Returns
    -------
    instances : list[cls]
        All instances of ``cls`` found inside ``obj`` (not pickled).
    rest
        Opaque object containing the pickled bytes plus all other objects where
        ``__reduce__`` / ``__reduce_ex__`` is either not implemented or raised.
        These are unpickleable objects, types, modules, and functions.

        This object is *typically* hashable save for fairly exotic objects
        that are neither pickleable nor hashable.

        This object is pickleable if everything except ``instances`` was pickleable
        in the input object.

    See Also
    --------
    pickle_unflatten : Reverse function.

    Examples
    --------
    >>> class A:
    ...     def __repr__(self):
    ...         return "<A>"
    >>> class NS:
    ...     def __repr__(self):
    ...         return "<NS>"
    ...     def __reduce__(self):
    ...         assert False, "not serializable"
    >>> obj = {1: A(), 2: [A(), NS(), A()]}
    >>> instances, rest = pickle_flatten(obj, A)
    >>> instances
    [<A>, <A>, <A>]
    >>> pickle_unflatten(instances, rest)
    {1: <A>, 2: [<A>, <NS>, <A>]}

    This can be also used to swap inner objects; the only constraint is that
    the number of objects in and out must be the same:

    >>> pickle_unflatten(["foo", "bar", "baz"], rest)
    {1: "foo", 2: ["bar", <NS>, "baz"]}
    """
    instances: list[T] = []
    rest: list[object] = []

    class Pickler(pickle.Pickler):  # numpydoc ignore=GL08
        """
        Use the `pickle.Pickler.persistent_id` hook to extract objects.
        """

        @override
        def persistent_id(
            self, obj: object
        ) -> Literal[0, 1, None]:  # numpydoc ignore=GL08
            if isinstance(obj, cls):
                instances.append(obj)  # type: ignore[arg-type]
                return 0

            typ_ = type(obj)
            if typ_ in _BASIC_PICKLED_TYPES:  # No subclasses!
                # If obj is a collection, recursively descend inside it
                return None
            if typ_ in _BASIC_REST_TYPES:
                rest.append(obj)
                return 1

            try:
                # Note: a class that defines __slots__ without defining __getstate__
                # cannot be pickled with __reduce__(), but can with __reduce_ex__(5)
                _ = obj.__reduce_ex__(pickle.HIGHEST_PROTOCOL)
            except Exception:  # pylint: disable=broad-exception-caught
                rest.append(obj)
                return 1

            # Object can be pickled. Let the Pickler recursively descend inside it.
            return None

    f = io.BytesIO()
    p = Pickler(f, protocol=pickle.HIGHEST_PROTOCOL)
    p.dump(obj)
    return instances, (f.getvalue(), *rest)


def pickle_unflatten(instances: Iterable[object], rest: FlattenRest) -> Any:
    """
    Reverse of ``pickle_flatten``.

    Parameters
    ----------
    instances : Iterable
        Inner objects to be reinserted into the flattened container.
    rest : FlattenRest
        Extra bits, as returned by ``pickle_flatten``.

    Returns
    -------
    object
        The outer object originally passed to ``pickle_flatten`` after a
        pickle->unpickle round-trip.

    See Also
    --------
    pickle_flatten : Serializing function.
    pickle.loads : Standard unpickle function.

    Notes
    -----
    The `instances` iterable must yield at least the same number of elements as the ones
    returned by ``pickle_flatten``, but the elements do not need to be the same objects
    or even the same types of objects. Excess elements, if any, will be left untouched.
    """
    iters = iter(instances), iter(rest)
    pik = cast(bytes, next(iters[1]))

    class Unpickler(pickle.Unpickler):  # numpydoc ignore=GL08
        """Mirror of the overridden Pickler in pickle_flatten."""

        @override
        def persistent_load(self, pid: Literal[0, 1]) -> object:  # numpydoc ignore=GL08
            try:
                return next(iters[pid])
            except StopIteration as e:
                msg = "Not enough objects to unpickle"
                raise ValueError(msg) from e

    f = io.BytesIO(pik)
    return Unpickler(f).load()


class _AutoJITWrapper(Generic[T]):  # numpydoc ignore=PR01
    """
    Helper of :func:`jax_autojit`.

    Wrap arbitrary inputs and outputs of the jitted function and
    convert them to/from PyTrees.
    """

    _obj: Any
    _is_iter: bool
    _registered: ClassVar[bool] = False
    __slots__: tuple[str, ...] = ("_is_iter", "_obj")

    def __init__(self, obj: T) -> None:  # numpydoc ignore=GL08
        self._register()
        if isinstance(obj, Iterator):
            self._obj = list(obj)
            self._is_iter = True
        else:
            self._obj = obj
            self._is_iter = False

    @property
    def obj(self) -> T:  # numpydoc ignore=RT01
        """Return wrapped object."""
        return iter(self._obj) if self._is_iter else self._obj

    @classmethod
    def _register(cls) -> None:  # numpydoc ignore=SS06
        """
        Register upon first use instead of at import time, to avoid
        globally importing JAX.
        """
        if not cls._registered:
            import jax

            jax.tree_util.register_pytree_node(
                cls,
                lambda instance: pickle_flatten(instance, jax.Array),  # pyright: ignore[reportUnknownArgumentType]
                lambda aux_data, children: pickle_unflatten(children, aux_data),  # pyright: ignore[reportUnknownArgumentType]
            )
            cls._registered = True


def jax_autojit(
    func: Callable[P, T],
) -> Callable[P, T]:  # numpydoc ignore=PR01,RT01,SS03
    """
    Wrap `func` with ``jax.jit``, with the following differences:

    - Python scalar arguments and return values are not automatically converted to
      ``jax.Array`` objects.
    - All non-array arguments are automatically treated as static.
      Unlike ``jax.jit``, static arguments must be either hashable or serializable with
      ``pickle``.
    - Unlike ``jax.jit``, non-array arguments and return values are not limited to
      tuple/list/dict, but can be any object serializable with ``pickle``.
    - Automatically descend into non-array arguments and find ``jax.Array`` objects
      inside them, then rebuild the arguments when entering `func`, swapping the JAX
      concrete arrays with tracer objects.
    - Automatically descend into non-array return values and find ``jax.Array`` objects
      inside them, then rebuild them downstream of exiting the JIT, swapping the JAX
      tracer objects with concrete arrays.
    - Returned iterators are immediately completely consumed.

    See Also
    --------
    jax.jit : JAX JIT compilation function.

    Notes
    -----
    These are useful choices *for testing purposes only*, which is how this function is
    intended to be used. The output of ``jax.jit`` is a C++ level callable, that
    directly dispatches to the compiled kernel after the initial call. In comparison,
    ``jax_autojit`` incurs a much higher dispatch time.

    Additionally, consider::

        def f(x: Array, y: float, plus: bool) -> Array:
            return x + y if plus else x - y

        j1 = jax.jit(f, static_argnames="plus")
        j2 = jax_autojit(f)

    In the above example, ``j2`` requires a lot less setup to be tested effectively than
    ``j1``, but on the flip side it means that it will be re-traced for every different
    value of ``y``, which likely makes it not fit for purpose in production.
    """
    import jax

    @jax.jit  # type: ignore[misc]  # pyright: ignore[reportUntypedFunctionDecorator]
    def inner(  # numpydoc ignore=GL08
        wargs: _AutoJITWrapper[Any],
    ) -> _AutoJITWrapper[T]:
        args, kwargs = wargs.obj
        res = func(*args, **kwargs)  # pyright: ignore[reportCallIssue]
        return _AutoJITWrapper(res)

    @wraps(func)
    def outer(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
        wargs = _AutoJITWrapper((args, kwargs))
        return inner(wargs).obj

    return outer
