from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path

import dill as pickle
import pytest
from inference.core.workflows.execution_engine.constants import \
    NODE_COMPILATION_OUTPUT_PROPERTY
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    NodeCategory
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    separate_control_flow_predecessors_from_data_providers
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}######!')
    iteration = os.environ['CODEFLASH_TEST_ITERATION']
    with Path('{codeflash_run_tmp_dir_client_side}', f'test_return_values_{iteration}.bin').open('ab') as f:
        pickled_values = pickle.dumps((args, kwargs, exception)) if exception else pickle.dumps((args, kwargs, return_value))
        _test_name = f"{codeflash_test_module_name}:{(codeflash_test_class_name + '.' if codeflash_test_class_name else '')}{codeflash_test_name}:{codeflash_function_name}:{codeflash_line_id}".encode('ascii')
        f.write(len(_test_name).to_bytes(4, byteorder='big'))
        f.write(_test_name)
        f.write(codeflash_duration.to_bytes(8, byteorder='big'))
        f.write(len(pickled_values).to_bytes(4, byteorder='big'))
        f.write(pickled_values)
        f.write(codeflash_loop_index.to_bytes(8, byteorder='big'))
        f.write(len(invocation_id).to_bytes(4, byteorder='big'))
        f.write(invocation_id.encode('ascii'))
    if exception:
        raise exception
    return return_value

class MockNodeOutput:
    """Real class to represent node compilation output with category and flow control info."""

    def __init__(self, node_category, controls_flow_value):
        self.node_category = node_category
        self._controls_flow = controls_flow_value

    def controls_flow(self):
        return self._controls_flow

def create_test_graph_with_nodes(nodes_config):
    """
    Helper function to create a DiGraph with configured nodes.
    nodes_config: list of tuples (node_name, node_category, controls_flow)
    Returns: DiGraph with nodes and their properties set.
    """
    graph = DiGraph()
    for node_name, node_category, controls_flow in nodes_config:
        node_output = MockNodeOutput(node_category, controls_flow)
        graph.add_node(node_name, **{NODE_COMPILATION_OUTPUT_PROPERTY: node_output})
    return graph

def test_empty_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a node that has no predecessors.'
    graph = create_test_graph_with_nodes([('node1', NodeCategory.STEP_NODE, False)])
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='node1')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_empty_predecessors', 'separate_control_flow_predecessors_from_data_providers', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_single_control_flow_predecessor():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a single control flow predecessor.'
    graph = create_test_graph_with_nodes([('cf_pred', NodeCategory.STEP_NODE, True), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_single_control_flow_predecessor', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_single_data_provider_predecessor():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a single data provider (non-control-flow) predecessor.'
    graph = create_test_graph_with_nodes([('data_pred', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('data_pred', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_single_data_provider_predecessor', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_mixed_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with both control flow and data provider predecessors.'
    graph = create_test_graph_with_nodes([('cf_pred1', NodeCategory.STEP_NODE, True), ('data_pred1', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred1', 'target')
    graph.add_edge('data_pred1', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_mixed_predecessors', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_multiple_control_flow_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple control flow predecessors.'
    graph = create_test_graph_with_nodes([('cf_pred1', NodeCategory.STEP_NODE, True), ('cf_pred2', NodeCategory.STEP_NODE, True), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred1', 'target')
    graph.add_edge('cf_pred2', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_multiple_control_flow_predecessors', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_multiple_data_provider_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple data provider predecessors.'
    graph = create_test_graph_with_nodes([('data_pred1', NodeCategory.STEP_NODE, False), ('data_pred2', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('data_pred1', 'target')
    graph.add_edge('data_pred2', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_multiple_data_provider_predecessors', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_multiple_mixed_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple control flow and data provider predecessors mixed.'
    graph = create_test_graph_with_nodes([('cf_pred1', NodeCategory.STEP_NODE, True), ('cf_pred2', NodeCategory.STEP_NODE, True), ('data_pred1', NodeCategory.STEP_NODE, False), ('data_pred2', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred1', 'target')
    graph.add_edge('cf_pred2', 'target')
    graph.add_edge('data_pred1', 'target')
    graph.add_edge('data_pred2', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_multiple_mixed_predecessors', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_input_node_predecessor():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with an INPUT_NODE predecessor (should not be treated as control flow).'
    graph = create_test_graph_with_nodes([('input_node', NodeCategory.INPUT_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('input_node', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_input_node_predecessor', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_output_node_predecessor():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with an OUTPUT_NODE predecessor (should not be treated as control flow).'
    graph = create_test_graph_with_nodes([('output_node', NodeCategory.OUTPUT_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('output_node', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_output_node_predecessor', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_control_flow_step_that_doesnt_control_flow():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with STEP_NODE that has controls_flow=False.'
    graph = create_test_graph_with_nodes([('step', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('step', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_control_flow_step_that_doesnt_control_flow', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_return_type_is_tuple():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that return type is a tuple.'
    graph = create_test_graph_with_nodes([('node1', NodeCategory.STEP_NODE, False)])
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='node1')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_return_type_is_tuple', 'separate_control_flow_predecessors_from_data_providers', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_return_lists_are_lists():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that return values are lists, not other iterables.'
    graph = create_test_graph_with_nodes([('cf_pred', NodeCategory.STEP_NODE, True), ('node1', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred', 'node1')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='node1')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_return_lists_are_lists', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_node_with_self_reference():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a node that references itself (cycle).'
    graph = create_test_graph_with_nodes([('self_ref', NodeCategory.STEP_NODE, True)])
    graph.add_edge('self_ref', 'self_ref')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='self_ref')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_node_with_self_reference', 'separate_control_flow_predecessors_from_data_providers', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_very_long_node_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with very long node names.'
    long_name_cf = 'a' * 1000
    long_name_data = 'b' * 1000
    long_name_target = 'c' * 1000
    graph = create_test_graph_with_nodes([(long_name_cf, NodeCategory.STEP_NODE, True), (long_name_data, NodeCategory.STEP_NODE, False), (long_name_target, NodeCategory.STEP_NODE, False)])
    graph.add_edge(long_name_cf, long_name_target)
    graph.add_edge(long_name_data, long_name_target)
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node=long_name_target)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_very_long_node_names', 'separate_control_flow_predecessors_from_data_providers', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_node_names_with_special_characters():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with node names containing special characters.'
    graph = create_test_graph_with_nodes([('node-1:step[0]', NodeCategory.STEP_NODE, True), ('node-2:step[1]', NodeCategory.STEP_NODE, False), ('target_node', NodeCategory.STEP_NODE, False)])
    graph.add_edge('node-1:step[0]', 'target_node')
    graph.add_edge('node-2:step[1]', 'target_node')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target_node')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_node_names_with_special_characters', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_numeric_node_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with numeric node names.'
    graph = create_test_graph_with_nodes([(1, NodeCategory.STEP_NODE, True), (2, NodeCategory.STEP_NODE, False), (3, NodeCategory.STEP_NODE, False)])
    graph.add_edge(1, 3)
    graph.add_edge(2, 3)
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node=3)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_numeric_node_names', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_complex_graph_structure():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a more complex graph structure (diamond).'
    graph = create_test_graph_with_nodes([('cf_pred', NodeCategory.STEP_NODE, True), ('data_pred', NodeCategory.STEP_NODE, False), ('intermediate1', NodeCategory.STEP_NODE, False), ('intermediate2', NodeCategory.STEP_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('cf_pred', 'intermediate1')
    graph.add_edge('data_pred', 'intermediate2')
    graph.add_edge('intermediate1', 'target')
    graph.add_edge('intermediate2', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_complex_graph_structure', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_predecessor_order_preserved():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that predecessor order is preserved from graph.'
    graph = create_test_graph_with_nodes([('z_pred', NodeCategory.STEP_NODE, True), ('a_pred', NodeCategory.STEP_NODE, True), ('m_pred', NodeCategory.STEP_NODE, True), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('z_pred', 'target')
    graph.add_edge('a_pred', 'target')
    graph.add_edge('m_pred', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_predecessor_order_preserved', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_many_control_flow_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a large number of control flow predecessors.'
    nodes = [(f'cf_pred_{i}', NodeCategory.STEP_NODE, True) for i in range(500)]
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    for i in range(500):
        graph.add_edge(f'cf_pred_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_many_control_flow_predecessors', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_many_data_provider_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a large number of data provider predecessors.'
    nodes = [(f'data_pred_{i}', NodeCategory.STEP_NODE, False) for i in range(500)]
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    for i in range(500):
        graph.add_edge(f'data_pred_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_many_data_provider_predecessors', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_large_mixed_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with large numbers of both control flow and data provider predecessors.'
    nodes = [(f'cf_pred_{i}', NodeCategory.STEP_NODE, True) for i in range(250)]
    nodes += [(f'data_pred_{i}', NodeCategory.STEP_NODE, False) for i in range(250)]
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    for i in range(250):
        graph.add_edge(f'cf_pred_{i}', 'target')
        graph.add_edge(f'data_pred_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_large_mixed_predecessors', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_large_linear_chain():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a large linear chain of nodes.'
    nodes = [(i, NodeCategory.STEP_NODE, i % 2 == 0) for i in range(1000)]
    graph = create_test_graph_with_nodes(nodes)
    for i in range(999):
        graph.add_edge(i, i + 1)
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node=999)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_large_linear_chain', 'separate_control_flow_predecessors_from_data_providers', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_large_complete_graph_subset():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a large subset where all nodes connect to target.'
    nodes = [(f'node_{i}', NodeCategory.STEP_NODE, i % 3 == 0) for i in range(100)]
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    for i in range(100):
        graph.add_edge(f'node_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_large_complete_graph_subset', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value
    expected_cf_count = len([i for i in range(100) if i % 3 == 0])
    expected_data_count = 100 - expected_cf_count

def test_multiple_levels_deep():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that only immediate predecessors are considered (not transitive).'
    nodes = [(f'node_{i}', NodeCategory.STEP_NODE, True) for i in range(3)]
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    graph.add_edge('node_0', 'node_1')
    graph.add_edge('node_1', 'node_2')
    graph.add_edge('node_2', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_multiple_levels_deep', 'separate_control_flow_predecessors_from_data_providers', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_all_node_categories_as_predecessors():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with predecessors of all node categories.'
    graph = create_test_graph_with_nodes([('input_node', NodeCategory.INPUT_NODE, False), ('step_node_cf', NodeCategory.STEP_NODE, True), ('step_node_data', NodeCategory.STEP_NODE, False), ('output_node', NodeCategory.OUTPUT_NODE, False), ('target', NodeCategory.STEP_NODE, False)])
    graph.add_edge('input_node', 'target')
    graph.add_edge('step_node_cf', 'target')
    graph.add_edge('step_node_data', 'target')
    graph.add_edge('output_node', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_all_node_categories_as_predecessors', 'separate_control_flow_predecessors_from_data_providers', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value

def test_wide_graph_with_many_unrelated_nodes():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with a large graph where most nodes are unrelated to target.'
    all_nodes = [(f'node_{i}', NodeCategory.STEP_NODE, i % 5 == 0) for i in range(500)]
    all_nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(all_nodes)
    for i in range(50):
        graph.add_edge(f'node_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_wide_graph_with_many_unrelated_nodes', 'separate_control_flow_predecessors_from_data_providers', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value
    expected_cf = [i for i in range(50) if i % 5 == 0]

def test_alternating_control_flow_and_data_providers():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with alternating control flow and data provider pattern (100 of each).'
    nodes = []
    for i in range(0, 200, 2):
        nodes.append((f'node_{i}', NodeCategory.STEP_NODE, True))
    for i in range(1, 200, 2):
        nodes.append((f'node_{i}', NodeCategory.STEP_NODE, False))
    nodes.append(('target', NodeCategory.STEP_NODE, False))
    graph = create_test_graph_with_nodes(nodes)
    for i in range(200):
        graph.add_edge(f'node_{i}', 'target')
    _call__bound__arguments = inspect.signature(separate_control_flow_predecessors_from_data_providers).bind(execution_graph=graph, node='target')
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(separate_control_flow_predecessors_from_data_providers, 'tests.inference.unit_tests.test_separate_control_flow_predecessors_from_data_providers__unit_test_1', None, 'test_alternating_control_flow_and_data_providers', 'separate_control_flow_predecessors_from_data_providers', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    control_flow, data_providers = codeflash_return_value