"""Function builders for generating complete JAX functions.

This module provides builders for:
- Init functions (compute cache from parameters)
- Eval functions (compute residuals/Jacobian from cache + voltages)
"""

import ast
import logging
import re
from collections import defaultdict
from typing import Any, Dict, List, Optional, Set, Tuple

logger = logging.getLogger(__name__)

from ..mir.cfg import CFGAnalyzer, LoopInfo
from ..mir.constprop import SCCP
from ..mir.ssa import PHIResolution, PHIResolutionType, SSAAnalyzer
from ..mir.types import Block, MIRFunction, MIRInstruction, ValueId
from ..openvaf_ast import (
    assign,
    assign_tuple,
    attr,
    binop,
    function_def,
    import_from,
    jnp_bool,
    jnp_call,
    jnp_where,
    list_expr,
    return_stmt,
    subscript,
    tuple_expr,
    unaryop,
)
from ..openvaf_ast import (
    call as ast_call,
)
from ..openvaf_ast import (
    const as ast_const,
)
from ..openvaf_ast import (
    name as ast_name,
)
from ..openvaf_ast.statements import build_module
from .context import CodeGenContext, build_context_from_mir
from .instruction import InstructionTranslator


class FunctionBuilder:
    """Base class for function builders."""

    def __init__(self, mir_func: MIRFunction):
        """Initialize builder with MIR function.

        Args:
            mir_func: Parsed MIR function
        """
        self.mir_func = mir_func
        self.cfg = CFGAnalyzer(mir_func)
        self.ssa = SSAAnalyzer(mir_func, self.cfg)
        self.codegen_warnings: list[str] = []  # Warnings from code generation
        self.simparam_metadata: dict = {}  # Simparam registry metadata after build

    def _emit_preamble(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit function preamble: imports."""
        # JAX imports: import jax; from jax import numpy as jnp, lax
        # The 'import jax' is needed for jax.debug.print used by $display
        body.append(ast.Import(names=[ast.alias(name="jax", asname=None)]))
        body.append(import_from("jax", [("numpy", "jnp"), "lax"]))
        # Note: 0.0 and 1.0 constants are now inlined via ctx.zero() and ctx.one()

    def _emit_constants(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit constant definitions."""
        # Float constants
        for name, value in self.mir_func.constants.items():
            var_name = f"{ctx.var_prefix}{name}"
            if value == float("inf"):
                expr = attr(ast_name("jnp"), "inf")  # jnp.inf (constant, not call)
            elif value == float("-inf"):
                expr = unaryop(ast.USub(), attr(ast_name("jnp"), "inf"))
            elif value != value:  # NaN
                expr = attr(ast_name("jnp"), "nan")  # jnp.nan (constant, not call)
            else:
                expr = ast_const(value)
            body.append(assign(var_name, expr))
            ctx.defined_vars.add(var_name)

        # Boolean constants
        for name, value in self.mir_func.bool_constants.items():
            var_name = f"{ctx.var_prefix}{name}"
            body.append(assign(var_name, jnp_bool(ast_const(value))))
            ctx.defined_vars.add(var_name)

        # Integer constants
        for name, value in self.mir_func.int_constants.items():
            var_name = f"{ctx.var_prefix}{name}"
            body.append(assign(var_name, ast_const(value)))
            ctx.defined_vars.add(var_name)

    def _emit_block(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        block: Block,
        translator: InstructionTranslator,
        loop_info: Optional[LoopInfo] = None,
    ):
        """Emit code for a basic block with PHI batching optimization.

        PHIs with the same condition are batched to reduce XLA graph complexity.
        """
        # Separate PHIs from regular instructions
        phi_insts = block.phi_nodes
        regular_insts = [
            inst for inst in block.instructions if not inst.is_phi and not inst.is_terminator
        ]

        # Batch PHIs with same condition
        self._emit_batched_phis(body, ctx, translator, phi_insts, loop_info)

        # Emit regular instructions
        for inst in regular_insts:
            expr = translator.translate(inst, loop_info)
            if expr and inst.result:
                var_name = ctx.define_var(inst.result)
                body.append(assign(var_name, expr))
            elif expr and not inst.result:
                # Side-effect only instruction (e.g., $display)
                body.append(ast.Expr(value=expr))

    def _emit_batched_phis(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        translator: InstructionTranslator,
        phi_insts: List[MIRInstruction],
        loop_info: Optional[LoopInfo] = None,
    ):
        """Emit PHI nodes with batching optimization.

        Groups PHIs by their condition and emits batched jnp.where calls
        for groups with 2+ PHIs sharing the same condition.
        """
        if not phi_insts:
            return

        # Resolve all PHIs and group by condition
        # Key: (condition_str, is_negated) -> list of (inst, resolution, true_val, false_val)
        cond_groups: Dict[Tuple[str, bool], List[Tuple[MIRInstruction, PHIResolution]]] = (
            defaultdict(list)
        )
        non_batchable: List[Tuple[MIRInstruction, PHIResolution]] = []

        for inst in phi_insts:
            if not inst.result:
                continue

            resolution = self.ssa.resolve_phi(inst, loop_info)

            # Only batch simple TWO_WAY PHIs (no nested resolutions)
            if (
                resolution.type == PHIResolutionType.TWO_WAY
                and resolution.condition
                and resolution.nested_true is None
                and resolution.nested_false is None
                and resolution.true_value
                and resolution.false_value
            ):
                # Normalize negated conditions for grouping
                cond_str = resolution.condition
                is_negated = cond_str.startswith("!")
                base_cond = cond_str[1:] if is_negated else cond_str

                cond_groups[(base_cond, is_negated)].append((inst, resolution))
            else:
                non_batchable.append((inst, resolution))

        # Emit batched groups (2+ PHIs with same condition)
        for (base_cond, is_negated), group in cond_groups.items():
            if len(group) >= 2:
                self._emit_phi_batch(body, ctx, translator, base_cond, is_negated, group)
            else:
                # Single PHI, emit normally
                inst, resolution = group[0]
                expr = translator._apply_phi_resolution(resolution)
                var_name = ctx.define_var(inst.result)
                body.append(assign(var_name, expr))

        # Emit non-batchable PHIs individually
        for inst, resolution in non_batchable:
            expr = translator._apply_phi_resolution(resolution)
            var_name = ctx.define_var(inst.result)
            body.append(assign(var_name, expr))

    def _emit_phi_batch(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        translator: InstructionTranslator,
        base_cond: str,
        is_negated: bool,
        group: List[Tuple[MIRInstruction, PHIResolution]],
    ):
        """Emit a batch of PHIs with the same condition using tree_map.

        Generates:
            (v1, v2, ...) = jax.tree_util.tree_map(
                lambda t, f: jnp.where(cond, t, f),
                (true1, true2, ...),
                (false1, false2, ...)
            )
        """
        # Build the condition expression
        cond_expr = ctx.get_operand(base_cond)
        if is_negated:
            cond_expr = jnp_call("logical_not", cond_expr)

        # Collect variable names and values
        var_names = []
        true_vals = []
        false_vals = []

        for inst, resolution in group:
            var_name = ctx.define_var(inst.result)
            var_names.append(var_name)
            true_vals.append(ctx.get_operand(resolution.true_value))
            false_vals.append(ctx.get_operand(resolution.false_value))

        # Build tree_map call:
        # jax.tree_util.tree_map(lambda t, f: jnp.where(cond, t, f), (t1, t2, ...), (f1, f2, ...))

        # Lambda: lambda t, f: jnp.where(cond, t, f)
        t_param = ast.arg(arg="t", annotation=None)
        f_param = ast.arg(arg="f", annotation=None)
        lambda_body = jnp_where(cond_expr, ast_name("t"), ast_name("f"))
        lambda_node = ast.Lambda(
            args=ast.arguments(
                posonlyargs=[],
                args=[t_param, f_param],
                vararg=None,
                kwonlyargs=[],
                kw_defaults=[],
                kwarg=None,
                defaults=[],
            ),
            body=lambda_body,
        )

        # jax.tree_util.tree_map(lambda_node, (true_vals), (false_vals))
        tree_map_call = ast_call(
            attr(attr(ast_name("jax"), "tree_util"), "tree_map"),
            [lambda_node, tuple_expr(true_vals), tuple_expr(false_vals)],
        )

        # Emit tuple assignment: (v1, v2, ...) = tree_map(...)
        body.append(assign_tuple(var_names, tree_map_call))

    def _emit_loop(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        loop: LoopInfo,
        translator: InstructionTranslator,
    ) -> None:
        """Emit code for a loop using lax.while_loop."""
        header_block = self.mir_func.blocks.get(loop.header)
        if not header_block:
            return

        # Find PHI nodes in header
        phi_nodes = header_block.phi_nodes
        if not phi_nodes:
            # No loop-carried values, emit blocks linearly
            for block_name in sorted(loop.body):
                block = self.mir_func.blocks.get(block_name)
                if block:
                    self._emit_block(body, ctx, block, translator)
            return

        # Extract loop-carried state from PHIs
        loop_state: List[Tuple[ValueId, ValueId, ValueId]] = []  # (result, init_val, update_val)
        for phi in phi_nodes:
            resolution = self.ssa.resolve_phi(phi, loop)
            if resolution.type == PHIResolutionType.LOOP_INIT:
                # All values must be present for LOOP_INIT resolution
                assert phi.result is not None
                assert resolution.init_value is not None
                assert resolution.update_value is not None
                loop_state.append((phi.result, resolution.init_value, resolution.update_value))

        if not loop_state:
            # Couldn't extract loop state, emit linearly
            for block_name in sorted(loop.body):
                block = self.mir_func.blocks.get(block_name)
                if block:
                    self._emit_block(body, ctx, block, translator)
            return

        # Build initial state tuple
        # Cast all values to float32 to ensure type consistency
        init_vals = [jnp_call("float32", ctx.get_operand(lc[1])) for lc in loop_state]
        init_state = tuple_expr(init_vals) if len(init_vals) > 1 else init_vals[0]

        # Pre-initialize PHI variables from non-header loop body blocks
        # These may be used in post-loop PHI resolutions (e.g., as branch conditions)
        # and must exist in the outer scope before the lax.while_loop
        self._pre_initialize_loop_body_phis(body, ctx, loop)

        # Find branch condition in header
        branch_cond = None
        header_term = header_block.terminator
        if header_term and header_term.is_branch:
            branch_cond = header_term.condition

        # Build condition function
        cond_body = self._build_loop_cond(ctx, loop, loop_state, translator, branch_cond)
        cond_fn = function_def("_loop_cond", ["_state"], cond_body)
        body.append(cond_fn)

        # Build body function
        loop_body_stmts = self._build_loop_body(ctx, loop, loop_state, translator)
        body_fn = function_def("_loop_body", ["_state"], loop_body_stmts)
        body.append(body_fn)

        # Call while_loop
        loop_call = ast_call(
            attr(ast_name("lax"), "while_loop"),
            [ast_name("_loop_cond"), ast_name("_loop_body"), init_state],
        )
        body.append(assign("_loop_result", loop_call))

        # Unpack results
        for i, (result, _, _) in enumerate(loop_state):
            var_name = ctx.define_var(result)
            if len(loop_state) > 1:
                body.append(assign(var_name, subscript(ast_name("_loop_result"), ast_const(i))))
            else:
                body.append(assign(var_name, ast_name("_loop_result")))

    def _build_loop_cond(
        self,
        ctx: CodeGenContext,
        loop: LoopInfo,
        loop_state: List[Tuple[ValueId, ValueId, ValueId]],
        translator: InstructionTranslator,
        branch_cond: Optional[ValueId],
    ) -> List[ast.stmt]:
        """Build loop condition function body."""
        cond_body = []

        # Unpack state
        state_vars = [f"{ctx.var_prefix}{ls[0]}" for ls in loop_state]
        if len(state_vars) > 1:
            cond_body.append(assign_tuple(state_vars, ast_name("_state")))
        else:
            cond_body.append(assign(state_vars[0], ast_name("_state")))

        # Mark state vars as defined in ctx so get_operand finds them
        # This is critical: without this, SCCP constants take precedence
        # and loop variables get replaced with their initial constant values
        for var in state_vars:
            ctx.defined_vars.add(var)
        local_defined = set(state_vars)

        # Process header block instructions to compute condition
        header_block = self.mir_func.blocks.get(loop.header)
        if header_block:
            for inst in header_block.body_instructions:
                expr = translator.translate(inst)
                if expr and inst.result:
                    var_name = f"{ctx.var_prefix}{inst.result}"
                    cond_body.append(assign(var_name, expr))
                    local_defined.add(var_name)

        # Return condition
        if branch_cond:
            cond_var = f"{ctx.var_prefix}{branch_cond}"
            if cond_var in local_defined or cond_var in ctx.defined_vars:
                cond_body.append(return_stmt(ast_name(cond_var)))
            else:
                cond_body.append(return_stmt(ctx.get_operand(branch_cond)))
        else:
            cond_body.append(return_stmt(jnp_bool(ast_const(False))))

        return cond_body

    def _build_loop_body(
        self,
        ctx: CodeGenContext,
        loop: LoopInfo,
        loop_state: List[Tuple[ValueId, ValueId, ValueId]],
        translator: InstructionTranslator,
    ) -> List[ast.stmt]:
        """Build loop body function."""
        loop_body = []

        # Unpack state
        state_vars = [f"{ctx.var_prefix}{ls[0]}" for ls in loop_state]
        if len(state_vars) > 1:
            loop_body.append(assign_tuple(state_vars, ast_name("_state")))
        else:
            loop_body.append(assign(state_vars[0], ast_name("_state")))

        # Mark state vars as defined in ctx so get_operand finds them
        # This is critical: without this, SCCP constants take precedence
        # and loop variables get replaced with their initial constant values
        for var in state_vars:
            ctx.defined_vars.add(var)
        local_defined = set(state_vars)

        # Process all blocks in loop
        for block_name in sorted(loop.body):
            block = self.mir_func.blocks.get(block_name)
            if not block:
                continue

            # Process PHIs for non-header blocks (header PHIs are the loop state)
            # These are regular control flow PHIs within the loop body
            if block_name != loop.header and block.phi_nodes:
                self._emit_batched_phis(loop_body, ctx, translator, block.phi_nodes, loop)
                # Add PHI results to defined_vars
                for phi in block.phi_nodes:
                    if phi.result:
                        var_name = f"{ctx.var_prefix}{phi.result}"
                        local_defined.add(var_name)
                        ctx.defined_vars.add(var_name)

            for inst in block.instructions:
                if inst.is_terminator or inst.is_phi:
                    continue
                expr = translator.translate(inst)
                if expr and inst.result:
                    var_name = f"{ctx.var_prefix}{inst.result}"
                    loop_body.append(assign(var_name, expr))
                    local_defined.add(var_name)
                    # Also add to ctx.defined_vars so subsequent get_operand calls find it
                    ctx.defined_vars.add(var_name)

        # Build return tuple with updated values
        # Cast all values to float32 to ensure type consistency with init values
        update_vals = []
        for result, _, update in loop_state:
            update_var = f"{ctx.var_prefix}{update}"
            if update_var in local_defined:
                val_expr = ast_name(update_var)
            else:
                # Fallback to operand resolution
                val_expr = ctx.get_operand(update)
            # Cast to float32 to ensure type consistency
            cast_expr = jnp_call("float32", val_expr)
            update_vals.append(cast_expr)

        if len(update_vals) > 1:
            loop_body.append(return_stmt(tuple_expr(update_vals)))
        else:
            loop_body.append(return_stmt(update_vals[0]))

        return loop_body

    def _pre_initialize_loop_body_phis(
        self, body: List[ast.stmt], ctx: CodeGenContext, loop: LoopInfo
    ) -> None:
        """Pre-initialize PHI variables from non-header loop body blocks.

        PHI nodes inside loop bodies (not the header) may be used in post-loop
        PHI resolutions as branch conditions. Since lax.while_loop creates a
        separate function scope for the loop body, these variables won't exist
        in the outer scope after the loop unless we pre-initialize them.

        This fixes NameError for variables like v3401 in BSIMSOI that are
        computed inside the loop but used to determine post-loop control flow.
        """
        for block_name in loop.body:
            if block_name == loop.header:
                continue  # Skip header PHIs - those become loop state

            block = self.mir_func.blocks.get(block_name)
            if not block or not block.phi_nodes:
                continue

            for phi in block.phi_nodes:
                if phi.result:
                    var_name = f"{ctx.var_prefix}{phi.result}"
                    if var_name not in ctx.defined_vars:
                        # Pre-initialize to False (boolean) - these are typically
                        # branch condition tracking variables
                        body.append(assign(var_name, jnp_bool(ast_const(False))))
                        ctx.defined_vars.add(var_name)


class InitFunctionBuilder(FunctionBuilder):
    """Builder for init functions."""

    def __init__(
        self,
        mir_func: MIRFunction,
        cache_mapping: List[Dict[str, Any]],
        collapse_decision_outputs: List[Tuple[int, str]],
    ):
        """Initialize init function builder.

        Args:
            mir_func: Parsed init MIR function
            cache_mapping: List of {init_value, eval_param} mappings
            collapse_decision_outputs: List of (pair_idx, value_name) tuples
        """
        super().__init__(mir_func)
        self.cache_mapping = cache_mapping
        self.collapse_decision_outputs = collapse_decision_outputs

    def build_simple(self, param_indices: List[int]) -> Tuple[str, List[str]]:
        """Build init function with single input array.

        Args:
            param_indices: List of param indices (usually range(n_params))

        Returns:
            Tuple of (function_name, code_lines)
        """
        # Create context
        ctx = build_context_from_mir(self.mir_func, var_prefix="")

        # Build function body
        body: List[ast.stmt] = []

        # Preamble
        self._emit_preamble(body, ctx)

        # Constants
        self._emit_constants(body, ctx)

        # Ensure v3 exists (commonly used for zero)
        if "v3" not in ctx.defined_vars:
            body.append(assign("v3", ctx.zero()))
            ctx.defined_vars.add("v3")

        # Map init params from input array
        self._emit_simple_param_mapping(body, ctx, param_indices)

        # Pre-initialize all cache output variables to 0.0 to avoid NameError
        # for variables only assigned in conditional branches (NMOS/PMOS paths)
        self._pre_initialize_cache_vars(body, ctx)

        # Process blocks
        translator = InstructionTranslator(ctx, self.ssa)
        block_order = self.cfg.topological_order()

        for item in block_order:
            if isinstance(item, LoopInfo):
                self._emit_loop(body, ctx, item, translator)
            else:
                block = self.mir_func.blocks.get(item)
                if block:
                    self._emit_block(body, ctx, block, translator)

        # Collect warnings from translator
        self.codegen_warnings.extend(translator.simparam_warnings)
        self.codegen_warnings.extend(translator.discontinuity_warnings)

        # Build cache output array
        self._emit_cache_output(body, ctx)

        # Build collapse decisions
        self._emit_collapse_decisions(body, ctx)

        # Return statement
        body.append(return_stmt(tuple_expr([ast_name("cache"), ast_name("collapse_decisions")])))

        # Build function
        func = function_def("init_fn", ["inputs"], body)

        # Compile to code
        module = build_module([func])
        ast.fix_missing_locations(module)
        code_str = ast.unparse(module)

        return "init_fn", code_str.split("\n")

    def _emit_simple_param_mapping(
        self, body: List[ast.stmt], ctx: CodeGenContext, param_indices: List[int]
    ):
        """Emit parameter mapping from single input array."""
        for init_idx, param in enumerate(self.mir_func.params):
            var_name = f"{ctx.var_prefix}{param}"

            if init_idx < len(param_indices):
                body.append(assign(var_name, subscript(ast_name("inputs"), ast_const(init_idx))))
            else:
                # Fallback to zero
                body.append(assign(var_name, ctx.zero()))

            ctx.defined_vars.add(var_name)

    def build_split(
        self, shared_indices: List[int], varying_indices: List[int], init_to_eval: List[int]
    ) -> Tuple[str, List[str]]:
        """Build init function with split shared/device params.

        Args:
            shared_indices: Eval param indices that are constant across devices
            varying_indices: Eval param indices that vary per device
            init_to_eval: Mapping from init param index to eval param index

        Returns:
            Tuple of (function_name, code_lines)
        """
        # Build index mappings
        shared_set = set(shared_indices)
        varying_set = set(varying_indices)
        shared_to_pos = {idx: pos for pos, idx in enumerate(shared_indices)}
        varying_to_pos = {idx: pos for pos, idx in enumerate(varying_indices)}

        # Create context
        ctx = build_context_from_mir(self.mir_func, var_prefix="")

        # Build function body
        body: List[ast.stmt] = []

        # Preamble
        self._emit_preamble(body, ctx)

        # Constants
        self._emit_constants(body, ctx)

        # Ensure v3 exists (commonly used for zero)
        if "v3" not in ctx.defined_vars:
            body.append(assign("v3", ctx.zero()))
            ctx.defined_vars.add("v3")

        # Map init params from split arrays
        self._emit_split_param_mapping(
            body, ctx, shared_set, varying_set, shared_to_pos, varying_to_pos, init_to_eval
        )

        # Pre-initialize all cache output variables to 0.0 to avoid NameError
        # for variables only assigned in conditional branches (NMOS/PMOS paths)
        self._pre_initialize_cache_vars(body, ctx)

        # Process blocks
        translator = InstructionTranslator(ctx, self.ssa)
        block_order = self.cfg.topological_order()

        for item in block_order:
            if isinstance(item, LoopInfo):
                self._emit_loop(body, ctx, item, translator)
            else:
                block = self.mir_func.blocks.get(item)
                if block:
                    self._emit_block(body, ctx, block, translator)

        # Collect warnings from translator
        self.codegen_warnings.extend(translator.simparam_warnings)
        self.codegen_warnings.extend(translator.discontinuity_warnings)

        # Build cache output array
        self._emit_cache_output(body, ctx)

        # Build collapse decisions
        self._emit_collapse_decisions(body, ctx)

        # Return statement
        body.append(return_stmt(tuple_expr([ast_name("cache"), ast_name("collapse_decisions")])))

        # Build function
        func = function_def("init_fn_split", ["shared_params", "device_params"], body)

        # Compile to code
        module = build_module([func])
        ast.fix_missing_locations(module)
        code_str = ast.unparse(module)

        return "init_fn_split", code_str.split("\n")

    def _emit_split_param_mapping(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        shared_set: Set[int],
        varying_set: Set[int],
        shared_to_pos: Dict[int, int],
        varying_to_pos: Dict[int, int],
        init_to_eval: List[int],
    ):
        """Emit parameter mapping from split arrays."""
        for init_idx, param in enumerate(self.mir_func.params):
            eval_idx = init_to_eval[init_idx] if init_idx < len(init_to_eval) else -1
            var_name = f"{ctx.var_prefix}{param}"

            if eval_idx in shared_set:
                pos = shared_to_pos[eval_idx]
                body.append(assign(var_name, subscript(ast_name("shared_params"), ast_const(pos))))
            elif eval_idx in varying_set:
                pos = varying_to_pos[eval_idx]
                body.append(assign(var_name, subscript(ast_name("device_params"), ast_const(pos))))
            else:
                # Fallback to zero
                body.append(assign(var_name, ctx.zero()))

            ctx.defined_vars.add(var_name)

    def _pre_initialize_cache_vars(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Pre-initialize all variables that appear in cache output to 0.0.

        This fixes a bug where variables assigned in conditional branches (e.g.,
        NMOS vs PMOS paths) would cause NameError when referenced in cache arrays.
        By pre-initializing to 0.0, all variables are guaranteed to exist even if
        the conditional branch that assigns them isn't taken at runtime.
        """
        cache_vars = set()
        for mapping in self.cache_mapping:
            init_val = mapping["init_value"]
            var_name = f"{ctx.var_prefix}{init_val}"
            cache_vars.add(var_name)

        # Pre-initialize all cache variables to 0.0
        # Skip variables that are already defined (constants, etc.)
        for var_name in sorted(cache_vars):  # Sort for deterministic output
            if var_name not in ctx.defined_vars:
                body.append(assign(var_name, ctx.zero()))
                ctx.defined_vars.add(var_name)

    def _emit_cache_output(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit cache array construction."""
        cache_vals: List[ast.expr] = []

        for mapping in self.cache_mapping:
            init_val = mapping["init_value"]
            var_name = f"{ctx.var_prefix}{init_val}"
            if var_name in ctx.defined_vars or init_val in ctx.defined_vars:
                cache_vals.append(ast_name(var_name if var_name in ctx.defined_vars else init_val))
            else:
                cache_vals.append(ctx.zero())

        if cache_vals:
            body.append(assign("cache", jnp_call("array", list_expr(cache_vals))))
        else:
            body.append(assign("cache", jnp_call("array", list_expr([]))))

    def _emit_collapse_decisions(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit collapse decision array construction."""
        collapse_vals: List[ast.expr] = []

        for pair_idx, val_name in self.collapse_decision_outputs:
            if val_name.startswith("!"):
                actual_val = val_name[1:]
                negate = True
            else:
                actual_val = val_name
                negate = False

            var_name = f"{ctx.var_prefix}{actual_val}"
            if var_name in ctx.defined_vars or actual_val in ctx.defined_vars:
                val_expr = ast_name(var_name if var_name in ctx.defined_vars else actual_val)
                if negate:
                    val_expr = jnp_call("logical_not", val_expr)
                collapse_vals.append(ast_call(attr(ast_name("jnp"), "float32"), [val_expr]))
            else:
                # Default based on negation
                collapse_vals.append(ctx.one() if negate else ctx.zero())

        if collapse_vals:
            body.append(assign("collapse_decisions", jnp_call("array", list_expr(collapse_vals))))
        else:
            body.append(assign("collapse_decisions", jnp_call("array", list_expr([]))))


class EvalFunctionBuilder(FunctionBuilder):
    """Builder for eval functions."""

    def __init__(
        self,
        mir_func: MIRFunction,
        dae_data: Dict[str, Any],
        cache_mapping: List[Dict[str, Any]],
        param_idx_to_val: Dict[int, str],
        sccp_known_values: Optional[Dict[str, Any]] = None,
        eval_param_names: Optional[List[str]] = None,
    ):
        """Initialize eval function builder.

        Args:
            mir_func: Parsed eval MIR function
            dae_data: DAE system data (residuals, jacobian)
            cache_mapping: Cache slot to eval param mapping
            param_idx_to_val: Maps eval param index to value name
            sccp_known_values: Optional dict mapping MIR value IDs to constant values
                              for SCCP-based dead code elimination. Used to eliminate
                              branches based on compile-time known values like TYPE=1.
                              Example: {'v51588': 1} for NMOS.
            eval_param_names: Optional list of eval parameter names (e.g., ['V(A,CI)', 'R', ...]).
                             Used for tracing voltage params to node names for lim_rhs computation.
        """
        # Run SCCP if known values provided
        self.sccp: Optional[SCCP] = None
        if sccp_known_values:
            self.sccp = SCCP(mir_func, known_values=sccp_known_values)

        # Call parent init after SCCP so we can pass it to SSAAnalyzer
        super().__init__(mir_func)

        # Re-create SSA analyzer with SCCP if available
        if self.sccp is not None:
            self.ssa = SSAAnalyzer(mir_func, self.cfg, sccp=self.sccp)

        self.dae_data = dae_data
        self.cache_mapping = cache_mapping
        self.param_idx_to_val = param_idx_to_val
        self.eval_param_names = eval_param_names or []
        self.limit_metadata = {}  # Will be populated by build_with_cache_split if limit functions used

    def build_with_cache_split(
        self,
        shared_indices: List[int],
        varying_indices: List[int],
        shared_cache_indices: Optional[List[int]] = None,
        varying_cache_indices: Optional[List[int]] = None,
        simparam_params: Optional[Dict[int, str]] = None,
        use_limit_functions: bool = False,
        limit_param_map: Optional[Dict[int, Tuple[str, str]]] = None,
    ) -> Tuple[str, List[str]]:
        """Build eval function with split params and optional split cache.

        Args:
            shared_indices: Param indices constant across devices
            varying_indices: Param indices that vary per device
            shared_cache_indices: Cache indices constant across devices
            varying_cache_indices: Cache indices that vary per device
            simparam_params: Dict mapping original param indices to simparam names.
                             These params will be read from simparams array instead
                             of shared/device params. Used for $abstime, $mfactor.
            use_limit_functions: If True, generate calls to limit_funcs['pnjlim']
                                instead of passing through voltage unchanged. When
                                enabled, the generated function has an additional
                                'limit_funcs' parameter that should be a dict like:
                                {'pnjlim': pnjlim_fn, 'fetlim': fetlim_fn}
            limit_param_map: Dict mapping original param indices to (kind, name) tuples
                            for limit-related params (prev_state, enable_lim, new_state,
                            enable_integration). These are read from limit_state_in or
                            set to constants instead of shared/device params.

        Returns:
            Tuple of (function_name, code_lines)
        """
        # Always use cache split for uniform interface
        shared_cache_indices = shared_cache_indices or []
        varying_cache_indices = varying_cache_indices or []
        simparam_params = simparam_params or {}
        limit_param_map = limit_param_map or {}

        # Build index mappings
        idx_mapping: Dict[int, Tuple[str, int]] = {}
        for new_idx, orig_idx in enumerate(shared_indices):
            idx_mapping[orig_idx] = ("shared", new_idx)
        for new_idx, orig_idx in enumerate(varying_indices):
            idx_mapping[orig_idx] = ("device", new_idx)
        # Add simparam mappings - these will be handled specially in _emit_param_mapping
        for orig_idx, simparam_name in simparam_params.items():
            idx_mapping[orig_idx] = ("simparam", simparam_name)
        # Add limit param mappings - read from limit_state_in or use constants
        for orig_idx, (kind, name) in limit_param_map.items():
            if kind == "prev_state":
                # Extract state index from name (e.g., "prev_state_0" -> 0)
                state_idx = int(name.split("_")[-1])
                idx_mapping[orig_idx] = ("prev_state", state_idx)
            elif kind == "enable_lim":
                idx_mapping[orig_idx] = ("enable_lim", 1.0 if use_limit_functions else 0.0)
            elif kind == "new_state":
                # new_state is written via StoreLimit, initialize to 0
                idx_mapping[orig_idx] = ("new_state", 0)
            elif kind == "enable_integration":
                # 0.0 for DC, could be 1.0 for transient if needed
                idx_mapping[orig_idx] = ("enable_integration", 0.0)

        # Build cache index mappings
        cache_idx_mapping: Dict[int, Tuple[str, int]] = {}
        for new_idx, orig_idx in enumerate(shared_cache_indices):
            cache_idx_mapping[orig_idx] = ("shared_cache", new_idx)
        for new_idx, orig_idx in enumerate(varying_cache_indices):
            cache_idx_mapping[orig_idx] = ("device_cache", new_idx)

        # Create context
        ctx = build_context_from_mir(self.mir_func, var_prefix="")
        ctx.use_limit_functions = use_limit_functions

        # Build function body
        body: List[ast.stmt] = []

        # Preamble
        self._emit_preamble(body, ctx)

        # Constants
        self._emit_constants(body, ctx)

        # Ensure v3 exists
        if "v3" not in ctx.defined_vars:
            body.append(assign("v3", ctx.zero()))
            ctx.defined_vars.add("v3")

        # Map params from split arrays
        self._emit_param_mapping(body, ctx, idx_mapping)

        # Map cache values
        self._emit_cache_mapping(body, ctx, cache_idx_mapping)

        # Pre-initialize all output variables to 0.0 to avoid NameError
        # for variables only assigned in conditional branches (NMOS/PMOS paths)
        self._pre_initialize_output_vars(body, ctx)

        # Process blocks
        translator = InstructionTranslator(ctx, self.ssa)
        block_order = self.cfg.topological_order()

        for item in block_order:
            if isinstance(item, LoopInfo):
                self._emit_loop(body, ctx, item, translator)
            else:
                block = self.mir_func.blocks.get(item)
                if block:
                    self._emit_block(body, ctx, block, translator)

        # Collect warnings from translator
        self.codegen_warnings.extend(translator.simparam_warnings)
        self.codegen_warnings.extend(translator.discontinuity_warnings)

        # Collect simparam metadata from context
        self.simparam_metadata = ctx.get_simparam_metadata()

        # Collect limit metadata from context
        self.limit_metadata = ctx.get_limit_metadata()

        # Build output arrays
        self._emit_residual_arrays(body, ctx)
        self._emit_jacobian_arrays(body, ctx)
        self._emit_lim_rhs_arrays(body, ctx)
        self._emit_small_signal_arrays(body, ctx)

        # Build limit_state_out (always emitted - empty array when no limits used)
        # This ensures consistent function signature regardless of whether model uses limits
        self._emit_limit_state_out(body, ctx)

        # Return statement: (res_resist, res_react, jac_resist, jac_react,
        #                    lim_rhs_resist, lim_rhs_react,
        #                    small_signal_resist, small_signal_react, limit_state_out)
        # Note: limit_state_out is always returned (empty array if no limits) for uniform interface
        return_values = [
            ast_name("residuals_resist"),
            ast_name("residuals_react"),
            ast_name("jacobian_resist"),
            ast_name("jacobian_react"),
            ast_name("lim_rhs_resist"),
            ast_name("lim_rhs_react"),
            ast_name("small_signal_resist"),
            ast_name("small_signal_react"),
            ast_name("limit_state_out"),
        ]
        body.append(return_stmt(tuple_expr(return_values)))

        # Build function with uniform interface
        # Args: shared_params, device_params, shared_cache, device_cache, simparams, limit_state_in, limit_funcs
        # simparams layout: [analysis_type, mfactor, gmin]
        fn_name = "eval_fn"
        args = [
            "shared_params",
            "device_params",
            "shared_cache",
            "device_cache",
            "simparams",
            "limit_state_in",
            "limit_funcs",
        ]

        func = function_def(fn_name, args, body)

        # Compile to code
        module = build_module([func])
        ast.fix_missing_locations(module)
        code_str = ast.unparse(module)

        return fn_name, code_str.split("\n")

    def _emit_param_mapping(
        self, body: List[ast.stmt], ctx: CodeGenContext, idx_mapping: Dict[int, Tuple[str, any]]
    ):
        """Emit parameter mapping from split arrays.

        Args:
            body: List to append statements to
            ctx: Code generation context
            idx_mapping: Maps original param index to (source, value) where:
                - source='shared': value is new index in shared_params
                - source='device': value is new index in device_params
                - source='simparam': value is simparam name (e.g., '$abstime')
        """
        for i, param in enumerate(self.mir_func.params):
            var_name = f"{ctx.var_prefix}{param}"

            if i in idx_mapping:
                source, value = idx_mapping[i]
                if source == "shared":
                    body.append(
                        assign(var_name, subscript(ast_name("shared_params"), ast_const(value)))
                    )
                elif source == "device":
                    body.append(
                        assign(var_name, subscript(ast_name("device_params"), ast_const(value)))
                    )
                elif source == "simparam":
                    # Register simparam and emit simparams[idx]
                    simparam_name = value
                    simparam_idx = ctx.register_simparam(simparam_name)
                    body.append(
                        assign(var_name, subscript(ast_name("simparams"), ast_const(simparam_idx)))
                    )
                elif source == "prev_state":
                    # Read previous iteration's limited voltage from limit_state_in[N]
                    body.append(
                        assign(var_name, subscript(ast_name("limit_state_in"), ast_const(value)))
                    )
                elif source == "enable_lim":
                    # Constant flag: 1.0 when limiting enabled, 0.0 otherwise
                    body.append(assign(var_name, ast_const(value)))
                elif source == "new_state":
                    # Output state - initialized to 0, written via StoreLimit
                    body.append(assign(var_name, ast_const(0.0)))
                elif source == "enable_integration":
                    # Integration enable flag (0.0 for DC, 1.0 for transient)
                    body.append(assign(var_name, ast_const(value)))
            else:
                # Fallback: derivative selector params default to 0
                body.append(assign(var_name, ast_const(0)))

            ctx.defined_vars.add(var_name)

    def _emit_cache_mapping(
        self,
        body: List[ast.stmt],
        ctx: CodeGenContext,
        cache_idx_mapping: Dict[int, Tuple[str, int]],
    ):
        """Emit cache value mapping from split cache arrays.

        Always uses split cache format (shared_cache, device_cache) for uniform interface.
        """
        for cache_idx, mapping in enumerate(self.cache_mapping):
            eval_param_idx = mapping["eval_param"]
            eval_val = self.param_idx_to_val.get(eval_param_idx, f"cached_{eval_param_idx}")
            var_name = f"{ctx.var_prefix}{eval_val}"

            if cache_idx in cache_idx_mapping:
                source, new_idx = cache_idx_mapping[cache_idx]
                if source == "shared_cache":
                    body.append(
                        assign(var_name, subscript(ast_name("shared_cache"), ast_const(new_idx)))
                    )
                else:
                    body.append(
                        assign(var_name, subscript(ast_name("device_cache"), ast_const(new_idx)))
                    )
            else:
                # Cache index not in mapping - this shouldn't happen with proper setup
                # but fall back to zero for safety
                body.append(assign(var_name, ctx.zero()))

            ctx.defined_vars.add(var_name)

    def _emit_residual_arrays(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit residual output arrays."""
        resist_exprs: List[ast.expr] = []
        react_exprs: List[ast.expr] = []

        for res in self.dae_data["residuals"]:
            resist_var = self._mir_to_var(res["resist_var"], ctx)
            react_var = self._mir_to_var(res["react_var"], ctx)

            resist_exprs.append(
                ast_name(resist_var) if resist_var in ctx.defined_vars else ctx.zero()
            )
            react_exprs.append(ast_name(react_var) if react_var in ctx.defined_vars else ctx.zero())

        body.append(assign("residuals_resist", jnp_call("array", list_expr(resist_exprs))))
        body.append(assign("residuals_react", jnp_call("array", list_expr(react_exprs))))

    def _emit_jacobian_arrays(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit Jacobian output arrays."""
        resist_exprs: List[ast.expr] = []
        react_exprs: List[ast.expr] = []

        for entry in self.dae_data["jacobian"]:
            resist_var = self._mir_to_var(entry["resist_var"], ctx)
            react_var = self._mir_to_var(entry["react_var"], ctx)

            resist_exprs.append(
                ast_name(resist_var) if resist_var in ctx.defined_vars else ctx.zero()
            )
            react_exprs.append(ast_name(react_var) if react_var in ctx.defined_vars else ctx.zero())

        body.append(assign("jacobian_resist", jnp_call("array", list_expr(resist_exprs))))
        body.append(assign("jacobian_react", jnp_call("array", list_expr(react_exprs))))

    def _emit_lim_rhs_arrays(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit limiting RHS correction arrays.

        These corrections are subtracted from residuals during Newton-Raphson iteration
        when limiting is applied. The formula for each residual i is:

            lim_rhs[i] = sum_k J[i, hi_k] * (V_lim_k - V_raw_k)

        where:
            - k iterates over each limited branch
            - J[i, hi_k] is the Jacobian entry at (residual i, hi_node of limited branch k)
            - V_lim_k is the limited voltage (output of pnjlim/fetlim)
            - V_raw_k is the raw branch voltage before limiting
            - hi_k is the positive terminal node of limited branch k

        The corrected residual becomes:
            f_corrected = f_computed - lim_rhs

        This correction accounts for the fact that the device is evaluated at V_lim
        but the NR update is applied at V_raw.
        """
        n_residuals = len(self.dae_data["residuals"])

        # Collect limited branches: list of (lim_state_idx, v_raw_var, v_lim_var, hi_node_name)
        limited_branches = self._collect_limited_branches(ctx)

        if not limited_branches:
            # No limited branches traced - emit zeros (safe fallback)
            resist_exprs: List[ast.expr] = [ctx.zero() for _ in range(n_residuals)]
            react_exprs: List[ast.expr] = [ctx.zero() for _ in range(n_residuals)]
            body.append(assign("lim_rhs_resist", jnp_call("array", list_expr(resist_exprs))))
            body.append(assign("lim_rhs_react", jnp_call("array", list_expr(react_exprs))))
            return

        # Build node_name -> residual index lookup
        node_to_res_idx: Dict[str, int] = {}
        for i, res in enumerate(self.dae_data["residuals"]):
            node_to_res_idx[res["node_name"]] = i

        # Build (row_node_name, col_node_name) -> jacobian entry index lookup
        jac_key_to_idx: Dict[Tuple[str, str], int] = {}
        for j, entry in enumerate(self.dae_data["jacobian"]):
            jac_key_to_idx[(entry["row_node_name"], entry["col_node_name"])] = j

        # For each limited branch, compute delta = V_lim - V_raw as a named variable
        for branch_idx, (lim_idx, v_raw_var, v_lim_var, _hi_node) in enumerate(limited_branches):
            delta_name = f"_lim_delta_{branch_idx}"
            delta_expr = binop(ast_name(v_lim_var), ast.Sub(), ast_name(v_raw_var))
            body.append(assign(delta_name, delta_expr))
            ctx.defined_vars.add(delta_name)

        # Build resist and react lim_rhs expressions for each residual
        resist_exprs = []
        react_exprs = []
        for i, res in enumerate(self.dae_data["residuals"]):
            row_node = res["node_name"]
            resist_terms: List[ast.expr] = []
            react_terms: List[ast.expr] = []

            for branch_idx, (lim_idx, v_raw_var, v_lim_var, hi_node) in enumerate(limited_branches):
                jac_key = (row_node, hi_node)
                if jac_key in jac_key_to_idx:
                    j_entry_idx = jac_key_to_idx[jac_key]
                    j_entry = self.dae_data["jacobian"][j_entry_idx]

                    delta_name = f"_lim_delta_{branch_idx}"

                    # Resistive: jacobian_resist[j] * delta
                    j_resist_var = self._mir_to_var(j_entry["resist_var"], ctx)
                    if j_resist_var in ctx.defined_vars:
                        resist_terms.append(
                            binop(ast_name(j_resist_var), ast.Mult(), ast_name(delta_name))
                        )

                    # Reactive: jacobian_react[j] * delta
                    j_react_var = self._mir_to_var(j_entry["react_var"], ctx)
                    if j_react_var in ctx.defined_vars:
                        react_terms.append(
                            binop(ast_name(j_react_var), ast.Mult(), ast_name(delta_name))
                        )

            # Sum all terms for this residual, or zero if none
            resist_exprs.append(self._sum_exprs(resist_terms, ctx))
            react_exprs.append(self._sum_exprs(react_terms, ctx))

        body.append(assign("lim_rhs_resist", jnp_call("array", list_expr(resist_exprs))))
        body.append(assign("lim_rhs_react", jnp_call("array", list_expr(react_exprs))))

    def _collect_limited_branches(self, ctx: CodeGenContext) -> List[Tuple[int, str, str, str]]:
        """Collect info about each limited branch for lim_rhs computation.

        For each LimState, traces back to find:
        - V_raw: the raw voltage MIR variable (before limiting)
        - V_lim: the limited voltage MIR variable (after pnjlim/fetlim)
        - hi_node: the positive terminal node name of the limited branch

        Returns:
            List of (lim_state_idx, v_raw_var_name, v_lim_var_name, hi_node_name).
            Only includes branches that could be fully traced.
        """
        result = []

        # Build MIR param ID -> original param index mapping
        # MIR params are listed in the same order as param indices
        param_id_to_idx: Dict[str, int] = {}
        for i, param in enumerate(self.mir_func.params):
            param_id_to_idx[param] = i

        for lim_idx in range(ctx.limit_next_index):
            v_raw_id = ctx.limit_to_raw_operand.get(lim_idx)
            if v_raw_id is None:
                logger.debug(f"LimState {lim_idx}: no V_raw mapping, skipping")
                continue

            # Get the V_lim variable (the operand stored by StoreLimit)
            v_lim_id = ctx.limit_store_operands.get(lim_idx)
            if v_lim_id is None:
                logger.debug(f"LimState {lim_idx}: no V_lim store operand, skipping")
                continue

            v_raw_var = f"{ctx.var_prefix}{v_raw_id}"
            v_lim_var = f"{ctx.var_prefix}{v_lim_id}"

            # Verify both variables are defined in generated code
            if v_raw_var not in ctx.defined_vars:
                logger.debug(f"LimState {lim_idx}: V_raw var {v_raw_var} not defined, skipping")
                continue
            if v_lim_var not in ctx.defined_vars:
                logger.debug(f"LimState {lim_idx}: V_lim var {v_lim_var} not defined, skipping")
                continue

            # Trace V_raw back to a voltage param to find hi_node
            # V_raw is typically a direct param read: V_raw = device_params[N]
            # The param index N maps to param_idx_to_val which has the voltage name
            hi_node = self._find_hi_node_for_operand(v_raw_id, param_id_to_idx)
            if hi_node is None:
                logger.debug(
                    f"LimState {lim_idx}: could not trace V_raw {v_raw_id} to a voltage param, "
                    f"skipping"
                )
                continue

            result.append((lim_idx, v_raw_var, v_lim_var, hi_node))

        return result

    def _find_hi_node_for_operand(
        self, v_raw_id: str, param_id_to_idx: Dict[str, int]
    ) -> Optional[str]:
        """Find the hi (positive) node name for a voltage operand.

        Traces a MIR value ID back to a voltage parameter name like V(A,CI)
        and extracts the first (hi) node name.

        Args:
            v_raw_id: MIR value ID of the raw voltage (e.g., 'v123')
            param_id_to_idx: Maps MIR param IDs to original param indices

        Returns:
            Hi node name (e.g., 'A') or None if not traceable.
        """
        # Check if V_raw is a direct param reference
        if v_raw_id in param_id_to_idx:
            orig_idx = param_id_to_idx[v_raw_id]
            return self._voltage_param_to_hi_node(orig_idx)

        # V_raw might be computed (e.g., optbarrier of a param).
        # Look through all params for a match - the param_idx_to_val
        # maps eval param index to MIR value name.
        for orig_idx, val_name in self.param_idx_to_val.items():
            if val_name == v_raw_id:
                return self._voltage_param_to_hi_node(orig_idx)

        return None

    def _voltage_param_to_hi_node(self, param_idx: int) -> Optional[str]:
        """Extract the hi node name from a voltage parameter.

        Voltage params have names like V(A,CI) or V(A).
        The hi node is the first node name.

        Args:
            param_idx: Original eval param index

        Returns:
            Hi node name or None if not a voltage param.
        """
        if param_idx >= len(self.eval_param_names):
            return None

        param_name = self.eval_param_names[param_idx]

        # Parse V(hi_node, lo_node) or V(hi_node)
        m = re.match(r"^V\((\w+)(?:,(\w+))?\)$", param_name)
        if m:
            return m.group(1)

        return None

    @staticmethod
    def _sum_exprs(terms: List[ast.expr], ctx: CodeGenContext) -> ast.expr:
        """Sum a list of AST expressions, returning 0.0 if empty."""
        if not terms:
            return ctx.zero()
        result = terms[0]
        for term in terms[1:]:
            result = binop(result, ast.Add(), term)
        return result

    def _emit_small_signal_arrays(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit small-signal residual output arrays.

        These values are known to be zero in large-signal (DC) analysis but
        are needed for AC small-signal analysis. The separation allows the
        simulator to avoid unnecessary derivative computations during DC.
        """
        resist_exprs: List[ast.expr] = []
        react_exprs: List[ast.expr] = []

        for res in self.dae_data["residuals"]:
            # Get small-signal variables if they exist
            resist_ss_var = self._mir_to_var(res.get("resist_small_signal_var", ""), ctx)
            react_ss_var = self._mir_to_var(res.get("react_small_signal_var", ""), ctx)

            resist_exprs.append(
                ast_name(resist_ss_var) if resist_ss_var in ctx.defined_vars else ctx.zero()
            )
            react_exprs.append(
                ast_name(react_ss_var) if react_ss_var in ctx.defined_vars else ctx.zero()
            )

        body.append(assign("small_signal_resist", jnp_call("array", list_expr(resist_exprs))))
        body.append(assign("small_signal_react", jnp_call("array", list_expr(react_exprs))))

    def _emit_limit_state_out(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Emit limit_state_out array from stored limit values.

        When device-level $limit functions are enabled, StoreLimit instructions
        track which values should be stored for the next NR iteration.
        This method builds the limit_state_out array from those tracked values.

        If no limit states are used, emits an empty array.
        """
        limit_count = ctx.limit_next_index
        if limit_count == 0:
            # No limit states used - emit empty array
            body.append(assign("limit_state_out", jnp_call("array", list_expr([]))))
            return

        # Build array of stored limit values
        # ctx.limit_store_operands maps limit index -> MIR operand ID
        limit_exprs: List[ast.expr] = []
        for idx in range(limit_count):
            if idx in ctx.limit_store_operands:
                operand_id = ctx.limit_store_operands[idx]
                # Get the variable name for this operand
                var_name = f"{ctx.var_prefix}{operand_id}"
                if var_name in ctx.defined_vars:
                    limit_exprs.append(ast_name(var_name))
                else:
                    raise ValueError(
                        f"Limit state {idx}: operand {operand_id} -> variable {var_name} "
                        f"not in defined_vars. This indicates a code generation bug."
                    )
            else:
                raise ValueError(
                    f"Limit state {idx} was registered (LoadLimit) but no StoreLimit found. "
                    f"Expected StoreLimit for all registered limit states."
                )

        body.append(assign("limit_state_out", jnp_call("array", list_expr(limit_exprs))))

    def _mir_to_var(self, mir_ref: str, ctx: CodeGenContext) -> str:
        """Convert MIR reference (e.g., 'mir_123') to variable name."""
        if mir_ref and mir_ref.startswith("mir_"):
            # Extract value ID
            val_id = mir_ref[4:]  # Remove 'mir_' prefix
            return f"{ctx.var_prefix}v{val_id}"
        return mir_ref or ""

    def _pre_initialize_output_vars(self, body: List[ast.stmt], ctx: CodeGenContext):
        """Pre-initialize all variables that appear in output arrays to 0.0.

        This fixes a bug where variables assigned in conditional branches (e.g.,
        NMOS vs PMOS paths) would cause NameError when referenced in output arrays.
        By pre-initializing to 0.0, all variables are guaranteed to exist even if
        the conditional branch that assigns them isn't taken at runtime.

        The correct value will be assigned when the appropriate branch executes.
        If no branch assigns a value, it remains 0.0 (safe default).
        """
        # Collect all variables from residuals
        output_vars = set()
        for res in self.dae_data.get("residuals", []):
            for key in [
                "resist_var",
                "react_var",
                "resist_lim_rhs_var",
                "react_lim_rhs_var",
                "resist_small_signal_var",
                "react_small_signal_var",
            ]:
                mir_ref = res.get(key, "")
                if mir_ref:
                    var_name = self._mir_to_var(mir_ref, ctx)
                    if var_name:
                        output_vars.add(var_name)

        # Collect all variables from jacobian
        for entry in self.dae_data.get("jacobian", []):
            for key in ["resist_var", "react_var"]:
                mir_ref = entry.get(key, "")
                if mir_ref:
                    var_name = self._mir_to_var(mir_ref, ctx)
                    if var_name:
                        output_vars.add(var_name)

        # Pre-initialize all output variables to 0.0
        # Skip variables that are already defined (constants, etc.)
        for var_name in sorted(output_vars):  # Sort for deterministic output
            if var_name not in ctx.defined_vars:
                body.append(assign(var_name, ctx.zero()))
                ctx.defined_vars.add(var_name)
