import functools as ft
from collections.abc import Callable
from typing import Any, Optional, Union

import jax
import jax._src.pretty_printer as pp
import jax._src.source_info_util as source_info_util
import jax.core
import jax.interpreters.ad as ad
import jax.interpreters.batching as batching
import jax.interpreters.mlir as mlir
import jax.numpy as jnp
import jax.tree_util as jtu
import numpy as np
from jaxtyping import Array, Int, PyTree

from .._ad import filter_jvp
from .._caches import internal_lru_caches
from .._compile_utils import hashable_combine, hashable_partition
from .._eval_shape import filter_eval_shape
from .._filters import combine, is_array, partition
from .._module import field, Module, module_update_wrapper
from .._vmap_pmap import filter_vmap
from . import _primitive
from ._primitive import (
    filter_primitive_batching,
    filter_primitive_bind,
    filter_primitive_def,
    filter_primitive_jvp,
    filter_primitive_transpose,
    materialise_zeros,
)


def _is_shapedarray(x):
    return isinstance(x, jax.core.ShapedArray)


def _to_shapedarray(x):
    if isinstance(x, jax.ShapeDtypeStruct):
        return jax.core.ShapedArray(x.shape, x.dtype)
    else:
        return x


def _only_shapedarrays(rule):
    def _rule_wrapper(*args, **params):
        for arg in args:
            if not _is_shapedarray(arg):
                raise NotImplementedError(
                    "noinline only supports ShapedArrays in abstract evaluation"
                )
        return rule(*args, **params)

    return _rule_wrapper


@ft.lru_cache(maxsize=128)
def _cache_filter_eval_shape(key):
    flat_dynamic, treedef_dynamic, static = key
    flat_dynamic = [jax.ShapeDtypeStruct(shape, dtype) for shape, dtype in flat_dynamic]
    dynamic = jtu.tree_unflatten(treedef_dynamic, flat_dynamic)
    (
        abstract_fn,
        transforms,
        args,
    ) = hashable_combine(dynamic, static)
    assert transforms[0] is _impl_transform
    abstract_fn = _abstract_transform(abstract_fn)
    for transform in transforms[1:]:
        abstract_fn = transform(abstract_fn)
    return filter_eval_shape(abstract_fn, args)


internal_lru_caches.append(_cache_filter_eval_shape)


def _is_undefined(x):
    return type(x) is ad.UndefinedPrimal


def _is_none(x):
    return x is None


def _is_not_mapped(x):
    return x is batching.not_mapped


def _move_to_front(input, batch_axis):
    if batch_axis is batching.not_mapped:
        return input
    else:
        return jnp.swapaxes(input, 0, batch_axis)


def _int_to_zero(batch_axis):
    if batch_axis is batching.not_mapped:
        return batch_axis
    else:
        return 0


@ft.lru_cache(maxsize=None)
def _get_callback(treedef, static, is_float0):
    @ft.partial(jax.jit, static_argnums=0)
    def callback(static_fn, dynamic):
        # pure_callback casts float0 to bool, so need to cast back
        assert len(dynamic) == len(is_float0)
        dynamic = [
            np.broadcast_to(np.zeros((), dtype=jax.dtypes.float0), x.shape) if y else x
            for x, y in zip(dynamic, is_float0)
        ]
        iter_dynamic = iter(dynamic)
        flat = [next(iter_dynamic) if x is None else x for x in static]
        assert next(iter_dynamic, None) is None
        transforms, args = jtu.tree_unflatten(treedef, flat)
        for transform in transforms:
            static_fn = transform(static_fn)
        out = static_fn(args)
        return tuple(jtu.tree_leaves(out))

    def callback_lookup(dynamic_index, *dynamic):
        static_fn = _index_to_fn[dynamic_index.item()]
        return callback(static_fn, dynamic)

    return callback_lookup


internal_lru_caches.append(_get_callback)


def _impl_transform(static_fn):
    def _impl_transform_impl(args):
        dynamic_fn, args, kwargs = args
        fn = hashable_combine(dynamic_fn, static_fn)
        return fn(*args, **kwargs)

    return _impl_transform_impl


def _abstract_transform(abstract_fn):
    def _abstract_transform_impl(args):
        dynamic_fn, args, kwargs = args
        return abstract_fn(dynamic_fn, *args, **kwargs)

    return _abstract_transform_impl


def _jvp_transform(static_fn):
    def _jvp_transform_impl(args):
        args, t_args = args
        tang_is_none = jtu.tree_map(_is_none, t_args, is_leaf=_is_none)
        no_tang, has_tang = partition(args, tang_is_none, is_leaf=_is_none)

        def _to_jvp(_has_tang):
            _args = combine(_has_tang, no_tang)
            return static_fn(_args)

        _, t_out = filter_jvp(_to_jvp, (has_tang,), (t_args,))
        return t_out

    return _jvp_transform_impl


class _MetaTransposeTransform(Module):
    undefined: PyTree[jax.core.ShapedArray]

    def __call__(self, static_fn):
        def _transpose_transform_impl(args):
            defined, cts_out = args

            def _to_transpose(_undefined):
                _args = combine(defined, _undefined)
                return static_fn(_args)

            (cts_undefined,) = jax.linear_transpose(_to_transpose, self.undefined)(
                cts_out
            )
            return cts_undefined

        return _transpose_transform_impl


class _MetaBatchTransform(Module):
    batch_axes: PyTree[Union[batching.NotMapped, int]]

    def __call__(self, static_fn):
        return filter_vmap(static_fn, in_axes=(self.batch_axes,))


@filter_primitive_def
def _noinline_impl(dynamic_index, abstract_fn, transforms, args):
    del abstract_fn
    static_fn = _index_to_fn[dynamic_index.item()]
    for transform in transforms:
        static_fn = transform(static_fn)
    return static_fn(args)


@_only_shapedarrays
@filter_primitive_def
def _noinline_abstract(dynamic_index, abstract_fn, transforms, args):
    del dynamic_index
    dynamic, static = hashable_partition(
        (abstract_fn, transforms, args), _is_shapedarray
    )
    flat_dynamic, treedef_dynamic = jtu.tree_flatten(dynamic)
    key = tuple((x.shape, x.dtype) for x in flat_dynamic), treedef_dynamic, static
    out_struct = _cache_filter_eval_shape(key)
    return jtu.tree_map(_to_shapedarray, out_struct)


@filter_primitive_jvp
def _noinline_jvp(primals, tangents):
    # TODO: add custom partial-eval rule to avoid the double-noinline?
    dynamic_index, abstract_fn, transforms, args = primals
    t_dynamic_index, t_abstract_fn, t_transforms, t_args = tangents
    assert (
        len(jtu.tree_leaves((t_dynamic_index, t_abstract_fn, t_transforms))) == 0
    )  # all none
    del t_dynamic_index, t_abstract_fn, t_transforms
    tangents = jtu.tree_map(materialise_zeros, args, t_args, is_leaf=_is_none)
    primal_outs = filter_primitive_bind(noinline_p, *primals)
    tangent_outs = filter_primitive_bind(
        noinline_p,
        dynamic_index,
        abstract_fn,
        transforms + [_jvp_transform],
        (args, t_args),
    )
    return primal_outs, tangent_outs


@filter_primitive_transpose(materialise_zeros=True)  # pyright: ignore
def _noinline_transpose(inputs, cts_out):
    dynamic_index, abstract_fn, transforms, args = inputs
    assert all(
        not _is_undefined(x)
        for x in jtu.tree_leaves(
            (dynamic_index, abstract_fn, transforms), is_leaf=_is_undefined
        )
    )
    # Note that `defined` may also include non-JAX-arrays
    undefined, defined = partition(args, _is_undefined, is_leaf=_is_undefined)
    undefined = jtu.tree_map(lambda x: x.aval, undefined, is_leaf=_is_undefined)
    cts_args = filter_primitive_bind(
        noinline_p,
        dynamic_index,
        abstract_fn,
        transforms + [_MetaTransposeTransform(undefined)],
        (defined, cts_out),
    )
    cts_rest = jtu.tree_map(lambda _: None, (dynamic_index, abstract_fn, transforms))
    return cts_rest + (cts_args,)


@filter_primitive_batching
def _noinline_batch(inputs, batch_axes):
    dynamic_index, abstract_fn, transforms, args = inputs
    dynamic_index_bdim, abstract_fn_bdim, transforms_bdim, args_bdim = batch_axes
    assert len(jtu.tree_leaves((abstract_fn_bdim, transforms_bdim))) == 0  # all none
    if dynamic_index_bdim is not batching.not_mapped:
        # The batch rule for `lax.cond` with vmap'd predicate simply
        # broadcasts all constants in the branches. In particular it may broadcast
        # this. We simply need to ignore this and return to having a single dynamic
        # index.
        # This is actually a silent error if you do something exceptionally silly, and
        # try to manually combine two different `noinline`d functions. There's no real
        # way to catch this that wouldn't slow things down at runtime though, I don't
        # think.
        assert jnp.ndim(dynamic_index) == 1
        dynamic_index = dynamic_index[0]
    del dynamic_index_bdim, abstract_fn_bdim, transforms_bdim
    args = jtu.tree_map(_move_to_front, args, args_bdim, is_leaf=_is_not_mapped)
    args_bdim = jtu.tree_map(_int_to_zero, args_bdim, is_leaf=_is_not_mapped)
    out = filter_primitive_bind(
        noinline_p,
        dynamic_index,
        abstract_fn,
        transforms + [_MetaBatchTransform(args_bdim)],
        args,
    )
    return out, jtu.tree_map(lambda _: 0, out)


def _pp_transform(x):
    if x is _jvp_transform:
        return "jvp"
    elif type(x) is _MetaTransposeTransform:
        return "transpose"
    elif type(x) is _MetaBatchTransform:
        return "vmap"
    else:
        assert False


def _noinline_pretty_print(eqn, context, settings):
    _, abstract_fn, transforms, _ = jtu.tree_unflatten(
        eqn.params["treedef"], eqn.params["static"]
    )
    pretty_params = dict(abstract_fn=abstract_fn)
    if type(eqn.invars[0]) is jax.core.Literal:
        static_fn_leaves, static_fn_treedef = _index_to_fn[eqn.invars[0].val]
        pretty_params["static_fn"] = jtu.tree_unflatten(
            static_fn_treedef, static_fn_leaves
        )
    if len(transforms) > 1:
        # skip impl
        transforms = [_pp_transform(x) for x in transforms[1:]]
        pretty_params["transforms"] = transforms
    lhs = jax.core.pp_vars(eqn.outvars, context, print_shapes=settings.print_shapes)
    rhs = [
        pp.text(eqn.primitive.name),
        jax.core.pp_kv_pairs(sorted(pretty_params.items()), context, settings),
        pp.text(" ") + jax.core.pp_vars(eqn.invars, context),
    ]
    annotation = (
        source_info_util.summarize(eqn.source_info) if settings.source_info else None
    )
    return pp.concat([lhs, pp.text(" = ", annotation=annotation), *rhs])


# Not a PyTree
class _MlirWrapper:
    def __init__(self, val):
        self.val = val


def _noinline_mlir(ctx, *dynamic, treedef, static, flatten, **kwargs):
    assert len(kwargs) == 0
    assert flatten.called()
    dynamic = [_MlirWrapper(x) for x in dynamic]
    abstract_dynamic = [_MlirWrapper(x) for x in ctx.avals_in]
    # This is really an internal implementation detail of another component that we're
    # messing with here.
    # Fortunately `noinline` and `primitive` are both inside Equinox, so this isn't too
    # bad.
    flat = _primitive._combine(dynamic, static)
    abstract_flat = _primitive._combine(abstract_dynamic, static)
    index, _, transforms, args = jtu.tree_unflatten(treedef, flat)
    abstract_index, _, abstract_transforms, abstract_args = jtu.tree_unflatten(
        treedef, abstract_flat
    )
    flat, treedef = jtu.tree_flatten((transforms, args))
    abstract_flat, abstract_treedef = jtu.tree_flatten(
        (abstract_transforms, abstract_args)
    )
    assert treedef == abstract_treedef
    dynamic = [x.val for x in flat if type(x) is _MlirWrapper]
    abstract_dynamic = [x.val for x in abstract_flat if type(x) is _MlirWrapper]
    static = tuple(None if type(x) is _MlirWrapper else x for x in flat)
    is_float0 = tuple(x.dtype == jax.dtypes.float0 for x in abstract_dynamic)
    callback = _get_callback(treedef, static, is_float0)

    vals_in = [index.val] + dynamic
    avals_in = [abstract_index.val] + abstract_dynamic
    result, _, keepalive = mlir.emit_python_callback(
        ctx,
        callback,
        None,
        vals_in,
        avals_in,
        ctx.avals_out,
        False,
        sharding=None,
    )
    ctx.module_context.add_keepalive(keepalive)
    return result


noinline_p = jax.core.Primitive("noinline")
noinline_p.multiple_results = True
noinline_p.def_impl(_noinline_impl)
noinline_p.def_abstract_eval(_noinline_abstract)
ad.primitive_jvps[noinline_p] = _noinline_jvp
ad.primitive_transposes[noinline_p] = _noinline_transpose
batching.primitive_batchers[noinline_p] = _noinline_batch
jax.core.pp_eqn_rules[noinline_p] = _noinline_pretty_print
mlir.register_lowering(noinline_p, _noinline_mlir)


_fn_to_index = {}
_index_to_fn = []


class _NoInlineWrapper(Module):
    dynamic_index: Int[Array, ""]
    abstract_fn: Callable = field(static=True)
    dynamic_fn: Any

    @property
    def __wrapped__(self):
        return self.abstract_fn

    def __call__(self, *args, **kwargs):
        return filter_primitive_bind(
            noinline_p,
            self.dynamic_index,
            self.abstract_fn,
            [_impl_transform],
            (self.dynamic_fn, args, kwargs),
        )


def noinline(
    fn: Callable, abstract_fn: Optional[Callable] = None  # pyright: ignore
) -> Callable:
    """Marks a function as not being inlined into a larger computation.
    This can help to reduce compile time at the expense of increased runtime.

    `fn` can be any callable `PyTree[Any] -> PyTree[Any]`. In addition `fn`
    itself may be any (callable) PyTree; any JAX arrays in `fn`'s PyTree are
    also treated as inputs to the noinline'd computation.

    `abstract_fn` determines the shapes/dtypes/pytrees of the output. (And must
     return results consistent with `fn`.) If `fn` is called as
    `fn(*args, **kwargs)`, then `abstract_fn` is called as
    ```python
    eqx.filter_eval_shape(
        abstract_fn, eqx.filter(fn, eqx.is_array), *args, **kwargs)
    )
    ```
    If not passed then `abstract_fn` is automatically constructed from `fn`.
    Specifying it is useful as noinline'd functions sharing the same
    `abstract_fn` may be substituted for each other without recompiling the
    enclosing computation graph; see the second example below.

    !!! Example

        ```python
        @noinline
        def f(x, y):
            print("Tracing!")
            return x + y

        @jax.jit
        def g(x, y):
            a = f(x, jnp.array(1))
            b = f(x, jnp.array(2))
            c = f(y, jnp.array(1))
            d = f(y, jnp.array(2))
            return a + b + c + d

        g(1, 1)
        ```

        In this example, `Tracing!` is printed twice. Once just for abstract
        evaluation (only) to figure out out its return shape in `f`'s
        computation graph. (Which is computationally cheap.) And then once
        when it is actually compiled as part of its own computation graph.
        (Which is computationally expensive.)

        Without the `noinline`, we would see `Tracing!` printed four times:
        once for each call, and with it being compiled each time. (Every time
        being computationally expensive.)

        Note how the `1` and the `2` are wrapped in arrays. If they were still
        Python scalars then they would be treated as static arguments, which
        will lead to recompilations of `f`. (E.g. what if there was an
        `if y == 1` command in there?)

    !!! Example

        ```python
        def abstract(_, x, y):
            return jnp.broadcast_arrays(x, y)[0]

        def f(x, y):
            print("Compiling f!")
            return x + y

        def g(x, y):
            print("Compiling g!")
            return x * y

        f = noinline(f, abstract)
        g = noinline(g, abstract)

        @jax.jit
        def call(fn, x, y):
            print("Compiling call!")
            return fn(x, y)

        call(f, 1, 1)  # Compiling call! Compiling f!
        call(g, 1, 1)  # Compiling g!
        ```

        In this example, we see how noinline'd functions with the same
        `abstract_fn` may be passed as inputs to the main call graph,
        and swapped without needing to recompile the main call graph.

    """
    dynamic_fn, static_fn = hashable_partition(fn, is_array)
    if abstract_fn is None:

        def abstract_fn(__dynamic_fn, *args, **kwargs):
            _fn = hashable_combine(__dynamic_fn, static_fn)
            return _fn(*args, **kwargs)

    try:
        dynamic_index = _fn_to_index[static_fn]
    except KeyError:
        dynamic_index = len(_index_to_fn)
        _fn_to_index[static_fn] = dynamic_index
        _index_to_fn.append(static_fn)
    dynamic_index = jnp.array(dynamic_index)
    noinline_fn = _NoInlineWrapper(dynamic_index, abstract_fn, dynamic_fn)
    return module_update_wrapper(noinline_fn, abstract_fn)
