from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.workflows.execution_engine.v1.compiler.graph_constructor import \
    establish_step_execution_dimensionality


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

def test_returns_zero_when_no_inputs_and_no_controls():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {}
    controls = []
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_returns_zero_when_no_inputs_and_no_controls', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_returns_control_flow_length_when_no_nonzero_inputs_but_control_flow_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'a': {0, -1}, 'b': {0}}
    controls = ['branch1', 'branch2', 'branch3']
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_returns_control_flow_length_when_no_nonzero_inputs_but_control_flow_present', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_uses_min_non_zero_dimensionality():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'inp1': {3, 5}, 'inp2': {2}, 'inp3': {4, 2}}
    controls = []
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_uses_min_non_zero_dimensionality', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_applies_negative_output_offset():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'a': {4}, 'b': {2, 3}}
    controls = []
    offset = -1
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_applies_negative_output_offset', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_zero_result_when_min_equals_one_and_offset_negative():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'only': {1}}
    controls = []
    offset = -1
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_zero_result_when_min_equals_one_and_offset_negative', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_ignores_zero_and_negative_input_dimensionalities():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'a': {0, -2}, 'b': {0}, 'c': {5}}
    controls = ['cf']
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_ignores_zero_and_negative_input_dimensionalities', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_handles_duplicate_values_in_sets():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'a': {3, 3, 2}, 'b': {2, 2, 4}, 'c': {2}}
    controls = []
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_handles_duplicate_values_in_sets', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_inputs_with_positive_dimensionalities():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {}
    for i in range(1000):
        val = i % 10 + 1
        inputs[f'inp_{i}'] = {val}
    controls = []
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_large_scale_many_inputs_with_positive_dimensionalities', 'establish_step_execution_dimensionality', '4', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_inputs_without_positive_but_control_flow_present():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {f'inp_{i}': {0} for i in range(1000)}
    controls = [f'branch_{j}' for j in range(500)]
    offset = 0
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_large_scale_many_inputs_without_positive_but_control_flow_present', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_negative_offset_produces_negative_result():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    inputs = {'a': {2}, 'b': {3}}
    controls = []
    offset = -5
    codeflash_return_value = codeflash_wrap(establish_step_execution_dimensionality, 'tests.inference.unit_tests.test_establish_step_execution_dimensionality__unit_test_0', None, 'test_large_negative_offset_produces_negative_result', 'establish_step_execution_dimensionality', '3', codeflash_loop_index, inputs, controls, offset)
    codeflash_output = codeflash_return_value
    result = codeflash_output