# SPDX-License-Identifier: Apache-2.0
"""Control flow operators.

This module implements ONNX control flow operators like Loop and If.
"""

from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple

import onnx
import torch
import torch.nn as nn
import torch.fx
from onnx import numpy_helper

from ..utils.names import sanitize_name
from ..op_registry import register
from ..utils.attributes import get_attribute
from ..utils.op_helpers import get_optional_input

if TYPE_CHECKING:
    from ..graph_builder import GraphBuilder


def _collect_all_subgraph_inputs(node: onnx.NodeProto) -> set:
    """Recursively collect all inputs from a node including nested subgraph inputs.

    For control flow nodes like If and Loop, this also collects inputs that
    are referenced by nested subgraphs.

    Parameters
    ----------
    node : onnx.NodeProto
        The ONNX node to collect inputs from.

    Returns
    -------
    set
        Set of all input names referenced by this node and its nested subgraphs.
    """
    inputs = set(node.input)

    # Collect inputs from subgraphs (for If, Loop, etc.)
    for attr in node.attribute:
        if attr.type == onnx.AttributeProto.GRAPH:
            subgraph = attr.g
            # Collect subgraph's own initializers and inputs as local values
            local_values = set()
            for init in subgraph.initializer:
                local_values.add(init.name)
            for inp in subgraph.input:
                local_values.add(inp.name)

            # Recursively collect inputs from subgraph nodes
            for sub_node in subgraph.node:
                sub_inputs = _collect_all_subgraph_inputs(sub_node)
                # Add outputs of this subgraph node to local values
                for out in sub_node.output:
                    if out:
                        local_values.add(out)
                # Inputs not satisfied locally are outer references
                for sub_inp in sub_inputs:
                    if sub_inp and sub_inp not in local_values:
                        inputs.add(sub_inp)

    return inputs


def _build_subgraph_module(
    body_graph: onnx.GraphProto,
    parent_env: Dict[str, torch.fx.Node],
    parent_opset_versions: Dict[str, int],
    parent_type_info: Optional[Dict[str, bool]] = None,
    tensor_loader: Optional[Callable[[onnx.TensorProto], torch.Tensor]] = None,
) -> Tuple[torch.fx.GraphModule, List[str], List[str], List[str]]:
    """Build an FX GraphModule from an ONNX subgraph.

    Parameters
    ----------
    body_graph : onnx.GraphProto
        The ONNX subgraph to convert.
    parent_env : Dict[str, torch.fx.Node]
        Environment from parent graph (for accessing outer scope values).
    parent_opset_versions : Dict[str, int]
        Opset versions from parent model.
    parent_type_info : Optional[Dict[str, bool]]
        Mapping from value names to whether they are optional types in parent scope.

    Returns
    -------
    Tuple[torch.fx.GraphModule, List[str], List[str], List[str]]
        The FX GraphModule, list of input names, list of output names, and outer refs.
    """
    from ..op_registry import get_handler

    graph = torch.fx.Graph()
    env: Dict[str, torch.fx.Node] = {}
    constants: Dict[str, torch.Tensor] = {}

    # Initialize parent type info
    if parent_type_info is None:
        parent_type_info = {}

    # Get input and output names
    input_names = [inp.name for inp in body_graph.input]
    output_names = [out.name for out in body_graph.output]

    # Load initializers from subgraph
    initializer_map: Dict[str, torch.Tensor] = {}
    for initializer in body_graph.initializer:
        if tensor_loader is not None:
            initializer_map[initializer.name] = tensor_loader(initializer)
        else:
            np_array = numpy_helper.to_array(initializer)
            initializer_map[initializer.name] = torch.from_numpy(np_array.copy())

    # Register initializers as constants
    for name, tensor in initializer_map.items():
        safe_name = sanitize_name(name)
        constants[safe_name] = tensor
        fx_node = graph.get_attr(safe_name)
        env[name] = fx_node

    # Create placeholders for subgraph inputs
    for inp in body_graph.input:
        if inp.name in env:
            continue  # Skip if already loaded as initializer
        safe_name = sanitize_name(inp.name)
        placeholder = graph.placeholder(safe_name)
        env[inp.name] = placeholder

    # Collect all values that will be produced by nodes in this subgraph
    subgraph_outputs = set()
    for node in body_graph.node:
        for out in node.output:
            if out:
                subgraph_outputs.add(out)

    # Add references to parent scope values that are used in the subgraph
    # (including nested subgraphs). These will be passed as additional inputs.
    outer_refs: List[str] = []
    for node in body_graph.node:
        # Collect all inputs including from nested subgraphs
        all_inputs = _collect_all_subgraph_inputs(node)
        for inp_name in all_inputs:
            # Skip if empty, already in env, or will be produced by a node in this subgraph
            if not inp_name or inp_name in env or inp_name in subgraph_outputs:
                continue
            if inp_name in parent_env:
                if inp_name not in outer_refs:  # Avoid duplicates
                    outer_refs.append(inp_name)
                    safe_name = sanitize_name(inp_name)
                    placeholder = graph.placeholder(f"outer_{safe_name}")
                    env[inp_name] = placeholder

    # Create a minimal builder-like object for handler calls
    class SubgraphBuilder:
        def __init__(self):
            self.graph = graph
            self.env = env
            self._opset_versions = parent_opset_versions
            self._constants = constants
            self._submodules: Dict[str, nn.Module] = {}
            self.initializer_map = initializer_map
            self._body_graph = body_graph
            self._parent_type_info = parent_type_info
            self._tensor_loader = tensor_loader
            # Build type info for this subgraph (to pass to nested subgraphs)
            self._type_info = self._build_type_info()

        def load_tensor(self, tensor: onnx.TensorProto) -> torch.Tensor:
            if self._tensor_loader is not None:
                return self._tensor_loader(tensor)
            np_array = numpy_helper.to_array(tensor)
            return torch.from_numpy(np_array.copy())

        def _build_type_info(self) -> Dict[str, bool]:
            """Build a mapping of value names to whether they are optional types."""
            info: Dict[str, bool] = {}
            # Include parent type info
            info.update(self._parent_type_info)
            # Add types from this subgraph
            for value_info in self._body_graph.input:
                info[value_info.name] = value_info.type.HasField("optional_type")
            for value_info in self._body_graph.value_info:
                info[value_info.name] = value_info.type.HasField("optional_type")
            for value_info in self._body_graph.output:
                info[value_info.name] = value_info.type.HasField("optional_type")
            return info

        @property
        def opset_version(self) -> int:
            return self._opset_versions.get("", 1)

        def is_optional_type(self, name: str) -> bool:
            """Check if a value has optional type in the subgraph or parent scope."""
            # First check the combined type info (includes parent scope)
            if name in self._type_info:
                return self._type_info[name]
            return False

        def get_opset_version(self, domain: str = "") -> int:
            return self._opset_versions.get(domain, 1)

        def get_value(self, name: str) -> torch.fx.Node:
            if name not in self.env:
                raise KeyError(f"Value '{name}' not found in subgraph environment")
            return self.env[name]

        def has_value(self, name: str) -> bool:
            return name in self.env

        def call_function(
            self,
            func,
            args: tuple = (),
            kwargs: Optional[Dict[str, Any]] = None,
        ) -> torch.fx.Node:
            return self.graph.call_function(func, args=tuple(args), kwargs=kwargs or {})

        def register_submodule(self, name: str, module: nn.Module) -> str:
            safe_name = sanitize_name(name)
            if safe_name in self._submodules:
                counter = 0
                while f"{safe_name}_{counter}" in self._submodules:
                    counter += 1
                safe_name = f"{safe_name}_{counter}"
            self._submodules[safe_name] = module
            return safe_name

        def call_module(
            self,
            module_name: str,
            args: tuple = (),
            kwargs: Optional[Dict[str, Any]] = None,
        ) -> torch.fx.Node:
            return self.graph.call_module(
                module_name, args=tuple(args), kwargs=kwargs or {}
            )

    builder = SubgraphBuilder()

    # Convert nodes
    for node in body_graph.node:
        domain = node.domain if node.domain else ""
        opset = builder.get_opset_version(domain)
        handler = get_handler(node.op_type, domain, opset)
        if handler is None:
            raise ValueError(
                f"Unsupported operator in subgraph: {node.op_type} (domain={domain})"
            )
        fx_node = handler(builder, node)

        # Handle outputs
        if len(node.output) == 1:
            env[node.output[0]] = fx_node
        else:
            for i, output_name in enumerate(node.output):
                if output_name:
                    getitem_node = graph.call_function(
                        lambda x, idx=i: x[idx] if isinstance(x, (tuple, list)) else x,
                        args=(fx_node, i),
                    )
                    env[output_name] = getitem_node

    # Create output - return tuple of output values
    output_nodes = []
    for out_name in output_names:
        if out_name in env:
            output_nodes.append(env[out_name])
        else:
            raise KeyError(f"Output '{out_name}' not found in subgraph environment")

    if len(output_nodes) == 1:
        graph.output(output_nodes[0])
    else:
        graph.output(tuple(output_nodes))

    # Create the module
    root_module = nn.Module()
    for name, tensor in constants.items():
        root_module.register_buffer(name, tensor)
    for name, submod in builder._submodules.items():
        root_module.add_module(name, submod)

    module = torch.fx.GraphModule(root_module, graph)
    module.graph.lint()

    return module, input_names, output_names, outer_refs


class LoopModule(nn.Module):
    """Module that executes an ONNX Loop."""

    def __init__(
        self,
        body_module: torch.fx.GraphModule,
        n_loop_carried: int,
        n_scan_outputs: int,
        n_loop_vars: int,
        n_outer_vars: int,
    ):
        super().__init__()
        self.body = body_module
        self.n_loop_carried = n_loop_carried
        self.n_scan_outputs = n_scan_outputs
        self.n_loop_vars = n_loop_vars
        self.n_outer_vars = n_outer_vars

    def forward(self, max_iters, init_cond, *args) -> Tuple[torch.Tensor, ...]:
        """Execute the loop.

        Args are: loop_vars..., outer_vals...
        Returns final loop-carried values followed by concatenated scan outputs.
        """
        # Split args into loop_vars and outer_vals
        loop_vars = list(args[: self.n_loop_vars])
        outer_vals = list(args[self.n_loop_vars :])

        # Determine max iterations
        if max_iters is not None:
            max_i = (
                int(max_iters.item()) if hasattr(max_iters, "item") else int(max_iters)
            )
        else:
            max_i = 2**63 - 1  # Very large number

        # Initial condition
        if init_cond is not None:
            cond = (
                bool(init_cond.item())
                if hasattr(init_cond, "item")
                else bool(init_cond)
            )
        else:
            cond = True

        # Current loop-carried values
        current_vars = list(loop_vars)

        # Scan output accumulators
        scan_outputs: List[List[torch.Tensor]] = [
            [] for _ in range(self.n_scan_outputs)
        ]

        i = 0
        while i < max_i and cond:
            # Prepare inputs for body: iteration_num, condition, loop_carried..., outer...
            iter_tensor = torch.tensor(i, dtype=torch.int64)
            cond_tensor = torch.tensor(cond, dtype=torch.bool)

            # Call body function
            body_inputs = [iter_tensor, cond_tensor] + current_vars + outer_vals
            outputs = self.body(*body_inputs)

            # Handle single vs multiple outputs
            if not isinstance(outputs, tuple):
                outputs = (outputs,)

            # First output is new condition
            new_cond = outputs[0]
            cond = (
                bool(new_cond.item()) if hasattr(new_cond, "item") else bool(new_cond)
            )

            # Next n_loop_carried outputs are updated loop-carried values
            current_vars = list(outputs[1 : 1 + self.n_loop_carried])

            # Remaining outputs are scan outputs for this iteration
            for j in range(self.n_scan_outputs):
                scan_outputs[j].append(outputs[1 + self.n_loop_carried + j])

            i += 1

        # Prepare final outputs: loop-carried values, then stacked scan outputs
        final_outputs = list(current_vars)
        for scan_list in scan_outputs:
            if scan_list:
                # Stack scan outputs along a new first dimension
                final_outputs.append(torch.stack(scan_list, dim=0))
            else:
                # Empty scan output - create empty tensor
                final_outputs.append(torch.tensor([]))

        return tuple(final_outputs)


class IfModule(nn.Module):
    """Module that executes an ONNX If conditional."""

    def __init__(
        self,
        then_module: torch.fx.GraphModule,
        else_module: torch.fx.GraphModule,
        n_then_outer: int,
        n_else_outer: int,
    ):
        super().__init__()
        self.then_branch = then_module
        self.else_branch = else_module
        self.n_then_outer = n_then_outer
        self.n_else_outer = n_else_outer

    def forward(self, condition, *args) -> Any:
        """Execute the conditional.

        Args are: then_outer..., else_outer...
        Returns the outputs of the selected branch.
        """
        then_outer = list(args[: self.n_then_outer])
        else_outer = list(args[self.n_then_outer :])

        cond_val = (
            bool(condition.item()) if hasattr(condition, "item") else bool(condition)
        )

        if cond_val:
            result = self.then_branch(*then_outer)
        else:
            result = self.else_branch(*else_outer)

        return result


@register("Loop")
def loop_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """ONNX Loop operator.

    Loop has inputs: M (max trip count), cond (initial condition), v_initial... (loop-carried deps)
    Loop has attribute: body (GraphProto)
    Body inputs: iteration_num, condition, loop_carried_deps...
    Body outputs: condition, loop_carried_deps..., scan_outputs...
    """
    # Get body subgraph
    body_graph = get_attribute(node, "body")
    if body_graph is None:
        raise ValueError("Loop operator requires 'body' attribute")

    # Get inputs
    max_trip_count = builder.get_value(node.input[0]) if node.input[0] else None
    initial_cond = get_optional_input(builder, node, 1)
    loop_carried_inputs = [
        builder.get_value(node.input[i]) for i in range(2, len(node.input))
    ]

    # Get parent type info if available (for nested subgraphs)
    parent_type_info = getattr(builder, "_type_info", None)

    # Build subgraph module
    body_module, body_input_names, body_output_names, outer_refs = (
        _build_subgraph_module(
            body_graph,
            builder.env,
            builder._opset_versions,
            parent_type_info,
            tensor_loader=builder.load_tensor,
        )
    )

    # Get outer scope values that the body references
    outer_values = [builder.get_value(name) for name in outer_refs]

    # Number of loop-carried dependencies (excluding iteration_num and condition in body inputs)
    n_loop_carried = len(body_input_names) - 2
    # Number of scan outputs
    n_scan_outputs = len(body_output_names) - 1 - n_loop_carried

    # Create the loop module
    loop_module = LoopModule(
        body_module,
        n_loop_carried,
        n_scan_outputs,
        len(loop_carried_inputs),
        len(outer_values),
    )

    # Register the loop module
    module_name = builder.register_submodule(f"loop_{node.name or 'op'}", loop_module)

    # Build flat args for call_module: max_iters, init_cond, loop_vars..., outer_vals...
    args = [max_trip_count, initial_cond] + loop_carried_inputs + outer_values

    result = builder.call_module(module_name, args=tuple(args))

    return result


class ScanModule(nn.Module):
    """Module that executes an ONNX Scan operation."""

    def __init__(
        self,
        body_module: torch.fx.GraphModule,
        n_state_vars: int,
        n_scan_inputs: int,
        n_scan_outputs: int,
        n_outer_vars: int,
        scan_input_axes: List[int],
        scan_output_axes: List[int],
        scan_input_directions: List[int],
        scan_output_directions: List[int],
    ):
        super().__init__()
        self.body = body_module
        self.n_state_vars = n_state_vars
        self.n_scan_inputs = n_scan_inputs
        self.n_scan_outputs = n_scan_outputs
        self.n_outer_vars = n_outer_vars
        self.scan_input_axes = scan_input_axes
        self.scan_output_axes = scan_output_axes
        self.scan_input_directions = scan_input_directions
        self.scan_output_directions = scan_output_directions

    def forward(self, *args) -> Tuple[torch.Tensor, ...]:
        """Execute the scan.

        Args are: state_vars..., scan_inputs..., outer_vals...
        Returns final state variables followed by scan outputs.
        """
        # Split args
        state_vars = list(args[: self.n_state_vars])
        scan_inputs = list(
            args[self.n_state_vars : self.n_state_vars + self.n_scan_inputs]
        )
        outer_vals = list(args[self.n_state_vars + self.n_scan_inputs :])

        # Determine sequence length from first scan input
        if self.n_scan_inputs > 0:
            first_input = scan_inputs[0]
            axis = self.scan_input_axes[0] if self.scan_input_axes else 0
            sequence_length = first_input.shape[axis]
        else:
            sequence_length = 0

        # Initialize scan output accumulators
        scan_outputs: List[List[torch.Tensor]] = [
            [] for _ in range(self.n_scan_outputs)
        ]

        # Current state
        current_state = list(state_vars)

        # Execute loop
        for t in range(sequence_length):
            # Extract scan input elements for this iteration
            scan_input_elts = []
            for i, scan_input in enumerate(scan_inputs):
                axis = self.scan_input_axes[i] if i < len(self.scan_input_axes) else 0
                direction = (
                    self.scan_input_directions[i]
                    if i < len(self.scan_input_directions)
                    else 0
                )
                # Reverse direction: 0 = forward, 1 = reverse
                idx = sequence_length - 1 - t if direction == 1 else t
                # Select along the axis
                elt = torch.select(scan_input, axis, idx)
                scan_input_elts.append(elt)

            # Call body: inputs are state_vars..., scan_input_elts..., outer_vals...
            body_inputs = current_state + scan_input_elts + outer_vals
            outputs = self.body(*body_inputs)

            # Handle single vs multiple outputs
            if not isinstance(outputs, tuple):
                outputs = (outputs,)

            # First n_state_vars outputs are updated state
            current_state = list(outputs[: self.n_state_vars])

            # Remaining outputs are scan output elements
            for j in range(self.n_scan_outputs):
                scan_outputs[j].append(outputs[self.n_state_vars + j])

        # Prepare final outputs: state variables, then stacked scan outputs
        final_outputs = list(current_state)
        for j, scan_list in enumerate(scan_outputs):
            if scan_list:
                axis = self.scan_output_axes[j] if j < len(self.scan_output_axes) else 0
                direction = (
                    self.scan_output_directions[j]
                    if j < len(self.scan_output_directions)
                    else 0
                )
                # Stack along the specified axis
                stacked = torch.stack(scan_list, dim=axis)
                # Reverse if direction is 1 (prepending = reverse order)
                if direction == 1:
                    stacked = torch.flip(stacked, dims=[axis])
                final_outputs.append(stacked)
            else:
                # Empty scan output
                final_outputs.append(torch.tensor([]))

        return tuple(final_outputs)


@register("Scan", since_version=9)
def scan_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """ONNX Scan operator (version 9+).

    Scan iterates over one or more scan_input tensors, constructing scan_output tensors.
    It combines ideas from general recurrences, functional programming constructs
    such as scan, fold, map, and zip.

    Inputs: initial_state_and_scan_inputs (variadic) - N state vars followed by M scan inputs
    Outputs: final_state_and_scan_outputs (variadic) - N final states followed by K scan outputs

    Attributes:
    - body: The graph run each iteration
    - num_scan_inputs: Number of scan inputs M
    - scan_input_axes: Axis to scan for each scan input (default: 0)
    - scan_input_directions: Direction for each scan input (0=forward, 1=reverse)
    - scan_output_axes: Axis for each scan output (default: 0)
    - scan_output_directions: Direction for each scan output (0=append, 1=prepend)
    """
    # Get body subgraph
    body_graph = get_attribute(node, "body")
    if body_graph is None:
        raise ValueError("Scan operator requires 'body' attribute")

    # Get num_scan_inputs attribute (required)
    num_scan_inputs = get_attribute(node, "num_scan_inputs")
    if num_scan_inputs is None:
        raise ValueError("Scan operator requires 'num_scan_inputs' attribute")

    # Get optional attributes
    scan_input_axes = get_attribute(node, "scan_input_axes") or []
    scan_input_directions = get_attribute(node, "scan_input_directions") or []
    scan_output_axes = get_attribute(node, "scan_output_axes") or []
    scan_output_directions = get_attribute(node, "scan_output_directions") or []

    # Parse inputs: first (len(node.input) - num_scan_inputs) are state variables
    n_state_vars = len(node.input) - num_scan_inputs

    state_inputs = [builder.get_value(node.input[i]) for i in range(n_state_vars)]
    scan_inputs = [
        builder.get_value(node.input[i]) for i in range(n_state_vars, len(node.input))
    ]

    # Get parent type info if available (for nested subgraphs)
    parent_type_info = getattr(builder, "_type_info", None)

    # Build subgraph module
    body_module, body_input_names, body_output_names, outer_refs = (
        _build_subgraph_module(
            body_graph,
            builder.env,
            builder._opset_versions,
            parent_type_info,
            tensor_loader=builder.load_tensor,
        )
    )

    # Get outer scope values that the body references
    outer_values = [builder.get_value(name) for name in outer_refs]

    # Number of scan outputs
    n_scan_outputs = len(body_output_names) - n_state_vars

    # Create the scan module
    scan_module = ScanModule(
        body_module,
        n_state_vars,
        num_scan_inputs,
        n_scan_outputs,
        len(outer_values),
        list(scan_input_axes),
        list(scan_output_axes),
        list(scan_input_directions),
        list(scan_output_directions),
    )

    # Register the scan module
    module_name = builder.register_submodule(f"scan_{node.name or 'op'}", scan_module)

    # Build args: state_vars..., scan_inputs..., outer_vals...
    args = state_inputs + scan_inputs + outer_values

    result = builder.call_module(module_name, args=tuple(args))

    return result


@register("Scan", since_version=8)
def scan_op_v8(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """ONNX Scan operator (version 8).

    Version 8 has batching support and different input format:
    - First input is optional sequence_lens
    - Requires batch dimension (axis 0) and sequence dimension (axis 1)

    For simplicity, this implementation handles the common case where
    sequence_lens is empty (all sequences have same length).
    """
    # Get body subgraph
    body_graph = get_attribute(node, "body")
    if body_graph is None:
        raise ValueError("Scan operator requires 'body' attribute")

    # Get num_scan_inputs attribute (required)
    num_scan_inputs = get_attribute(node, "num_scan_inputs")
    if num_scan_inputs is None:
        raise ValueError("Scan operator requires 'num_scan_inputs' attribute")

    # Get optional directions attribute (v8 uses 'directions' instead of scan_input_directions)
    directions = get_attribute(node, "directions") or []

    # In v8, first input is optional sequence_lens, rest are state vars + scan inputs
    # Check if first input is empty (sequence_lens is optional)
    has_sequence_lens = node.input[0] != ""

    if has_sequence_lens:
        # sequence_lens provided - we ignore it for now (assume fixed length)
        start_idx = 1
    else:
        start_idx = 1  # Skip the empty sequence_lens input

    # Parse remaining inputs: state variables followed by scan inputs
    remaining_inputs = list(node.input[start_idx:])
    n_state_vars = len(remaining_inputs) - num_scan_inputs

    state_inputs = [builder.get_value(remaining_inputs[i]) for i in range(n_state_vars)]
    scan_inputs = [
        builder.get_value(remaining_inputs[i])
        for i in range(n_state_vars, len(remaining_inputs))
    ]

    # Get parent type info if available (for nested subgraphs)
    parent_type_info = getattr(builder, "_type_info", None)

    # Build subgraph module
    body_module, body_input_names, body_output_names, outer_refs = (
        _build_subgraph_module(
            body_graph,
            builder.env,
            builder._opset_versions,
            parent_type_info,
            tensor_loader=builder.load_tensor,
        )
    )

    # Get outer scope values that the body references
    outer_values = [builder.get_value(name) for name in outer_refs]

    # Number of scan outputs
    n_scan_outputs = len(body_output_names) - n_state_vars

    # Create the scan module for v8 (handles batching)
    # Note: In v8, batch axis is 0 and sequence axis is 1 (handled in ScanModuleV8)
    scan_module = ScanModuleV8(
        body_module,
        n_state_vars,
        num_scan_inputs,
        n_scan_outputs,
        len(outer_values),
        list(directions),
    )

    # Register the scan module
    module_name = builder.register_submodule(
        f"scan_v8_{node.name or 'op'}", scan_module
    )

    # Build args: state_vars..., scan_inputs..., outer_vals...
    args = state_inputs + scan_inputs + outer_values

    result = builder.call_module(module_name, args=tuple(args))

    return result


class ScanModuleV8(nn.Module):
    """Module that executes an ONNX Scan operation (version 8 with batching)."""

    def __init__(
        self,
        body_module: torch.fx.GraphModule,
        n_state_vars: int,
        n_scan_inputs: int,
        n_scan_outputs: int,
        n_outer_vars: int,
        directions: List[int],
    ):
        super().__init__()
        self.body = body_module
        self.n_state_vars = n_state_vars
        self.n_scan_inputs = n_scan_inputs
        self.n_scan_outputs = n_scan_outputs
        self.n_outer_vars = n_outer_vars
        self.directions = directions

    def forward(self, *args) -> Tuple[torch.Tensor, ...]:
        """Execute the scan with batching.

        In v8, tensors have shape [batch, sequence, ...].
        State variables have shape [batch, ...].

        Args are: state_vars..., scan_inputs..., outer_vals...
        Returns final state variables followed by scan outputs.
        """
        # Split args
        state_vars = list(args[: self.n_state_vars])
        scan_inputs = list(
            args[self.n_state_vars : self.n_state_vars + self.n_scan_inputs]
        )
        outer_vals = list(args[self.n_state_vars + self.n_scan_inputs :])

        # Get batch size and sequence length from first scan input
        if self.n_scan_inputs > 0:
            first_input = scan_inputs[0]
            batch_size = first_input.shape[0]
            sequence_length = first_input.shape[1]
        else:
            batch_size = state_vars[0].shape[0] if state_vars else 1
            sequence_length = 0

        # Process each batch
        batch_final_states: List[List[torch.Tensor]] = [
            [] for _ in range(self.n_state_vars)
        ]
        batch_scan_outputs: List[List[torch.Tensor]] = [
            [] for _ in range(self.n_scan_outputs)
        ]

        for batch in range(batch_size):
            # Get batch slice of state variables
            current_state = [sv[batch] for sv in state_vars]

            # Initialize scan output accumulators for this batch
            scan_outputs: List[List[torch.Tensor]] = [
                [] for _ in range(self.n_scan_outputs)
            ]

            # Execute loop over sequence
            for t in range(sequence_length):
                # Extract scan input elements for this batch and time step
                scan_input_elts = []
                for i, scan_input in enumerate(scan_inputs):
                    direction = self.directions[i] if i < len(self.directions) else 0
                    idx = sequence_length - 1 - t if direction == 1 else t
                    # scan_input has shape [batch, sequence, ...]
                    elt = scan_input[batch, idx]
                    scan_input_elts.append(elt)

                # Call body
                body_inputs = current_state + scan_input_elts + outer_vals
                outputs = self.body(*body_inputs)

                if not isinstance(outputs, tuple):
                    outputs = (outputs,)

                # Update state
                current_state = list(outputs[: self.n_state_vars])

                # Collect scan outputs
                for j in range(self.n_scan_outputs):
                    scan_outputs[j].append(outputs[self.n_state_vars + j])

            # Store final state for this batch
            for i, state in enumerate(current_state):
                batch_final_states[i].append(state)

            # Stack scan outputs for this batch
            for j, scan_list in enumerate(scan_outputs):
                if scan_list:
                    stacked = torch.stack(scan_list, dim=0)
                    batch_scan_outputs[j].append(stacked)

        # Stack across batches
        final_outputs: List[torch.Tensor] = []

        # Final state variables: stack across batch
        for states in batch_final_states:
            final_outputs.append(torch.stack(states, dim=0))

        # Scan outputs: stack across batch
        for outputs in batch_scan_outputs:
            if outputs:
                final_outputs.append(torch.stack(outputs, dim=0))

        return tuple(final_outputs)


@register("If")
def if_op(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
    """ONNX If operator.

    If has input: cond (boolean condition)
    If has attributes: then_branch (GraphProto), else_branch (GraphProto)
    Both branches must have the same number and types of outputs.
    """
    # Get condition
    cond = builder.get_value(node.input[0])

    # Get branch subgraphs
    then_graph = get_attribute(node, "then_branch")
    else_graph = get_attribute(node, "else_branch")

    if then_graph is None or else_graph is None:
        raise ValueError(
            "If operator requires 'then_branch' and 'else_branch' attributes"
        )

    # Get parent type info if available (for nested subgraphs)
    parent_type_info = getattr(builder, "_type_info", None)

    # Build subgraph modules for both branches
    then_module, then_input_names, then_output_names, then_outer_refs = (
        _build_subgraph_module(
            then_graph,
            builder.env,
            builder._opset_versions,
            parent_type_info,
            tensor_loader=builder.load_tensor,
        )
    )
    else_module, else_input_names, else_output_names, else_outer_refs = (
        _build_subgraph_module(
            else_graph,
            builder.env,
            builder._opset_versions,
            parent_type_info,
            tensor_loader=builder.load_tensor,
        )
    )

    # Get outer scope values for both branches
    then_outer_values = [builder.get_value(name) for name in then_outer_refs]
    else_outer_values = [builder.get_value(name) for name in else_outer_refs]

    # Create the if module
    if_module = IfModule(
        then_module,
        else_module,
        len(then_outer_values),
        len(else_outer_values),
    )

    # Register the if module
    module_name = builder.register_submodule(f"if_{node.name or 'op'}", if_module)

    # Build flat args for call_module: condition, then_outer..., else_outer...
    args = [cond] + then_outer_values + else_outer_values

    result = builder.call_module(module_name, args=tuple(args))

    return result
