"""
AUTORAY - backend agnostic array operations.


Copyright 2019-2021 Johnnie Gray

Licensed 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.
"""

import importlib
import functools
import itertools
import contextlib
from inspect import signature
from collections import OrderedDict, defaultdict

import numpy as _numpy


def do(fn, *args, like=None, **kwargs):
    """Do function named ``fn`` on ``(*args, **kwargs)``, peforming single
    dispatch to retrieve ``fn`` based on whichever library defines the class of
    the ``args[0]``, or the ``like`` keyword argument if specified.

    Examples
    --------

    Works on numpy arrays:

        >>> import numpy as np
        >>> x_np = np.random.uniform(size=[5])
        >>> y_np = do('sqrt', x_np)
        >>> y_np
        array([0.32464973, 0.90379787, 0.85037325, 0.88729814, 0.46768083])

        >>> type(y_np)
        numpy.ndarray

    Works on cupy arrays:

        >>> import cupy as cp
        >>> x_cp = cp.random.uniform(size=[5])
        >>> y_cp = do('sqrt', x_cp)
        >>> y_cp
        array([0.44541656, 0.88713113, 0.92626237, 0.64080557, 0.69620767])

        >>> type(y_cp)
        cupy.core.core.ndarray

    Works on tensorflow arrays:

        >>> import tensorflow as tf
        >>> x_tf = tf.random.uniform(shape=[5])
        >>> y_tf = do('sqrt', x_tf)
        >>> y_tf
        <tf.Tensor 'Sqrt_1:0' shape=(5,) dtype=float32>

        >>> type(y_tf)
        tensorflow.python.framework.ops.Tensor

    You get the idea.

    For functions that don't dispatch on the first argument you can use the
    ``like`` keyword:

        >>> do('eye', 3, like=x_tf)
        <tf.Tensor: id=91, shape=(3, 3), dtype=float32>
    """
    backend = choose_backend(fn, *args, like=like, **kwargs)
    return get_lib_fn(backend, fn)(*args, **kwargs)


# ------------------------- efficiently dispatching ------------------------- #


def _default_infer_from_sig(fn, *args, **kwargs):
    dispatch_arg = _DISPATCHERS[fn](*args, **kwargs)
    return infer_backend(dispatch_arg)


_infer_auto = _default_infer_from_sig
_global_backend = None


def _always_the_same(*args, x, **kwargs):
    return x


def get_backend():
    """Return the global backend if it has been set, otherwise ``None``.
    """
    global _global_backend
    return _global_backend


def set_backend(like):
    """Set a default global backend. The argument ``like`` can be an explicit
    backend name or an ``array``.
    """
    global _global_backend
    global _infer_auto

    if like is None:
        _global_backend = None
        _infer_auto = _default_infer_from_sig
    elif isinstance(like, str):
        _global_backend = like
        _infer_auto = functools.partial(_always_the_same, x=_global_backend)
    else:
        _global_backend = infer_backend(like)
        _infer_auto = functools.partial(_always_the_same, x=_global_backend)


@contextlib.contextmanager
def backend_like(like):
    """Context manager for setting a default backend. Currently not thread
    safe. The argument ``like`` can be an explicit backend name or an
    ``array`` to infer it from.
    """
    old_backend = get_backend()
    try:
        set_backend(like)
        yield
    finally:
        set_backend(old_backend)


_CUSTOM_BACKENDS = {}


def register_backend(cls, name):
    """Register the name (and by default the module or submodule) of a custom
    array class.
    """
    if not isinstance(cls, type):
        raise TypeError("The array class itself should be supplied.")

    global _CUSTOM_BACKENDS
    _CUSTOM_BACKENDS[cls] = name


@functools.lru_cache(None)
def _infer_class_backend_cached(T):
    if issubclass(T, _numpy.ndarray):
        return "numpy"

    if T in _CUSTOM_BACKENDS:
        return _CUSTOM_BACKENDS[T]

    lib = T.__module__.split(".")[0]

    # check if lib should mapped entirely to another lib
    backend = _BACKEND_ALIASES.get(lib, lib)

    return backend


def infer_backend(array):
    """Get the name of the library that defined the class of ``array`` - unless
    ``array`` is directly a subclass of ``numpy.ndarray``, in which case assume
    ``numpy`` is the desired backend.
    """
    return _infer_class_backend_cached(array.__class__)


def choose_backend(fn, *args, like=None, **kwargs):
    """Choose a backend based on function name, arguments, and the ``like``
    keyword argument. The default, if ``like`` is not specified, is to infer
    the backend from the function call, the default of which is simply to use
    the first argument, if no custom dispatcher is found. Otherwise the
    backend is chosen based on the ``like`` argument - which can be an explicit
    backend name or an arbitrary object.
    """
    if like is None:
        return _infer_auto(fn, *args, **kwargs)
    elif isinstance(like, str):
        return like
    else:
        return infer_backend(like)


# ------------------- importing and caching the function -------------------- #


def import_lib_fn(backend, fn):

    # first check explicitly composed functions -> if the function hasn't been
    # called directly yet, it won't have been loaded into the cache, and needs
    # generating before e.g. the ``do`` verrsion will work
    if fn in _COMPOSED_FUNCTION_GENERATORS:
        return _COMPOSED_FUNCTION_GENERATORS[fn](backend)

    try:
        # alias for global module,
        #     e.g. 'decimal' -> 'math'
        module = _MODULE_ALIASES.get(backend, backend)

        # submodule where function is found for backend,
        #     e.g. ['tensorflow', trace'] -> 'tensorflow.linalg'
        try:
            full_location = _SUBMODULE_ALIASES[backend, fn]

            # if explicit submodule alias given, don't use prepended location
            #     for example, ('torch', 'linalg.svd') -> torch.svd
            only_fn = fn.split(".")[-1]

        except KeyError:
            full_location = module

            # move any prepended location into the full module path
            #     e.g. 'fn=linalg.eigh' -> ['linalg', 'eigh']
            split_fn = fn.split(".")
            full_location = ".".join([full_location] + split_fn[:-1])
            only_fn = split_fn[-1]

        # cached lookup of custom name function might take
        #     e.g. ['tensorflow', 'sum'] -> 'reduce_sum'
        fn_name = _FUNC_ALIASES.get((backend, fn), only_fn)

        # import the function into the cache
        try:
            lib = importlib.import_module(full_location)
        except ImportError:
            if "." in full_location:
                # sometimes libraries hack an attribute to look like submodule
                mod, submod = full_location.split(".")
                lib = getattr(importlib.import_module(mod), submod)
            else:
                # failed to import library at all -> catch + raise ImportError
                raise AttributeError

        # check for a custom wrapper but default to identity
        wrapper = _CUSTOM_WRAPPERS.get((backend, fn), lambda fn: fn)

        # store the function!
        lib_fn = _FUNCS[backend, fn] = wrapper(getattr(lib, fn_name))

    except AttributeError:

        # check if there is a backup function (e.g. for older library version)
        backend_alt = backend + "[alt]"
        if backend_alt in _MODULE_ALIASES:
            return import_lib_fn(backend_alt, fn)

        raise ImportError(
            f"autoray couldn't find function '{fn}' for backend '{backend}'.")

    return lib_fn


def get_lib_fn(backend, fn):
    """Cached retrieval of correct function for backend, all the logic for
    finding the correct funtion only runs the first time.

    Parameters
    ----------
    backend : str
        The module defining the array class to dispatch on.
    fn : str
        The function to retrieve.

    Returns
    -------
    callable
    """

    try:
        lib_fn = _FUNCS[backend, fn]
    except KeyError:
        lib_fn = import_lib_fn(backend, fn)
    return lib_fn


# ---------------------- special top level functions ------------------------ #


def conj(x):
    """Array conjugate."""
    return do("conj", x)


def transpose(x, *args):
    """Array transpose."""
    return do("transpose", x, *args)


def dag(x):
    """Array Hermitian transpose."""
    try:
        return x.H
    except AttributeError:
        return do("conj", do("transpose", x))


def real(x):
    """Array real part."""
    return do("real", x)


def imag(x):
    """Array imaginary part."""
    return do("imag", x)


def reshape(x, shape):
    """Array reshaped."""
    try:
        return x.reshape(shape)
    except AttributeError:
        return do("reshape", x, shape)


def to_backend_dtype(dtype_name, like):
    """Turn string specifier ``dtype_name`` into dtype of backend ``like``."""
    if not isinstance(like, str):
        like = infer_backend(like)

    try:
        return get_lib_fn(like, dtype_name)
    except ImportError:
        return dtype_name


def get_dtype_name(x):
    """Find string specifier ``dtype_name`` of array ``x``."""
    try:
        return x.dtype.name
    except AttributeError:
        # let modules provide their own
        return do("get_dtype_name", x, like=x)


_COMPLEX_DTYPES = {'complex64', 'complex128'}
_DOUBLE_DTYPES = {'float64', 'complex128'}
_DTYPE_MAP = {
    (False, False): 'float32',
    (False, True): 'float64',
    (True, False): 'complex64',
    (True, True): 'complex128',
}


def get_common_dtype(*arrays):
    """Compute the minimal dtype sufficient for ``arrays``.
    """
    dtypes = set(map(get_dtype_name, arrays))
    has_complex = not _COMPLEX_DTYPES.isdisjoint(dtypes)
    has_double = not _DOUBLE_DTYPES.isdisjoint(dtypes)
    return _DTYPE_MAP[has_complex, has_double]


def astype(x, dtype_name, **kwargs):
    """Cast array as type ``dtype_name`` - tries ``x.astype`` first."""
    dtype = to_backend_dtype(dtype_name, like=x)
    try:
        return x.astype(dtype, **kwargs)
    except AttributeError:
        return do("astype", x, dtype, **kwargs)


def to_numpy(x):
    """Get a numpy version of array ``x``."""
    return do("to_numpy", x)


# -------------------------- some common wrappers --------------------------- #


def svd_not_full_matrices_wrapper(fn):
    @functools.wraps(fn)
    def default_not_full_matrices(*args, **kwargs):
        kwargs.setdefault("full_matrices", False)
        return fn(*args, **kwargs)

    return default_not_full_matrices


def svd_sUV_to_UsVH_wrapper(fn):
    @functools.wraps(fn)
    def numpy_like(*args, **kwargs):
        s, U, V = fn(*args, **kwargs)
        return U, s, dag(V)

    return numpy_like


def svd_UsV_to_UsVH_wrapper(fn):
    @functools.wraps(fn)
    def numpy_like(*args, **kwargs):
        U, s, V = fn(*args, **kwargs)
        return U, s, dag(V)

    return numpy_like


def svd_manual_full_matrices_kwarg(fn):
    @functools.wraps(fn)
    def numpy_like(*args, full_matrices=False, **kwargs):
        U, s, VH = fn(*args, **kwargs)

        if not full_matrices:
            U, VH = U[:, : s.size], VH[: s.size, :]

        return U, s, VH

    return numpy_like


def qr_allow_fat(fn):
    @functools.wraps(fn)
    def numpy_like(a, **kwargs):
        m, n = a.shape

        if m >= n:
            # square or thin
            return fn(a, **kwargs)

        Q, R_sq = fn(a[:, :m])
        R_r = dag(Q) @ a[:, m:]
        R = do("concatenate", (R_sq, R_r), axis=1, like=a)

        return Q, R

    return numpy_like


def tril_to_band_part(fn):
    @functools.wraps(fn)
    def numpy_like(x, k=0):

        if k < 0:
            raise ValueError(
                "'k' must be positive to recreate 'numpy.tril' "
                "behaviour with 'tensorflow.matrix_band_part'."
            )

        return fn(x, -1, k)

    return numpy_like


def triu_to_band_part(fn):
    @functools.wraps(fn)
    def numpy_like(x, k=0):

        if k > 0:
            raise ValueError(
                "'k' must be negative to recreate 'numpy.triu' "
                "behaviour with 'tensorflow.matrix_band_part'."
            )

        return fn(x, -k, -1)

    return numpy_like


def cholesky_lower(fn):
    @functools.wraps(fn)
    def cholesky_numpy_like(a):
        return fn(a, lower=True)

    return cholesky_numpy_like


def binary_allow_1d_rhs_wrap(fn):
    @functools.wraps(fn)
    def allow_1d_rhs(a, b):
        need_to_convert = len(a.shape) != len(b.shape)
        if need_to_convert:
            b = reshape(b, (*b.shape, 1))
        x = fn(a, b)
        if need_to_convert:
            x = reshape(x, x.shape[:-1])
        return x

    return allow_1d_rhs


def scale_random_uniform_manually(fn):
    @functools.wraps(fn)
    def numpy_like(low=0.0, high=1.0, size=None, dtype=None, **kwargs):
        if size is None:
            size = ()

        x = fn(size, **kwargs)

        if (low != 0.0) or (high != 1.0):
            x = (high - low) * x + low

        if (dtype is not None) and get_dtype_name(x) != dtype:
            x = astype(x, dtype)
        return x

    return numpy_like


def scale_random_normal_manually(fn):
    @functools.wraps(fn)
    def numpy_like(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs):
        if size is None:
            size = ()

        x = fn(size=size, **kwargs)

        if (loc != 0.0) or (scale != 1.0):
            x = scale * x + loc

        if (dtype is not None) and get_dtype_name(x) != dtype:
            x = astype(x, dtype)
        return x

    return numpy_like


def with_dtype_wrapper(fn):
    """Add ability to handle `dtype` keyword.
    If not None, `dtype` should be specified as a string, otherwise conversion
    will happen regardless.
    """

    @functools.wraps(fn)
    def with_dtype(*args, dtype=None, **kwargs):
        A = fn(*args, **kwargs)
        if (dtype is not None) and (dtype != get_dtype_name(A)):
            A = astype(A, dtype)
        return A

    return with_dtype


def translate_wrapper(fn, translator):
    """Wrap a function to match the api of another according to a translation.
    The ``translator`` entries in the form of an ordered dict should have
    entries like:

        (desired_kwarg: (backend_kwarg, default_value))

    with the order defining the args of the function.
    """

    @functools.wraps(fn)
    def translated_function(*args, **kwargs):
        new_kwargs = {}
        translation = translator.copy()

        # convert args
        for arg_value in args:
            new_arg_name = translation.popitem(last=False)[1][0]
            new_kwargs[new_arg_name] = arg_value

        # convert kwargs -  but only those in the translation
        for key, value in kwargs.items():
            try:
                new_kwargs[translation.pop(key)[0]] = value
            except KeyError:
                new_kwargs[key] = value

        # set remaining default kwargs
        for key, value in translation.items():
            new_kwargs[value[0]] = value[1]

        return fn(**new_kwargs)

    return translated_function


def make_translator(t):
    return functools.partial(translate_wrapper, translator=OrderedDict(t))


def complex_add_re_im(re, im):
    return re + 1j * im


def allclose(x, y, rtol=1e-05, atol=1e-08):
    return do("all", do("abs", x - y) <= atol + rtol * do("abs", y))


# ----------------------------- Custom dispatchers -------------------------- #


def register_dispatch(fun, dispatcher):
    """Register a new dispatcher.

    This is useful in case the backend to be used by a function cannot be
    inferred from the first argument.
    """
    _DISPATCHERS[fun] = dispatcher


def default_dispatcher(*args, **kwargs):
    """Try to infer backend from first argument passed to function."""
    return args[0]


# lookup of custom dispatcher methods, for cases when backend cannot be
#     inferred accurately from first argument.
_DISPATCHERS = defaultdict(lambda: default_dispatcher)


def join_array_dispatcher(*args, **kwargs):
    """Dispatcher for functions where first argument is a sequence."""
    try:
        return args[0][0]
    except (TypeError, ValueError):
        # user passed an empty sequence, or something non-iterable
        # try to infer backend from first argument as fallback
        return args[0]


# List of functions listed in numpy API as array joining operations
register_dispatch("concatenate", join_array_dispatcher)
register_dispatch("stack", join_array_dispatcher)
register_dispatch("block", join_array_dispatcher)
register_dispatch("vstack", join_array_dispatcher)
register_dispatch("hstack", join_array_dispatcher)
register_dispatch("dstack", join_array_dispatcher)
register_dispatch("column_stack", join_array_dispatcher)
register_dispatch("row_stack", join_array_dispatcher)


def einsum_dispatcher(*args, **kwargs):
    """Dispatcher for handling einsum.

    einsum can either take string as first argument, in which case backend
    should be inferred from second argument. Or it can take an array as first
    argument, which should be used to infer backend.
    """
    if isinstance(args[0], str):
        return args[1]
    return args[0]


register_dispatch("einsum", einsum_dispatcher)

# --------------- object to act as drop-in replace for numpy ---------------- #

_partial_functions = {}


class NumpyMimic:
    """A class to mimic the syntax of using `numpy` directly."""

    def __init__(self, submodule=None):
        self.submodule = submodule

    def __getattribute__(self, fn):

        # look out for certain submodules which are not functions
        if fn == "linalg":
            return numpy_linalg
        if fn == "random":
            return numpy_random

        # if this is the e.g. linalg mimic, preprend 'linalg.'
        submod = object.__getattribute__(self, "submodule")
        if submod is not None:
            fn = ".".join((submod, fn))

        # cache the correct partial function
        try:
            pfn = _partial_functions[fn]
        except KeyError:
            pfn = _partial_functions[fn] = functools.partial(do, fn)

        return pfn

    @staticmethod
    def __repr__():
        return "<autoray.numpy>"


numpy = NumpyMimic()
numpy_linalg = NumpyMimic("linalg")
numpy_random = NumpyMimic("random")


# --------------------------------------------------------------------------- #
#                    storage & backend specific functions                     #
# --------------------------------------------------------------------------- #

# lookup for mapping entire lib to another
_BACKEND_ALIASES = {}

# global (non function specific) aliases
_MODULE_ALIASES = {}

# lookup for when functions are elsewhere than the expected location
_SUBMODULE_ALIASES = {}

# lookup for when functions are simply called something else
_FUNC_ALIASES = {}

# custom wrappers for when functions don't just have different location or
#     name. For example, when kwargs need to be translated or results modified
_CUSTOM_WRAPPERS = {}

# actual cache of funtions to use - this is populated lazily and can be used
#     to directly set an implementation of a function for a specific backend
_FUNCS = {}

_COMPOSED_FUNCTION_GENERATORS = {}

# ------------------------------ standard-lib ------------------------------- #

_MODULE_ALIASES["decimal"] = "math"
_MODULE_ALIASES["builtins"] = "numpy"


_builtin_dtype_lookup = {
    int: 'int64',
    float: 'float64',
    complex: 'complex128',
}


def builtins_get_dtype_name(x):
    return _builtin_dtype_lookup[x.__class__]


_FUNCS["builtins", "get_dtype_name"] = builtins_get_dtype_name


# ---------------------------------- numpy ---------------------------------- #


def numpy_to_numpy(x):
    return do("asarray", x, like="numpy")


_FUNCS["numpy", "to_numpy"] = numpy_to_numpy
_FUNCS["numpy", "complex"] = complex_add_re_im
_FUNCS["builtins", "to_numpy"] = numpy_to_numpy
_SUBMODULE_ALIASES["numpy", "linalg.expm"] = "scipy.linalg"
_CUSTOM_WRAPPERS["numpy", "linalg.svd"] = svd_not_full_matrices_wrapper
_CUSTOM_WRAPPERS["numpy", "random.normal"] = with_dtype_wrapper
_CUSTOM_WRAPPERS["numpy", "random.uniform"] = with_dtype_wrapper

# ---------------------------------- cupy ----------------------------------- #


def cupy_to_numpy(x):  # pragma: no cover
    return x.get()


_FUNCS["cupy", "to_numpy"] = cupy_to_numpy
_FUNCS["cupy", "complex"] = complex_add_re_im
_CUSTOM_WRAPPERS["cupy", "linalg.svd"] = svd_not_full_matrices_wrapper


# ----------------------------------- jax ----------------------------------- #


_JAX_RANDOM_KEY = None


def jax_random_seed(seed=None):
    from jax.random import PRNGKey

    global _JAX_RANDOM_KEY
    if seed is None:
        from random import SystemRandom

        seed = SystemRandom().randint(
            -(2 ** 63), 2 ** 63 - 1
        )  # inclusive high
    _JAX_RANDOM_KEY = PRNGKey(seed)


def jax_random_get_key():
    from jax.random import split

    global _JAX_RANDOM_KEY
    if _JAX_RANDOM_KEY is None:
        jax_random_seed()
    _JAX_RANDOM_KEY, subkey = split(_JAX_RANDOM_KEY)
    return subkey


def jax_random_uniform(low=0.0, high=1.0, size=None, **kwargs):
    from jax.random import uniform

    if size is None:
        size = ()
    return uniform(
        jax_random_get_key(), shape=size, minval=low, maxval=high, **kwargs
    )


def jax_random_normal(loc=0.0, scale=1.0, size=None, **kwargs):
    from jax.random import normal

    if size is None:
        size = ()
    x = normal(jax_random_get_key(), shape=size, **kwargs)
    if scale != 1.0:
        x *= scale
    if loc != 0.0:
        x += loc
    return x


def jax_to_numpy(x):
    return do("asarray", x, like="numpy")


_BACKEND_ALIASES["jaxlib"] = "jax"
_MODULE_ALIASES["jax"] = "jax.numpy"
_SUBMODULE_ALIASES["jax", "complex"] = "jax.lax"
_SUBMODULE_ALIASES["jax", "linalg.expm"] = "jax.scipy.linalg"
_CUSTOM_WRAPPERS["jax", "linalg.qr"] = qr_allow_fat
_CUSTOM_WRAPPERS["jax", "linalg.svd"] = svd_not_full_matrices_wrapper
_FUNCS["jax", "to_numpy"] = jax_to_numpy
_FUNCS["jax", "random.seed"] = jax_random_seed
_FUNCS["jax", "random.uniform"] = jax_random_uniform
_FUNCS["jax", "random.normal"] = jax_random_normal


# -------------------------------- autograd --------------------------------- #

_MODULE_ALIASES["autograd"] = "autograd.numpy"
_CUSTOM_WRAPPERS["autograd", "linalg.svd"] = svd_not_full_matrices_wrapper
_FUNCS["autograd", "complex"] = complex_add_re_im


# ---------------------------------- dask ----------------------------------- #


def dask_to_numpy(x):
    return x.compute()


def dask_eye_wrapper(fn):
    # Make M work as positional argument
    @functools.wraps(fn)
    def numpy_like(N, M=None, **kwargs):
        return fn(N, M=M, **kwargs)

    return numpy_like


_FUNCS["dask", "to_numpy"] = dask_to_numpy
_FUNCS["dask", "complex"] = complex_add_re_im
_FUNC_ALIASES["dask", "abs"] = "absolute"
_MODULE_ALIASES["dask"] = "dask.array"
_CUSTOM_WRAPPERS["dask", "linalg.svd"] = svd_manual_full_matrices_kwarg
_CUSTOM_WRAPPERS["dask", "linalg.cholesky"] = cholesky_lower
_CUSTOM_WRAPPERS["dask", "random.normal"] = with_dtype_wrapper
_CUSTOM_WRAPPERS["dask", "random.uniform"] = with_dtype_wrapper
_CUSTOM_WRAPPERS["dask", "eye"] = dask_eye_wrapper

# ---------------------------------- mars ----------------------------------- #


def mars_to_numpy(x):
    return x.to_numpy()


_FUNCS["mars", "to_numpy"] = mars_to_numpy
_FUNCS["mars", "complex"] = complex_add_re_im
_MODULE_ALIASES["mars"] = "mars.tensor"
_CUSTOM_WRAPPERS["mars", "linalg.cholesky"] = cholesky_lower


# ----------------------------------- ctf ----------------------------------- #


def ctf_array(x):
    return do("astensor", x, like="ctf")


def ctf_to_numpy(x):
    return x.to_nparray()


def ctf_count_nonzero(x):
    return (x != 0).astype(int).sum()


def ctf_get_dtype_name(x):
    return x.dtype.__name__


_FUNCS["ctf", "array"] = ctf_array
_FUNCS["ctf", "complex"] = complex_add_re_im
_FUNCS["ctf", "allclose"] = allclose
_FUNCS["ctf", "to_numpy"] = ctf_to_numpy
_FUNCS["ctf", "count_nonzero"] = ctf_count_nonzero
_FUNCS["ctf", "get_dtype_name"] = ctf_get_dtype_name

_SUBMODULE_ALIASES["ctf", "float32"] = "numpy"
_SUBMODULE_ALIASES["ctf", "float64"] = "numpy"
_SUBMODULE_ALIASES["ctf", "complex64"] = "numpy"
_SUBMODULE_ALIASES["ctf", "complex128"] = "numpy"
_SUBMODULE_ALIASES["ctf", "linalg.svd"] = "ctf"
_SUBMODULE_ALIASES["ctf", "linalg.eigh"] = "ctf"
_SUBMODULE_ALIASES["ctf", "linalg.qr"] = "ctf"
_SUBMODULE_ALIASES["ctf", "linalg.norm"] = "ctf"

_FUNC_ALIASES["ctf", "random.uniform"] = "random"

_CUSTOM_WRAPPERS["ctf", "random.uniform"] = scale_random_uniform_manually


# ------------------------------- sparse------------------------------------- #


def sparse_array(x):
    return do("COO.from_numpy", x, like="sparse")


def sparse_to_numpy(x):
    return x.todense()


def sparse_complex(x, y):
    return x + 1j * y


def sparse_transpose(x, axes=None):
    return x.transpose(axes)


def sparse_reshape(x, shape):
    return x.reshape(shape)


def sparse_sum(x, axis=None, keepdims=False, dtype=None, out=None):
    return x.sum(axis=axis, keepdims=keepdims, dtype=dtype, out=out)


def sparse_prod(x, axis=None, keepdims=False, dtype=None, out=None):
    return x.prod(axis=axis, keepdims=keepdims, dtype=dtype, out=out)


def sparse_conj(x):
    return x.conj()


def sparse_real(x):
    return x.real


def sparse_imag(x):
    return x.imag


def sparse_count_nonzero(x):
    return x.nnz


def sparse_random_uniform(low=0.0, high=1.0, size=None, dtype=None, **kwargs):
    def rvs(nnz):
        return do(
            "random.uniform", low, high, (nnz,), dtype=dtype, like="numpy"
        )

    return do("random", size, data_rvs=rvs, **kwargs, like="sparse")


def sparse_random_normal(loc=0.0, scale=1.0, size=None, dtype=None, **kwargs):
    def rvs(nnz):
        return do(
            "random.normal", loc, scale, (nnz,), dtype=dtype, like="numpy"
        )

    return do("random", size, data_rvs=rvs, **kwargs, like="sparse")


_FUNCS["sparse", "array"] = sparse_array
_FUNCS["sparse", "to_numpy"] = sparse_to_numpy
_FUNCS["sparse", "transpose"] = sparse_transpose
_FUNCS["sparse", "reshape"] = sparse_reshape
_FUNCS["sparse", "sum"] = sparse_sum
_FUNCS["sparse", "prod"] = sparse_prod
_FUNCS["sparse", "conj"] = sparse_conj
_FUNCS["sparse", "real"] = sparse_real
_FUNCS["sparse", "imag"] = sparse_imag
_FUNCS["sparse", "complex"] = sparse_complex
_FUNCS["sparse", "count_nonzero"] = sparse_count_nonzero
_FUNCS["sparse", "random.uniform"] = sparse_random_uniform
_FUNCS["sparse", "random.normal"] = sparse_random_normal

# sparse uses numpys __array_func__ interface
for f in (
    "log",
    "log2",
    "log10",
    "exp",
    "sqrt",
    "sign",
    "sin",
    "cos",
    "tan",
    "arcsin",
    "arccos",
    "arctan",
    "sinh",
    "cosh",
    "tanh",
    "arcsinh",
    "arccosh",
    "arctanh",
    "tensordot",
    # NB put tensordot here, as sparse.tensordot can produce dense (numpy)
    # arrays but errors when both inputs are dense - we want nested calls to
    # tensordot to handle this
):
    _SUBMODULE_ALIASES["sparse", f] = "numpy"


# ------------------------------- tensorflow -------------------------------- #


def tensorflow_to_numpy(x):
    return x.numpy()


def tensorflow_pad_wrap(tf_pad):
    def numpy_like(array, pad_width, mode="constant", constant_values=0):
        if mode != "constant":
            raise NotImplementedError

        try:
            if len(pad_width) == 1:
                pad_width = pad_width * len(array.shape)
        except TypeError:
            pad_width = ((pad_width, pad_width),) * len(array.shape)

        return tf_pad(
            array, pad_width, mode="CONSTANT", constant_values=constant_values
        )

    return numpy_like


def tensorflow_where_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(condition, x=None, y=None, **kwargs):
        return tuple(transpose(fn(condition, x, y, **kwargs)))

    return numpy_like


def tensorflow_split_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(ary, indices_or_sections, axis=0, **kwargs):
        if isinstance(indices_or_sections, int):
            return fn(ary, indices_or_sections, axis=axis, **kwargs)
        else:
            diff = do(
                "diff",
                indices_or_sections,
                prepend=0,
                append=ary.shape[axis],
                like="numpy",
            )
            diff = list(diff)
            return fn(ary, diff, axis=axis)

    return numpy_like


_FUNCS["tensorflow", "to_numpy"] = tensorflow_to_numpy

_SUBMODULE_ALIASES["tensorflow", "log"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "conj"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "real"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "imag"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "power"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "count_nonzero"] = "tensorflow.math"
_SUBMODULE_ALIASES["tensorflow", "diag"] = "tensorflow.linalg"
_SUBMODULE_ALIASES["tensorflow", "trace"] = "tensorflow.linalg"
_SUBMODULE_ALIASES["tensorflow", "tril"] = "tensorflow.linalg"
_SUBMODULE_ALIASES["tensorflow", "triu"] = "tensorflow.linalg"

_FUNC_ALIASES["tensorflow", "sum"] = "reduce_sum"
_FUNC_ALIASES["tensorflow", "min"] = "reduce_min"
_FUNC_ALIASES["tensorflow", "max"] = "reduce_max"
_FUNC_ALIASES["tensorflow", "mean"] = "reduce_mean"
_FUNC_ALIASES["tensorflow", "prod"] = "reduce_prod"
_FUNC_ALIASES["tensorflow", "concatenate"] = "concat"
_FUNC_ALIASES["tensorflow", "clip"] = "clip_by_value"
_FUNC_ALIASES["tensorflow", "arange"] = "range"
_FUNC_ALIASES["tensorflow", "tril"] = "band_part"
_FUNC_ALIASES["tensorflow", "triu"] = "band_part"
_FUNC_ALIASES["tensorflow", "diag"] = "tensor_diag"
_FUNC_ALIASES["tensorflow", "array"] = "convert_to_tensor"
_FUNC_ALIASES["tensorflow", "astype"] = "cast"
_FUNC_ALIASES["tensorflow", "power"] = "pow"
_FUNC_ALIASES["tensorflow", "take"] = "gather"

_CUSTOM_WRAPPERS["tensorflow", "linalg.svd"] = svd_sUV_to_UsVH_wrapper
_CUSTOM_WRAPPERS["tensorflow", "linalg.qr"] = qr_allow_fat
_CUSTOM_WRAPPERS["tensorflow", "linalg.solve"] = binary_allow_1d_rhs_wrap
_CUSTOM_WRAPPERS["tensorflow", "matmul"] = binary_allow_1d_rhs_wrap
_CUSTOM_WRAPPERS["tensorflow", "tril"] = tril_to_band_part
_CUSTOM_WRAPPERS["tensorflow", "triu"] = triu_to_band_part
_CUSTOM_WRAPPERS["tensorflow", "pad"] = tensorflow_pad_wrap
_CUSTOM_WRAPPERS["tensorflow", "where"] = tensorflow_where_wrap
_CUSTOM_WRAPPERS["tensorflow", "split"] = tensorflow_split_wrap
_CUSTOM_WRAPPERS["tensorflow", "random.uniform"] = make_translator(
    [
        ("low", ("minval", 0.0)),
        ("high", ("maxval", 1.0)),
        ("size", ("shape", ())),
    ]
)
_CUSTOM_WRAPPERS["tensorflow", "random.normal"] = make_translator(
    [
        ("loc", ("mean", 0.0)),
        ("scale", ("stddev", 1.0)),
        ("size", ("shape", ())),
    ]
)
_CUSTOM_WRAPPERS["tensorflow", "clip"] = make_translator(
    [
        ("a", ("t", 0.0)),
        ("a_min", ("clip_value_min",)),
        ("a_max", ("clip_value_max",)),
    ]
)


# ---------------------------------- torch ---------------------------------- #


def torch_to_numpy(x):
    return x.detach().cpu().numpy()


def torch_transpose(x, axes=None):
    if axes is None:
        axes = reversed(range(0, x.ndimension()))
    return x.permute(*axes)


def torch_count_nonzero(x):
    return do("sum", x != 0, like="torch")


def torch_astype(x, dtype):
    return x.to(dtype=to_backend_dtype(dtype, like=x))


@functools.lru_cache(None)
def _torch_get_dtype_name(dtype):
    return str(dtype).split(".")[-1]


def torch_get_dtype_name(x):
    return _torch_get_dtype_name(x.dtype)


def torch_real(x):
    # torch doesn't support calling real on real arrays
    try:
        if x.is_complex():
            return x.real
    except AttributeError:
        pass
    return x


def torch_imag(x):
    # torch doesn't support calling imag on real arrays
    try:
        if x.is_complex():
            return x.imag
    except AttributeError:
        pass
    return do("zeros_like", x, like="torch")


def torch_linalg_solve_wrap(fn):
    @binary_allow_1d_rhs_wrap
    def numpy_like(a, b):
        return fn(b, a)[0]

    return numpy_like


def torch_linalg_eigh(x):
    return tuple(do("symeig", x, eigenvectors=True, like="torch"))


def torch_linalg_eigvalsh(x):
    return do("symeig", x, eigenvectors=False, like="torch")[0]


def torch_tensordot_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(a, b, axes=2):
        return fn(a, b, dims=axes)

    return numpy_like


def torch_pad(array, pad_width, mode="constant", constant_values=0):
    if mode != "constant":
        raise NotImplementedError

    try:
        # numpy takes pads like ((0, 0), (1, 1), ... (n-1, n-1))
        # torch takes pads like (n-1, n-1, n-2, n-2, n-3, n-3, ...)
        pad = tuple(itertools.chain.from_iterable(pad_width))[::-1]

        # a single tuple was specified ((a, b),) - use for all axes
        if len(pad) == 2:
            pad = pad * array.ndimension()

    except TypeError:
        # assume int
        pad = (pad_width,) * 2 * array.ndimension()

    return do(
        "nn.functional.pad",
        array,
        pad=pad,
        mode=mode,
        value=constant_values,
        like="torch",
    )


def torch_split_wrap(fn):
    # for torch >=1.8 we can use tensor_split instead, but in current stable
    # release this function has not been added
    @functools.wraps(fn)
    def numpy_like(ary, indices_or_sections, axis=0, **kwargs):
        if isinstance(indices_or_sections, int):
            split_size = ary.shape[axis] // indices_or_sections
            return fn(ary, split_size, dim=axis, **kwargs)
        else:
            diff = do(
                "diff",
                indices_or_sections,
                prepend=0,
                append=ary.shape[axis],
                like="numpy",
            )
            diff = list(diff)
            return fn(ary, diff, dim=axis)

    return numpy_like


def torch_zeros_ones_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(shape, dtype=None, **kwargs):
        if dtype is not None:
            dtype = to_backend_dtype(dtype, like="torch")
        return fn(shape, dtype=dtype)

    return numpy_like


def torch_eye_wrap(fn):
    @functools.wraps(fn)
    def numpy_like(N, M=None, dtype=None, **kwargs):
        if dtype is not None:
            dtype = to_backend_dtype(dtype, like="torch")
        if M is not None:
            return fn(N, m=M, dtype=dtype)
        else:
            return fn(N, dtype=dtype)

    return numpy_like


_FUNCS["torch", "pad"] = torch_pad
_FUNCS["torch", "real"] = torch_real
_FUNCS["torch", "imag"] = torch_imag
_FUNCS["torch", "astype"] = torch_astype
_FUNCS["torch", "to_numpy"] = torch_to_numpy
_FUNCS["torch", "complex"] = complex_add_re_im
_FUNCS["torch", "transpose"] = torch_transpose
_FUNCS["torch", "count_nonzero"] = torch_count_nonzero
_FUNCS["torch", "get_dtype_name"] = torch_get_dtype_name

_FUNC_ALIASES["torch", "clip"] = "clamp"
_FUNC_ALIASES["torch", "power"] = "pow"
_FUNC_ALIASES["torch", "array"] = "tensor"
_FUNC_ALIASES["torch", "concatenate"] = "cat"
_FUNC_ALIASES["torch", "random.normal"] = "randn"
_FUNC_ALIASES["torch", "random.uniform"] = "rand"
_FUNC_ALIASES["torch", "take"] = "index_select"
_FUNC_ALIASES["torch", "linalg.expm"] = "matrix_exp"

_SUBMODULE_ALIASES["torch", "linalg.expm"] = "torch"
_SUBMODULE_ALIASES["torch", "random.normal"] = "torch"
_SUBMODULE_ALIASES["torch", "random.uniform"] = "torch"

_CUSTOM_WRAPPERS["torch", "linalg.svd"] = svd_not_full_matrices_wrapper
_CUSTOM_WRAPPERS["torch", "random.normal"] = scale_random_normal_manually
_CUSTOM_WRAPPERS["torch", "random.uniform"] = scale_random_uniform_manually
_CUSTOM_WRAPPERS["torch", "split"] = torch_split_wrap
_CUSTOM_WRAPPERS["torch", "tensordot"] = torch_tensordot_wrap
_CUSTOM_WRAPPERS["torch", "stack"] = make_translator(
    [("arrays", ("tensors",)), ("axis", ("dim", 0))]
)
_CUSTOM_WRAPPERS["torch", "concatenate"] = make_translator(
    [("arrays", ("tensors",)), ("axis", ("dim", 0))]
)
_CUSTOM_WRAPPERS["torch", "tril"] = make_translator(
    [("m", ("input",)), ("k", ("diagonal", 0))]
)
_CUSTOM_WRAPPERS["torch", "triu"] = make_translator(
    [("m", ("input",)), ("k", ("diagonal", 0))]
)
_CUSTOM_WRAPPERS["torch", "clip"] = make_translator(
    [("a", ("input",)), ("a_min", ("min",)), ("a_max", ("max",))]
)
_CUSTOM_WRAPPERS["torch", "ones"] = torch_zeros_ones_wrap
_CUSTOM_WRAPPERS["torch", "zeros"] = torch_zeros_ones_wrap
_CUSTOM_WRAPPERS["torch", "eye"] = torch_eye_wrap
_CUSTOM_WRAPPERS["torch", "empty"] = make_translator([("shape", ("size",))])
_CUSTOM_WRAPPERS["torch", "take"] = make_translator(
    [("a", ("input",)), ("indices", ("index",)), ("axis", ("dim",))]
)

# for torch < 1.9
_MODULE_ALIASES['torch[alt]'] = 'torch'

_FUNCS["torch[alt]", "linalg.eigh"] = torch_linalg_eigh
_FUNCS["torch[alt]", "linalg.eigvalsh"] = torch_linalg_eigvalsh

_SUBMODULE_ALIASES["torch[alt]", "linalg.qr"] = "torch"
_SUBMODULE_ALIASES["torch[alt]", "linalg.svd"] = "torch"
_SUBMODULE_ALIASES["torch[alt]", "linalg.norm"] = "torch"
_SUBMODULE_ALIASES["torch[alt]", "linalg.solve"] = "torch"

_CUSTOM_WRAPPERS["torch[alt]", "linalg.svd"] = svd_UsV_to_UsVH_wrapper
_CUSTOM_WRAPPERS["torch[alt]", "linalg.qr"] = qr_allow_fat
_CUSTOM_WRAPPERS["torch[alt]", "linalg.solve"] = torch_linalg_solve_wrap

# --------------------------- register your own! ---------------------------- #


def register_function(backend, name, fn, wrap=False):
    """Directly provide your own function.

    Parameters
    ----------
    backend : str
        The name of the backend to register the function for.
    name : str
        Name of the function, e.g. `'sum'` or `'linalg.svd'`.
    fn : callable
        The function to register.
    wrap : bool, optional
        Whether to wrap the old function like ``fn(old_fn)`` rather than
        directly supply the entire new function.
    """
    if wrap:
        old = get_lib_fn(backend, name)
        _FUNCS[backend, name] = fn(old)
    else:
        _FUNCS[backend, name] = fn


class Composed:
    """Compose an ``autoray.do`` using function. See the main wrapper
    ``compose``.
    """

    def __init__(self, fn, name=None):
        self._default_fn = fn
        if name is None:
            name = fn.__name__
        self._name = name
        self._supply_backend = "backend" in signature(fn).parameters
        _COMPOSED_FUNCTION_GENERATORS[self._name] = self.make_function

    def register(self, backend, fn):
        register_function(backend, self._name, fn)

    def make_function(self, backend):
        if self._supply_backend:
            fn = functools.partial(self._default_fn, backend=backend)
        else:
            fn = self._default_fn
        self.register(backend, fn)
        return fn

    def __call__(self, *args, like=None, **kwargs):
        backend = choose_backend(self._name, *args, like=like, **kwargs)
        try:
            fn = get_lib_fn(backend, self._name)
        except ImportError:
            fn = self.make_function(backend)
        return fn(*args, **kwargs)

    def __repr__(self):
        return f"Composed('{self._name}')"


def compose(fn, *, name=None):
    """Take a function consisting of multiple ``autoray.do`` calls and compose
    it into a new, single, named function, registered with ``autoray.do``.

    This creates a default implementation of this function for each new backend
    encountered without explicitly having to write each out, but also allows
    for specific implementations to be overridden for specific backends.

    If the function takes a ``backend`` argument, it will be supplied with the
    backend name, to save having to re-choose the backend.
    """
    if fn is None:
        return functools.partial(compose, name=name)
    return functools.wraps(fn)(Composed(fn, name))
