import logging

from jax import lax
from jax.extend.core import Primitive, ClosedJaxpr
from jax2onnx.converter.jaxpr_converter import Jaxpr2OnnxConverter
from jax2onnx.plugin_system import PrimitiveLeafPlugin, register_primitive

logger = logging.getLogger(__name__)


lax.remat2_p = Primitive(
    "lax.remat2"
)  # need to be "lax.remat2" there is a mapping to "remat2" in jaxpr_converter.py
lax.remat2_p.multiple_results = True


@register_primitive(
    jaxpr_primitive=lax.remat2_p.name,
    jax_doc="https://docs.jax.dev/en/latest/jep/11830-new-remat-checkpoint.html",
    onnx=[],
    since="v0.6.5",
    context="primitives.lax",
    component="remat2",
)
class RematPlugin(PrimitiveLeafPlugin):
    """
    Handles the `remat2` primitive, which is generated by `jax.checkpoint`.

    For ONNX conversion, which represents a static forward pass, the concept
    of re-materialization is not applicable. This handler effectively
    "inlines" the jaxpr contained within the remat2 primitive, treating its
    computation as part of the main graph.
    """

    """
    Handles the `remat2` primitive, which is generated by `jax.checkpoint`.
    """

    def to_onnx(self, s: Jaxpr2OnnxConverter, node_inputs, node_outputs, params):
        """
        Inlines the jaxpr from a `remat2` equation.
        """
        logger.debug("Inlining jaxpr from 'remat2' (jax.checkpoint).")

        # Get the sub-graph from the 'jaxpr' parameter
        sub_jaxpr_param = params["jaxpr"]

        # FIX: Check if the parameter is a ClosedJaxpr or a raw Jaxpr.
        if isinstance(sub_jaxpr_param, ClosedJaxpr):
            sub_jaxpr = sub_jaxpr_param.jaxpr
            consts = sub_jaxpr_param.consts
        else:
            # The traceback shows it's a raw Jaxpr object in this case.
            sub_jaxpr = sub_jaxpr_param
            consts = []

        # Map inputs from the main graph to the sub-graph's inputs.
        for sub_invar, main_invar in zip(sub_jaxpr.invars, node_inputs):
            onnx_name = s.get_name(main_invar)
            if onnx_name is None:
                raise ValueError(
                    f"Input to remat2 primitive '{main_invar}' has no ONNX name."
                )
            s.var_to_name[sub_invar] = onnx_name

        # Process the equations within the sub-graph.
        s._process_jaxpr(sub_jaxpr, consts)

        # Map outputs from the sub-graph back to the main graph.
        for main_outvar, sub_outvar in zip(node_outputs, sub_jaxpr.outvars):
            onnx_name = s.get_name(sub_outvar)
            if onnx_name is None:
                raise ValueError(
                    f"Output of remat's internal jaxpr '{sub_outvar}' was not assigned an ONNX name."
                )
            s.var_to_name[main_outvar] = onnx_name
