"""Defines a new transformation, "finalisation", which replaces all custom primitives
with the results of their `impl` rule.

This is useful prior to e.g. ONNX export, to replace all custom primitives (e.g.
`equinox.internal.unvmap_any`) with the JAX operations that they wrap.

Most end users will find `finalise_fn` and `finalise_make_jaxpr` to be the most useful
operations provided by this module.

Library authors may wish to register their primitives with `primitive_finalisations`.

!!! warning

    Note that you should not perform any other JAX transformations -- e.g. `jax.vmap`,
    `jax.grad` etc. -- after finalisation. The result may be silently incorrect.
"""

import functools as ft
from typing import Any, Dict

import jax
import jax.tree_util as jtu


def _safe_map(f, *args):
    args = list(map(list, args))
    n = len(args[0])
    for arg in args[1:]:
        assert len(arg) == n, f"length mismatch: {list(map(len, args))}"
    return list(map(f, *args))


def _maybe_finalise_jaxpr(val: Any):
    is_open_jaxpr = False
    if isinstance(val, jax.core.Jaxpr):
        if len(val.constvars) == 0:
            is_open_jaxpr = True
            val = jax.core.ClosedJaxpr(val, [])
        else:
            return val
    if isinstance(val, jax.core.ClosedJaxpr):
        val = finalise_jaxpr(val)
    if is_open_jaxpr:
        val = val.jaxpr
    return val


def _finalise_jaxprs_in_params(params):
    new_params = {}
    for key, val in params.items():
        if type(val) is tuple:  # not isinstance to avoid breaking namedtuples
            val = tuple(_maybe_finalise_jaxpr(v) for v in val)
        else:
            val = _maybe_finalise_jaxpr(val)
        new_params[key] = val
    return new_params


def _default_finalisation(prim: jax.core.Primitive, *args, **kwargs):
    return prim.bind(*args, **kwargs)


def _impl_finalisation(prim: jax.core.Primitive, *args, **kwargs):
    return prim.impl(*args, **kwargs)


primitive_finalisations = {}


def register_impl_finalisation(prim: jax.core.Primitive):
    primitive_finalisations[prim] = ft.partial(_impl_finalisation, prim)


def finalise_eval_jaxpr(jaxpr: jax.core.Jaxpr, consts, *args):
    """As jax.core.eval_jaxpr, but finalises (typically by calling `impl` rather than
    `bind` for custom primitives).
    """

    def read(v: jax.core.Atom) -> Any:
        return v.val if isinstance(v, jax.core.Literal) else env[v]

    def write(v: jax.core.Var, val: Any) -> None:
        env[v] = val

    env: Dict[jax.core.Var, Any] = {}
    _safe_map(write, jaxpr.constvars, consts)
    _safe_map(write, jaxpr.invars, args)
    for eqn in jaxpr.eqns:
        params = _finalise_jaxprs_in_params(eqn.params)
        subfuns, bind_params = eqn.primitive.get_bind_params(params)
        try:
            call = primitive_finalisations[eqn.primitive]
        except KeyError:
            call = ft.partial(_default_finalisation, eqn.primitive)
        ans = call(*subfuns, *_safe_map(read, eqn.invars), **bind_params)
        if eqn.primitive.multiple_results:
            _safe_map(write, eqn.outvars, ans)
        else:
            write(eqn.outvars[0], ans)
    return _safe_map(read, jaxpr.outvars)


def finalise_jaxpr_as_fn(jaxpr: jax.core.ClosedJaxpr):
    """As `jax.core.jaxpr_as_fn`, but the result is finalised."""
    return ft.partial(finalise_eval_jaxpr, jaxpr.jaxpr, jaxpr.consts)


def finalise_jaxpr(jaxpr: jax.core.ClosedJaxpr):
    """A jaxpr-to-jaxpr transformation that performs finalisation."""
    fn = finalise_jaxpr_as_fn(jaxpr)
    args = [
        jax.ShapeDtypeStruct(x.aval.shape, x.aval.dtype) for x in jaxpr.jaxpr.invars
    ]
    return jax.make_jaxpr(fn)(*args)


def finalise_fn(fn):
    """Wraps a function to perform finalisation. (In a manner similar to e.g. `jax.vmap`
    wrapping a function to perform batching.)
    """

    def _finalise_fn(*args):
        jaxpr, struct = jax.make_jaxpr(fn, return_shape=True)(*args)
        flat_args = jtu.tree_leaves(args)
        out = finalise_eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, *flat_args)
        treedef = jtu.tree_structure(struct)
        return jtu.tree_unflatten(treedef, out)

    return _finalise_fn


def finalise_make_jaxpr(fn, *, return_shape=False):
    """As `jax.make_jaxpr`, but finalises the result."""

    def _finalise_make_jaxpr(*args):
        jaxpr = jax.make_jaxpr(fn, return_shape=return_shape)(*args)
        if return_shape:
            jaxpr, struct = jaxpr
            jaxpr = finalise_jaxpr(jaxpr)
            return jaxpr, struct
        else:
            jaxpr = finalise_jaxpr(jaxpr)
            return jaxpr

    return _finalise_make_jaxpr


# Register finalisation rules for Equinox's custom primitives.
from .debug import announce_jaxpr_p
from .errors import branched_error_p
from .noinline import noinline_p
from .nontraceable import nonbatchable_p, nondifferentiable_backward_p, nontraceable_p
from .unvmap import unvmap_all_p, unvmap_any_p, unvmap_max_p


for prim in (
    unvmap_all_p,
    unvmap_any_p,
    unvmap_max_p,
    nondifferentiable_backward_p,
    announce_jaxpr_p,
    noinline_p,
    nontraceable_p,
    nonbatchable_p,
):
    register_impl_finalisation(prim)


# Assume no errors are raised after finalisation. Can't use the impl rule as that
# has a pure_callback.
def _branched_error_if_finalisation(pred, index, *x, msgs):
    return x


primitive_finalisations[branched_error_p] = _branched_error_if_finalisation


# To make this also useful as debugging tool, we also inline some calls.


def _jvp_call_p_finalisation(fun, jvp, *args):
    return fun.call_wrapped(*args)


def _xla_call_p_finalisation(fun, *args, **params):
    del params
    return fun.call_wrapped(*args)


def _stop_gradient_finalisation(x):
    return x


primitive_finalisations[
    jax.custom_derivatives.custom_jvp_call_p
] = _jvp_call_p_finalisation
primitive_finalisations[jax.interpreters.xla.xla_call_p] = _xla_call_p_finalisation
primitive_finalisations[jax.lax.stop_gradient_p] = _stop_gradient_finalisation
