from __future__ import annotations

import gc
import os
import time
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Set, Tuple

import networkx as nx
import pytest
from inference.core.workflows.errors import AssumptionError
from inference.core.workflows.execution_engine.constants import \
    NODE_COMPILATION_OUTPUT_PROPERTY
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    ExecutionGraphNode
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    _collect_unique_control_flow_lineages_with_step_mapping
from inference.core.workflows.execution_engine.v1.compiler.utils import (
    identify_lineage, node_as)
from networkx import DiGraph


def codeflash_wrap(codeflash_wrapped, codeflash_test_module_name, codeflash_test_class_name, codeflash_test_name, codeflash_function_name, codeflash_line_id, codeflash_loop_index, *args, **kwargs):
    test_id = f'{codeflash_test_module_name}:{codeflash_test_class_name}:{codeflash_test_name}:{codeflash_line_id}:{codeflash_loop_index}'
    if not hasattr(codeflash_wrap, 'index'):
        codeflash_wrap.index = {}
    if test_id in codeflash_wrap.index:
        codeflash_wrap.index[test_id] += 1
    else:
        codeflash_wrap.index[test_id] = 0
    codeflash_test_index = codeflash_wrap.index[test_id]
    invocation_id = f'{codeflash_line_id}_{codeflash_test_index}'
    test_stdout_tag = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_loop_index}:{invocation_id}"
    print(f'!$######{test_stdout_tag}######$!')
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = codeflash_wrapped(*args, **kwargs)
        codeflash_duration = time.perf_counter_ns() - counter
    except Exception as e:
        codeflash_duration = time.perf_counter_ns() - counter
        exception = e
    gc.enable()
    print(f'!######{test_stdout_tag}:{codeflash_duration}######!')
    if exception:
        raise exception
    return return_value

@dataclass
class StepNode(ExecutionGraphNode):
    """Concrete implementation of ExecutionGraphNode for testing"""
    node_category: str = 'step'
    name: str = ''
    selector: str = ''
    data_lineage: List[str] = field(default_factory=list)

def _create_step_node(selector: str, name: str, data_lineage: List[str]) -> StepNode:
    """Helper to create a StepNode with the given properties"""
    return StepNode(node_category='step', name=name, selector=selector, data_lineage=data_lineage)

def _add_step_to_graph(graph: nx.DiGraph, selector: str, name: str, data_lineage: List[str]) -> None:
    """Helper to add a step node to the execution graph"""
    step_node = _create_step_node(selector=selector, name=name, data_lineage=data_lineage)
    graph.add_node(selector)
    graph.nodes[selector][NODE_COMPILATION_OUTPUT_PROPERTY] = step_node

class TestCollectUniqueControlFlowLineagesBasic:
    """Basic functionality tests for normal conditions"""

    def test_single_step_with_lineage(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test collecting lineage from a single step with non-empty lineage'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1', 'input_2'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_single_step_with_lineage', '_collect_unique_control_flow_lineages_with_step_mapping', '3', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage(['input_1', 'input_2'])

    def test_single_step_with_empty_lineage(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test collecting from a single step with empty lineage - should not appear in unique_lineages'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', [])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_single_step_with_empty_lineage', '_collect_unique_control_flow_lineages_with_step_mapping', '3', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage([])

    def test_two_steps_different_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test collecting from two steps with different lineages'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1'])
        _add_step_to_graph(graph, 'step_2', 'step_2', ['input_2'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_two_steps_different_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1', 'step_2'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_two_steps_same_lineage(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test collecting from two steps with identical lineages - should deduplicate'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1', 'input_2'])
        _add_step_to_graph(graph, 'step_2', 'step_2', ['input_1', 'input_2'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_two_steps_same_lineage', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1', 'step_2'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage(['input_1', 'input_2'])

    def test_three_steps_two_identical_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with three steps where two share the same lineage'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_a', 'input_b'])
        _add_step_to_graph(graph, 'step_2', 'step_2', ['input_x', 'input_y'])
        _add_step_to_graph(graph, 'step_3', 'step_3', ['input_a', 'input_b'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_three_steps_two_identical_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, ['step_1', 'step_2', 'step_3'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        step_counts = sorted([len(steps) for steps in lineage_mapping.values()])

    def test_preserves_lineage_order(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that the function preserves the order of elements within a lineage'
        graph = nx.DiGraph()
        ordered_lineage = ['z', 'a', 'm', 'b']
        _add_step_to_graph(graph, 'step_1', 'step_1', ordered_lineage)
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesBasic', 'test_preserves_lineage_order', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, _ = codeflash_return_value

class TestCollectUniqueControlFlowLineagesEdge:
    """Edge case tests for unusual or boundary conditions"""

    def test_empty_control_flow_steps_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with empty list of control flow steps'
        graph = nx.DiGraph()
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_empty_control_flow_steps_list', '_collect_unique_control_flow_lineages_with_step_mapping', '2', codeflash_loop_index, [], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_step_not_in_graph_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that missing step in graph raises AssumptionError'
        graph = nx.DiGraph()
        with pytest.raises(AssumptionError):
            codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_step_not_in_graph_raises_error', '_collect_unique_control_flow_lineages_with_step_mapping', '2_0', codeflash_loop_index, ['nonexistent_step'], graph)

    def test_step_missing_compilation_output_raises_error(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that step without NODE_COMPILATION_OUTPUT_PROPERTY raises AssumptionError'
        graph = nx.DiGraph()
        graph.add_node('step_1')
        with pytest.raises(AssumptionError):
            codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_step_missing_compilation_output_raises_error', '_collect_unique_control_flow_lineages_with_step_mapping', '3_0', codeflash_loop_index, ['step_1'], graph)

    def test_lineage_with_single_element(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test lineage containing only one element'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['single_input'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_lineage_with_single_element', '_collect_unique_control_flow_lineages_with_step_mapping', '3', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_lineage_with_duplicate_elements(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test lineage containing duplicate elements'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input', 'input', 'input'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_lineage_with_duplicate_elements', '_collect_unique_control_flow_lineages_with_step_mapping', '3', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_very_long_lineage(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with a very long lineage list'
        graph = nx.DiGraph()
        long_lineage = [f'input_{i}' for i in range(100)]
        _add_step_to_graph(graph, 'step_1', 'step_1', long_lineage)
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_very_long_lineage', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_lineage_with_special_characters(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test lineage containing special characters in element names'
        graph = nx.DiGraph()
        special_lineage = ['input-1', 'input_2', 'input.3', 'input$4']
        _add_step_to_graph(graph, 'step_1', 'step_1', special_lineage)
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_lineage_with_special_characters', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_steps_with_same_name_different_selector(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that selector (not name) is used for deduplication in mapping'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'selector_1', 'same_name', ['input_1'])
        _add_step_to_graph(graph, 'selector_2', 'same_name', ['input_1'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_steps_with_same_name_different_selector', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['selector_1', 'selector_2'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage(['input_1'])

    def test_all_empty_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test when all steps have empty lineages'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', [])
        _add_step_to_graph(graph, 'step_2', 'step_2', [])
        _add_step_to_graph(graph, 'step_3', 'step_3', [])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_all_empty_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, ['step_1', 'step_2', 'step_3'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage([])

    def test_mixed_empty_and_non_empty_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test mix of empty and non-empty lineages'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1'])
        _add_step_to_graph(graph, 'step_2', 'step_2', [])
        _add_step_to_graph(graph, 'step_3', 'step_3', ['input_2'])
        _add_step_to_graph(graph, 'step_4', 'step_4', [])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesEdge', 'test_mixed_empty_and_non_empty_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '6', codeflash_loop_index, ['step_1', 'step_2', 'step_3', 'step_4'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

class TestCollectUniqueControlFlowLineagesLargeScale:
    """Large-scale tests with significant data volumes"""

    def test_many_steps_unique_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with many steps each having unique lineage'
        graph = nx.DiGraph()
        num_steps = 100
        for i in range(num_steps):
            _add_step_to_graph(graph, f'step_{i}', f'step_{i}', [f'input_{i}'])
        step_selectors = [f'step_{i}' for i in range(num_steps)]
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_many_steps_unique_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, step_selectors, graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        for steps in lineage_mapping.values():
            pass

    def test_many_steps_identical_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with many steps all having the same lineage'
        graph = nx.DiGraph()
        num_steps = 100
        shared_lineage = ['shared_input_1', 'shared_input_2']
        for i in range(num_steps):
            _add_step_to_graph(graph, f'step_{i}', f'step_{i}', shared_lineage)
        step_selectors = [f'step_{i}' for i in range(num_steps)]
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_many_steps_identical_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '6', codeflash_loop_index, step_selectors, graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        lineage_id = identify_lineage(shared_lineage)

    def test_large_lineage_list(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with lineage containing many elements'
        graph = nx.DiGraph()
        large_lineage = [f'input_{i}' for i in range(500)]
        _add_step_to_graph(graph, 'step_1', 'step_1', large_lineage)
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_large_lineage_list', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1'], graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_many_steps_few_unique_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with many steps but only a few unique lineages'
        graph = nx.DiGraph()
        num_steps = 100
        num_unique = 5
        for i in range(num_steps):
            lineage_idx = i % num_unique
            lineage = [f'input_{lineage_idx}_{j}' for j in range(3)]
            _add_step_to_graph(graph, f'step_{i}', f'step_{i}', lineage)
        step_selectors = [f'step_{i}' for i in range(num_steps)]
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_many_steps_few_unique_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '6', codeflash_loop_index, step_selectors, graph)
        unique_lineages, lineage_mapping = codeflash_return_value
        for steps in lineage_mapping.values():
            pass

    def test_alternating_empty_and_non_empty_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test with alternating pattern of empty and non-empty lineages'
        graph = nx.DiGraph()
        num_steps = 200
        for i in range(num_steps):
            if i % 2 == 0:
                lineage = [f'input_{i}']
            else:
                lineage = []
            _add_step_to_graph(graph, f'step_{i}', f'step_{i}', lineage)
        step_selectors = [f'step_{i}' for i in range(num_steps)]
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_alternating_empty_and_non_empty_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, step_selectors, graph)
        unique_lineages, lineage_mapping = codeflash_return_value

    def test_preserves_order_of_unique_lineages(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that unique lineages are collected in order of first appearance'
        graph = nx.DiGraph()
        num_steps = 50
        for i in range(num_steps):
            lineage_idx = i % 2
            lineage = [f'lineage_{lineage_idx}']
            _add_step_to_graph(graph, f'step_{i}', f'step_{i}', lineage)
        step_selectors = [f'step_{i}' for i in range(num_steps)]
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_preserves_order_of_unique_lineages', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, step_selectors, graph)
        unique_lineages, _ = codeflash_return_value

    def test_return_types_are_correct(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that return types match the specification'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1'])
        _add_step_to_graph(graph, 'step_2', 'step_2', [])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_return_types_are_correct', '_collect_unique_control_flow_lineages_with_step_mapping', '4', codeflash_loop_index, ['step_1', 'step_2'], graph)
        codeflash_output = codeflash_return_value
        result = codeflash_output
        unique_lineages, lineage_mapping = result
        for lineage in unique_lineages:
            for element in lineage:
                pass
        for lineage_id, selectors in lineage_mapping.items():
            for selector in selectors:
                pass

    def test_deterministic_output(self):
        codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
        'Test that function produces deterministic output for same input'
        graph = nx.DiGraph()
        _add_step_to_graph(graph, 'step_1', 'step_1', ['input_1', 'input_2'])
        _add_step_to_graph(graph, 'step_2', 'step_2', ['input_3'])
        _add_step_to_graph(graph, 'step_3', 'step_3', ['input_1', 'input_2'])
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_deterministic_output', '_collect_unique_control_flow_lineages_with_step_mapping', '5', codeflash_loop_index, ['step_1', 'step_2', 'step_3'], graph)
        codeflash_output = codeflash_return_value
        result1 = codeflash_output
        codeflash_return_value = codeflash_wrap(_collect_unique_control_flow_lineages_with_step_mapping, 'tests.inference.unit_tests.test__collect_unique_control_flow_lineages_with_step_mapping__unit_test_1', 'TestCollectUniqueControlFlowLineagesLargeScale', 'test_deterministic_output', '_collect_unique_control_flow_lineages_with_step_mapping', '7', codeflash_loop_index, ['step_1', 'step_2', 'step_3'], graph)
        codeflash_output = codeflash_return_value
        result2 = codeflash_output
        unique_lineages1, lineage_mapping1 = result1
        unique_lineages2, lineage_mapping2 = result2