import inspect
from numbers import Number
from textwrap import indent
from typing import Union

import numba
import numpy as np
from numba.cpython.unsafe.tuple import tuple_setitem

from aesara.graph.basic import Apply
from aesara.link.numba.dispatch import basic as numba_basic
from aesara.link.numba.dispatch.basic import (
    create_numba_signature,
    create_tuple_creator,
    numba_funcify,
)
from aesara.link.utils import (
    compile_function_src,
    get_name_for_object,
    unique_name_generator,
)
from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
from aesara.tensor.math import MaxAndArgmax
from aesara.tensor.nnet.basic import LogSoftmax, Softmax
from aesara.tensor.type import tensor


def create_vectorize_func(op, node, use_signature=False, identity=None, **kwargs):
    scalar_op_fn = numba_funcify(op.scalar_op, node=node, inline="always", **kwargs)

    if len(node.outputs) > 1:
        raise NotImplementedError(
            "Multi-output Elemwise Ops are not supported by the Numba backend"
        )

    if use_signature:
        signature = [create_numba_signature(node, force_scalar=True)]
    else:
        signature = []

    numba_vectorize = numba.vectorize(signature, identity=identity)
    elemwise_fn = numba_vectorize(scalar_op_fn)
    elemwise_fn.py_scalar_func = scalar_op_fn

    return elemwise_fn


@numba_funcify.register(Elemwise)
def numba_funcify_Elemwise(op, node, **kwargs):

    elemwise_fn = create_vectorize_func(op, node, use_signature=False)
    elemwise_fn_name = elemwise_fn.__name__

    if op.inplace_pattern:
        input_idx = op.inplace_pattern[0]
        sign_obj = inspect.signature(elemwise_fn.py_scalar_func)
        input_names = list(sign_obj.parameters.keys())

        unique_names = unique_name_generator([elemwise_fn_name, "np"], suffix_sep="_")
        input_names = [unique_names(i, force_unique=True) for i in input_names]

        updated_input_name = input_names[input_idx]

        inplace_global_env = {elemwise_fn_name: elemwise_fn, "np": np}

        inplace_elemwise_fn_name = f"{elemwise_fn_name}_inplace"

        input_signature_str = ", ".join(input_names)

        if node.inputs[input_idx].ndim > 0:
            inplace_elemwise_src = f"""
def {inplace_elemwise_fn_name}({input_signature_str}):
    return {elemwise_fn_name}({input_signature_str + ", " + updated_input_name})
            """
        else:
            # We can't perform in-place updates on Numba scalars, so we need to
            # convert them to NumPy scalars.
            # TODO: We should really prevent the rewrites from creating
            # in-place updates on scalars when the Numba mode is selected (or
            # in general?).
            inplace_elemwise_src = f"""
def {inplace_elemwise_fn_name}({input_signature_str}):
    {updated_input_name}_scalar = np.asarray({updated_input_name})
    return {elemwise_fn_name}({input_signature_str + ", " + updated_input_name}_scalar).item()
            """

        inplace_elemwise_fn = compile_function_src(
            inplace_elemwise_src, inplace_elemwise_fn_name, inplace_global_env
        )
        return numba.njit(inline="always")(inplace_elemwise_fn)

    return elemwise_fn


def create_axis_reducer(
    reduce_fn: numba.np.ufunc.dufunc.DUFunc,
    identity: Union[np.ndarray, Number],
    axis: int,
    ndim: int,
    dtype: numba.types.Type,
    keepdims: bool = False,
) -> numba.core.dispatcher.Dispatcher:
    r"""Create a Numba JITed function that performs a NumPy reduction on a given axis.

    The functions generated by this function take the following form:

    .. code-block:: python

        def careduce_axis(x):
            res_shape = tuple(shape[i] if i < axis else shape[i + 1] for i in range(ndim - 1))
            res = np.full(res_shape, identity, dtype=dtype)

            x_axis_first = x.transpose(reaxis_first)

            for m in range(x.shape[axis]):
                reduce_fn(res, x_axis_first[m], res)

            if keepdims:
                return np.expand_dims(res, axis)
            else:
                return res


    This can be removed/replaced when
    https://github.com/numba/numba/issues/4504 is implemented.

    Parameters
    ==========
    reduce_fn:
        The Numba ``ufunc`` representing a binary op that can perform the
        reduction on arbitrary ``ndarray``\s.
    identity:
        The identity value for the reduction.
    axis:
        The axis to reduce.
    ndim:
        The number of dimensions of the result.
    dtype:
        The data type of the result.
    keepdims:
        Determines whether or not the reduced dimension is retained.
    """
    if ndim > 1:

        if keepdims:

            @numba.njit(inline="always")
            def set_out_dims(x):
                return np.expand_dims(x, axis)

        else:

            @numba.njit(inline="always")
            def set_out_dims(x):
                return x

        res_shape_tuple_ctor = create_tuple_creator(
            lambda i, shape: shape[i] if i < axis else shape[i + 1], ndim - 1
        )

        reaxis_first = (axis,) + tuple(i for i in range(ndim) if i != axis)

        @numba.njit(boundscheck=False)
        def careduce_axis(x):
            res_shape = res_shape_tuple_ctor(x.shape)
            x_axis_first = x.transpose(reaxis_first)

            res = np.full(res_shape, numba_basic.to_scalar(identity), dtype=dtype)
            for m in range(x.shape[axis]):
                reduce_fn(res, x_axis_first[m], res)

            return set_out_dims(res)

    else:

        if keepdims:

            @numba.njit(inline="always")
            def set_out_dims(x):
                return np.array([x], dtype)

        else:

            @numba.njit(inline="always")
            def set_out_dims(x):
                return numba_basic.direct_cast(x, dtype)

        @numba.njit(boundscheck=False)
        def careduce_axis(x):
            res = numba_basic.to_scalar(identity)
            for val in x:
                res = reduce_fn(res, val)
            return set_out_dims(res)

    return careduce_axis


def create_multiaxis_reducer(
    reduce_fn, identity, axes, ndim, dtype, input_name="input"
):
    r"""Construct a function that reduces multiple axes.

    The functions generated by this function take the following form:

    .. code-block:: python

        def careduce_maximum(input):
            axis_0_res = careduce_axes_fn_0(input)
            axis_1_res = careduce_axes_fn_1(axis_0_res)
            ...
            axis_N_res = careduce_axes_fn_N(axis_N_minus_1_res)
            return axis_N_res

    The range 0-N is determined by the `axes` argument (i.e. the
    axes to be reduced).


    Parameters
    ==========
    reduce_fn:
        The Numba ``ufunc`` representing a binary op that can perform the
        reduction on arbitrary ``ndarray``\s.
    identity:
        The identity value for the reduction.
    axes:
        The axes to reduce.
    ndim:
        The number of dimensions of the result.
    dtype:
        The data type of the result.

    """
    if len(axes) == 1:
        return create_axis_reducer(reduce_fn, identity, axes[0], ndim, dtype)

    careduce_fn_name = f"careduce_{get_name_for_object(reduce_fn)}"
    global_env = {}
    to_reduce = reversed(sorted(axes))
    careduce_lines_src = []
    var_name = input_name

    for i, axis in enumerate(to_reduce):
        careducer_axes_fn_name = f"careduce_axes_fn_{i}"
        global_env[careducer_axes_fn_name] = create_axis_reducer(
            reduce_fn, identity, axis - i, ndim, dtype
        )
        ndim -= 1
        last_var_name = var_name
        var_name = f"axis_{i}_res"
        careduce_lines_src.append(
            f"{var_name} = {careducer_axes_fn_name}({last_var_name})"
        )

    careduce_assign_lines = indent("\n".join(careduce_lines_src), " " * 4)
    careduce_def_src = f"""
def {careduce_fn_name}({input_name}):
{careduce_assign_lines}
    return {var_name}
    """

    careduce_fn = compile_function_src(careduce_def_src, careduce_fn_name, global_env)
    return numba.njit(careduce_fn)


def create_axis_apply_fn(fn, axis, ndim, dtype):
    reaxis_first = tuple(i for i in range(ndim) if i != axis) + (axis,)

    @numba.njit(boundscheck=False)
    def axis_apply_fn(x):
        x_reaxis = x.transpose(reaxis_first)

        res = np.zeros(x_reaxis.shape[:-1], dtype=dtype)
        for m in np.ndindex(res.shape):
            v = fn(x_reaxis[m])
            res[m] = v
        return res

    return axis_apply_fn


@numba_funcify.register(CAReduce)
def numba_funcify_CAReduce(op, node, **kwargs):
    axes = op.axis
    if axes is None:
        axes = list(range(node.inputs[0].ndim))

    if hasattr(op, "acc_dtype") and op.acc_dtype is not None:
        acc_dtype = op.acc_dtype
    else:
        acc_dtype = node.outputs[0].type.dtype

    np_acc_dtype = np.dtype(acc_dtype)

    scalar_op_identity = np.asarray(op.scalar_op.identity, dtype=np_acc_dtype)

    scalar_nfunc_spec = op.scalar_op.nfunc_spec

    # We construct a dummy `Apply` that has the minimum required number of
    # inputs for the scalar `Op`.  Without this, we would get a scalar function
    # with too few arguments.
    dummy_node = Apply(
        op,
        [tensor(np_acc_dtype, [False]) for i in range(scalar_nfunc_spec[1])],
        [tensor(np_acc_dtype, [False]) for o in range(scalar_nfunc_spec[2])],
    )

    # TODO: Use `scalar_op_identity`?
    elemwise_fn = create_vectorize_func(op, dummy_node, use_signature=True, **kwargs)

    input_name = get_name_for_object(node.inputs[0])
    ndim = node.inputs[0].ndim
    careduce_fn = create_multiaxis_reducer(
        elemwise_fn, scalar_op_identity, axes, ndim, np_acc_dtype, input_name=input_name
    )

    return careduce_fn


@numba_funcify.register(DimShuffle)
def numba_funcify_DimShuffle(op, **kwargs):
    shuffle = tuple(op.shuffle)
    drop = tuple(op.drop)
    augment = tuple(op.augment)
    inplace = op.inplace

    ndim_new_shape = len(shuffle) + len(augment)
    create_zeros_tuple = numba_basic.create_tuple_creator(lambda _: 0, ndim_new_shape)

    if len(shuffle) > 0:

        @numba.njit
        def populate_new_shape(i, j, new_shape, shuffle_shape):
            if i in augment:
                new_shape = tuple_setitem(new_shape, i, 1)
                return j, new_shape
            else:
                new_shape = tuple_setitem(new_shape, i, shuffle_shape[j])
                return j + 1, new_shape

    else:
        # When `len(shuffle) == 0`, the `shuffle_shape[j]` expression above is
        # is typed as `getitem(Tuple(), int)`, which has no implementation
        # (since getting an item from an empty sequence doesn't make sense).
        # To avoid this compile-time error, we omit the expression altogether.
        @numba.njit(inline="always")
        def populate_new_shape(i, j, new_shape, shuffle_shape):
            return j, tuple_setitem(new_shape, i, 1)

    @numba.njit
    def dimshuffle_inner(x, shuffle):
        res = np.transpose(x, shuffle + drop)
        shuffle_shape = res.shape[: len(shuffle)]

        new_shape = create_zeros_tuple()

        j = 0
        for i in range(len(new_shape)):
            j, new_shape = populate_new_shape(i, j, new_shape, shuffle_shape)

        # FIXME: Numba's `array.reshape` only accepts C arrays.
        res_reshape = np.reshape(np.ascontiguousarray(res), new_shape)

        if not inplace:
            return res_reshape.copy()
        else:
            return res_reshape

    # Without the following wrapper function we would see this error:
    # E   No implementation of function Function(<built-in function getitem>) found for signature:
    # E
    # E    >>> getitem(UniTuple(int64 x 2), slice<a:b>)
    # E
    # E   There are 22 candidate implementations:
    # E      - Of which 22 did not match due to:
    # E      Overload of function 'getitem': File: <numerous>: Line N/A.
    # E        With argument(s): '(UniTuple(int64 x 2), slice<a:b>)':
    # E       No match.
    # ...(on this line)...
    # E           shuffle_shape = res.shape[: len(shuffle)]
    @numba.njit(inline="always")
    def dimshuffle(x):
        return dimshuffle_inner(np.asarray(x), shuffle)

    return dimshuffle


@numba_funcify.register(Softmax)
def numba_funcify_Softmax(op, node, **kwargs):

    x_at = node.inputs[0]
    x_dtype = x_at.type.numpy_dtype
    x_dtype = numba.np.numpy_support.from_dtype(x_dtype)

    # np.max(x, axis=1)
    reduce_max = create_axis_reducer(np.maximum, -np.inf, 1, x_at.ndim, x_dtype)
    # np.sum(x, axis=1)
    reduce_sum = create_axis_reducer(np.add, 0.0, 1, x_at.ndim, x_dtype)

    @numba.njit
    def softmax(x):
        z = np.expand_dims(reduce_max(x), -1)
        e_x = np.exp(x - z)
        w = np.expand_dims(reduce_sum(e_x), -1)
        sm = e_x / w
        return sm

    return softmax


@numba_funcify.register(LogSoftmax)
def numba_funcify_LogSoftmax(op, node, **kwargs):

    x_at = node.inputs[0]
    x_dtype = x_at.type.numpy_dtype
    x_dtype = numba.np.numpy_support.from_dtype(x_dtype)

    # np.max(x, axis=1)
    reduce_max = create_axis_reducer(np.maximum, -np.inf, 1, x_at.ndim, x_dtype)
    # np.sum(x, axis=1, keepdims=True)
    reduce_sum = create_axis_reducer(np.add, 0.0, 1, x_at.ndim, x_dtype, keepdims=True)

    @numba.njit
    def log_softmax(x):
        xdev = x - np.expand_dims(reduce_max(x), -1)
        lsm = xdev - np.log(reduce_sum(np.exp(xdev)))
        return lsm

    return log_softmax


@numba_funcify.register(MaxAndArgmax)
def numba_funcify_MaxAndArgmax(op, node, **kwargs):
    axis = op.axis
    x_at = node.inputs[0]
    x_dtype = x_at.type.numpy_dtype
    x_dtype = numba.np.numpy_support.from_dtype(x_dtype)
    x_ndim = x_at.ndim

    if x_ndim == 0:

        @numba.njit(inline="always")
        def maxandargmax(x):
            return x, 0

    else:

        axes = tuple(int(ax) for ax in axis)

        # NumPy does not support multiple axes for argmax; this is a
        # work-around
        keep_axes = tuple(i for i in range(x_ndim) if i not in axes)

        reduce_max = create_multiaxis_reducer(
            np.maximum, -np.inf, axes, x_ndim, x_dtype
        )
        reduced_x_ndim = x_ndim - len(axes) + 1
        argmax_axis = create_axis_apply_fn(
            np.argmax, reduced_x_ndim - 1, reduced_x_ndim, np.int64
        )

        reaxis_order = keep_axes + axes
        sl1 = slice(None, len(keep_axes))
        sl2 = slice(len(keep_axes), None)

        @numba.njit
        def maxandargmax(x):
            max_res = reduce_max(x)

            # Not-reduced axes in front
            transposed_x = np.ascontiguousarray(np.transpose(x, reaxis_order))
            kept_shape = transposed_x.shape[sl1]
            reduced_shape = transposed_x.shape[sl2]
            reduced_size = 1
            for s in reduced_shape:
                reduced_size *= s

            # Numpy.prod returns 1.0 when arg is empty, so we cast it to int64
            # Otherwise reshape would complain citing float arg
            new_shape = kept_shape + (reduced_size,)
            reshaped_x = transposed_x.reshape(new_shape)

            max_idx_res = argmax_axis(reshaped_x)

            return max_res, max_idx_res

    return maxandargmax
