from __future__ import annotations

import gc
import inspect
import os
import time
from collections import defaultdict
from copy import copy
from enum import Enum
from pathlib import Path
from typing import Dict, List, Set, Tuple

import dill as pickle
import networkx as nx
import pytest
from inference.core.workflows.errors import ControlFlowDefinitionError
from inference.core.workflows.execution_engine.constants import \
    NODE_COMPILATION_OUTPUT_PROPERTY
from inference.core.workflows.execution_engine.v1.compiler.entities import \
    StepNode
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import (
    _collect_unique_control_flow_lineages_with_step_mapping,
    get_all_batch_lineage_prefixes, get_batch_lineage_prefixes,
    verify_compatibility_of_input_data_lineage_with_control_flow_lineage)
from inference.core.workflows.execution_engine.v1.compiler.utils import \
    identify_lineage
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 NodeCategory(str, Enum):
    STEP = 'step'
    CONTROL_FLOW = 'control_flow'

def test_verify_compatibility_empty_inputs_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns None when inputs_lineage is empty (no verification needed).'
    graph = nx.DiGraph()
    step_node = StepNode(node_category=NodeCategory.STEP, name='test_step', selector='step_1', data_lineage=['step_0'])
    graph.add_node('step_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: step_node})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='test_step', inputs_lineage=[], control_flow_steps_selectors=['step_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_empty_inputs_lineage', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_empty_control_flow_steps():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns None when there are no control flow steps.'
    graph = nx.DiGraph()
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='test_step', inputs_lineage=[['input_step']], control_flow_steps_selectors=[], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_empty_control_flow_steps', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_matching_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that function returns None when control flow lineage matches input lineage.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['input_a'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_matching_lineages', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_incompatible_lineage_raises_error():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test that function raises ControlFlowDefinitionError when lineages don't match."
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['incompatible_data'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    with pytest.raises(ControlFlowDefinitionError) as exc_info:
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_incompatible_lineage_raises_error', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_verify_compatibility_prefix_match():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that control flow lineage prefix matches input lineage.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['input_a'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a', 'input_b']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_prefix_match', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_collect_unique_control_flow_lineages_single_step():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test collecting lineages from a single control flow step.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['data_a'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    unique_lineages, lineage_mapping = _collect_unique_control_flow_lineages_with_step_mapping(control_flow_steps_selectors=['control_1'], execution_graph=graph)
    lineage_id = identify_lineage(['data_a'])

def test_collect_unique_control_flow_lineages_duplicate_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that duplicate lineages are deduplicated.'
    graph = nx.DiGraph()
    step1 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step_1', selector='control_1', data_lineage=['data_a'])
    step2 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step_2', selector='control_2', data_lineage=['data_a'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: step1})
    graph.add_node('control_2', **{NODE_COMPILATION_OUTPUT_PROPERTY: step2})
    unique_lineages, lineage_mapping = _collect_unique_control_flow_lineages_with_step_mapping(control_flow_steps_selectors=['control_1', 'control_2'], execution_graph=graph)
    lineage_id = identify_lineage(['data_a'])

def test_get_batch_lineage_prefixes_simple():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating prefixes from a simple lineage.'
    lineage = ['a', 'b', 'c']
    prefixes = get_batch_lineage_prefixes(lineage=lineage)

def test_get_batch_lineage_prefixes_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating prefixes from an empty lineage.'
    lineage = []
    prefixes = get_batch_lineage_prefixes(lineage=lineage)

def test_get_batch_lineage_prefixes_single_element():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating prefixes from a single element lineage.'
    lineage = ['a']
    prefixes = get_batch_lineage_prefixes(lineage=lineage)

def test_get_all_batch_lineage_prefixes_multiple_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating all prefixes from multiple lineages.'
    lineages = [['a', 'b'], ['c']]
    all_prefixes = get_all_batch_lineage_prefixes(lineages=lineages)
    prefix_lists = [list(p) for p in all_prefixes]

def test_get_all_batch_lineage_prefixes_deduplication():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that duplicate prefixes are deduplicated.'
    lineages = [['a', 'b'], ['a', 'c']]
    all_prefixes = get_all_batch_lineage_prefixes(lineages=lineages)

def test_verify_compatibility_control_flow_empty_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with control flow step having empty data_lineage.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=[])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_control_flow_empty_lineage', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_multiple_control_flow_steps_all_compatible():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple control flow steps, all compatible.'
    graph = nx.DiGraph()
    step1 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_1', selector='c1', data_lineage=['input_a'])
    step2 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_2', selector='c2', data_lineage=['input_b'])
    graph.add_node('c1', **{NODE_COMPILATION_OUTPUT_PROPERTY: step1})
    graph.add_node('c2', **{NODE_COMPILATION_OUTPUT_PROPERTY: step2})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a'], ['input_b']], control_flow_steps_selectors=['c1', 'c2'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_multiple_control_flow_steps_all_compatible', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '6', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_multiple_control_flow_steps_one_incompatible():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multiple control flow steps, one incompatible.'
    graph = nx.DiGraph()
    step1 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_1', selector='c1', data_lineage=['input_a'])
    step2 = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_2', selector='c2', data_lineage=['incompatible'])
    graph.add_node('c1', **{NODE_COMPILATION_OUTPUT_PROPERTY: step1})
    graph.add_node('c2', **{NODE_COMPILATION_OUTPUT_PROPERTY: step2})
    with pytest.raises(ControlFlowDefinitionError) as exc_info:
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a']], control_flow_steps_selectors=['c1', 'c2'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_multiple_control_flow_steps_one_incompatible', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '6_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_verify_compatibility_nested_lineage_prefix():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with deeply nested lineage hierarchies.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['level1', 'level2'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['level1', 'level2', 'level3', 'level4']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_nested_lineage_prefix', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_special_characters_in_step_names():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with special characters in step names and lineage.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step_@#$', selector='control_1', data_lineage=['data_@_special', 'data_#_chars'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step_@#$', inputs_lineage=[['data_@_special', 'data_#_chars']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_special_characters_in_step_names', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_batch_lineage_prefixes_long_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with long lineage (100 elements).'
    lineage = [f'element_{i}' for i in range(100)]
    prefixes = get_batch_lineage_prefixes(lineage=lineage)

def test_collect_unique_lineages_empty_selector_list():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test collecting lineages with empty selector list.'
    graph = nx.DiGraph()
    unique_lineages, lineage_mapping = _collect_unique_control_flow_lineages_with_step_mapping(control_flow_steps_selectors=[], execution_graph=graph)

def test_verify_compatibility_control_flow_lineage_longer_than_input():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when control flow lineage is longer than any input prefix.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['a', 'b', 'c'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    with pytest.raises(ControlFlowDefinitionError):
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_control_flow_lineage_longer_than_input', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_verify_compatibility_similar_but_different_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with lineages that are similar but not matching.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['input_a'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    with pytest.raises(ControlFlowDefinitionError):
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_b']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_similar_but_different_lineage', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_error_message_contains_step_name():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that error message contains the step name for clarity.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['incompatible'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    step_name = 'my_unique_step_name_xyz'
    with pytest.raises(ControlFlowDefinitionError) as exc_info:
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name=step_name, inputs_lineage=[['input_a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_error_message_contains_step_name', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_error_message_contains_control_flow_step_selector():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that error message contains control flow step selector.'
    graph = nx.DiGraph()
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=['incompatible'])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    with pytest.raises(ControlFlowDefinitionError) as exc_info:
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[['input_a']], control_flow_steps_selectors=['control_1'], execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_error_message_contains_control_flow_step_selector', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '4_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_verify_compatibility_many_control_flow_steps():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 100 control flow steps, all compatible.'
    graph = nx.DiGraph()
    control_flow_selectors = []
    for i in range(100):
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'control_step_{i}', selector=f'c_{i}', data_lineage=[f'input_{i % 10}'])
        graph.add_node(f'c_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        control_flow_selectors.append(f'c_{i}')
    input_lineage = [[f'input_{i}'] for i in range(10)]
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=input_lineage, control_flow_steps_selectors=control_flow_selectors, execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_many_control_flow_steps', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_deep_lineage_hierarchy():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with deeply nested lineage (100 levels deep).'
    graph = nx.DiGraph()
    deep_lineage = [f'level_{i}' for i in range(100)]
    control_flow_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='control_step', selector='control_1', data_lineage=deep_lineage[:50])
    graph.add_node('control_1', **{NODE_COMPILATION_OUTPUT_PROPERTY: control_flow_step})
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[deep_lineage], control_flow_steps_selectors=['control_1'], execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_deep_lineage_hierarchy', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_collect_lineages_many_duplicate_steps():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test collecting lineages from 500 steps with many duplicates.'
    graph = nx.DiGraph()
    selectors = []
    for i in range(500):
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'step_{i}', selector=f's_{i}', data_lineage=[f'lineage_{i % 10}'])
        graph.add_node(f's_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        selectors.append(f's_{i}')
    unique_lineages, lineage_mapping = _collect_unique_control_flow_lineages_with_step_mapping(control_flow_steps_selectors=selectors, execution_graph=graph)
    for lineage_id, mapped_selectors in lineage_mapping.items():
        pass

def test_get_all_batch_lineage_prefixes_large_dataset():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating prefixes from 100 lineages with varying lengths.'
    lineages = []
    for i in range(100):
        length = i % 50 + 1
        lineage = [f'step_{j}' for j in range(length)]
        lineages.append(lineage)
    all_prefixes = get_all_batch_lineage_prefixes(lineages=lineages)

def test_verify_compatibility_many_input_lineages():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 100 different input lineages.'
    graph = nx.DiGraph()
    control_flow_selectors = []
    for i in range(50):
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'control_{i}', selector=f'c_{i}', data_lineage=[f'data_{i}'])
        graph.add_node(f'c_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        control_flow_selectors.append(f'c_{i}')
    input_lineages = [[f'data_{i % 50}'] for i in range(100)]
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=input_lineages, control_flow_steps_selectors=control_flow_selectors, execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_many_input_lineages', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_verify_compatibility_complex_prefix_matching_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test complex prefix matching with 50 control flow steps and 50 input lineages.'
    graph = nx.DiGraph()
    control_flow_selectors = []
    for i in range(50):
        lineage = [f'level_{j}' for j in range(i + 1)]
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'control_{i}', selector=f'c_{i}', data_lineage=lineage)
        graph.add_node(f'c_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        control_flow_selectors.append(f'c_{i}')
    input_lineage = [f'level_{i}' for i in range(50)]
    _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=[input_lineage], control_flow_steps_selectors=control_flow_selectors, execution_graph=graph)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_complex_prefix_matching_large_scale', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_get_batch_lineage_prefixes_very_long_lineage():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test generating prefixes from a 500-element lineage.'
    lineage = [f'element_{i}' for i in range(500)]
    prefixes = get_batch_lineage_prefixes(lineage=lineage)

def test_collect_lineages_mixed_empty_and_non_empty():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test collecting lineages with mix of empty and non-empty from 100 steps.'
    graph = nx.DiGraph()
    selectors = []
    for i in range(100):
        if i % 2 == 0:
            lineage = [f'data_{i}']
        else:
            lineage = []
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'step_{i}', selector=f's_{i}', data_lineage=lineage)
        graph.add_node(f's_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        selectors.append(f's_{i}')
    unique_lineages, lineage_mapping = _collect_unique_control_flow_lineages_with_step_mapping(control_flow_steps_selectors=selectors, execution_graph=graph)

def test_verify_compatibility_stress_test_many_incompatible():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test error handling with many steps where one is incompatible.'
    graph = nx.DiGraph()
    control_flow_selectors = []
    for i in range(200):
        step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name=f'control_{i}', selector=f'c_{i}', data_lineage=[f'input_{i % 10}'])
        graph.add_node(f'c_{i}', **{NODE_COMPILATION_OUTPUT_PROPERTY: step})
        control_flow_selectors.append(f'c_{i}')
    incompatible_step = StepNode(node_category=NodeCategory.CONTROL_FLOW, name='incompatible_step', selector='c_incompatible', data_lineage=['incompatible_data'])
    graph.add_node('c_incompatible', **{NODE_COMPILATION_OUTPUT_PROPERTY: incompatible_step})
    control_flow_selectors.append('c_incompatible')
    input_lineage = [[f'input_{i}'] for i in range(10)]
    with pytest.raises(ControlFlowDefinitionError) as exc_info:
        _call__bound__arguments = inspect.signature(verify_compatibility_of_input_data_lineage_with_control_flow_lineage).bind(step_name='main_step', inputs_lineage=input_lineage, control_flow_steps_selectors=control_flow_selectors, execution_graph=graph)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(verify_compatibility_of_input_data_lineage_with_control_flow_lineage, 'tests.inference.unit_tests.test_verify_compatibility_of_input_data_lineage_with_control_flow_lineage__unit_test_1', None, 'test_verify_compatibility_stress_test_many_incompatible', 'verify_compatibility_of_input_data_lineage_with_control_flow_lineage', '8_0', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)

def test_identify_lineage_hash_consistency_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that identify_lineage produces consistent hashes for large-scale data.'
    lineages = [[f'step_{i}', f'step_{j}'] for i in range(50) for j in range(20)]
    hashes_first = {identify_lineage(lineage) for lineage in lineages}
    hashes_second = {identify_lineage(lineage) for lineage in lineages}