from __future__ import annotations

import gc
import os
import time

import pytest
from inference.core.workflows.execution_engine.v1.executor.execution_data_manager.step_input_assembler import \
    get_masks_intersection_for_dimensions


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_empty_batch_masks_returns_none_for_each_dim():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    batch_masks = []
    dimensions = {1, 2, 3}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_empty_batch_masks_returns_none_for_each_dim', 'get_masks_intersection_for_dimensions', '2', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for dim in dimensions:
        pass

def test_single_dimension_intersection_ignores_empty_masks():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {(1,), (2,), (3,)}
    mask2 = set()
    mask3 = {(1,), (3,)}
    batch_masks = [mask1, mask2, mask3]
    dimensions = {1}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_single_dimension_intersection_ignores_empty_masks', 'get_masks_intersection_for_dimensions', '5', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_single_dimension_no_non_empty_masks_returns_empty_set():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {(1, 1), (2, 2)}
    mask2 = {(3, 3)}
    batch_masks = [mask1, mask2]
    dimensions = {1}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_single_dimension_no_non_empty_masks_returns_empty_set', 'get_masks_intersection_for_dimensions', '4', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_two_dimensions_basic_prefix_chain_filtering():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {(1,), (1, 1), (2,)}
    mask2 = {(1,), (1, 1), (2,)}
    mask3 = {(1,), (1, 1), (2,)}
    batch_masks = [mask1, mask2, mask3]
    dimensions = {1, 2}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_two_dimensions_basic_prefix_chain_filtering', 'get_masks_intersection_for_dimensions', '5', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_three_dimensions_partial_and_complete_chains():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask_common = {(1,), (2,), (1, 1), (2, 1), (1, 1, 1)}
    batch_masks = [set(mask_common) for _ in range(3)]
    dimensions = {1, 2, 3}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_three_dimensions_partial_and_complete_chains', 'get_masks_intersection_for_dimensions', '3', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_dimensions_returns_empty_mapping():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    batch_masks = [{(1,), (1, 1)}]
    dimensions = set()
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_empty_dimensions_returns_empty_mapping', 'get_masks_intersection_for_dimensions', '2', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_zero_length_indices_and_dim_zero_behavior():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    mask1 = {(), (1,), (1, 1)}
    mask2 = {(), (1,), (1, 1)}
    batch_masks = [mask1, mask2]
    dimensions = {0, 1, 2}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_zero_length_indices_and_dim_zero_behavior', 'get_masks_intersection_for_dimensions', '4', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_scale_many_chains_performance_and_correctness():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    K = 300
    chains = []
    for i in range(1, K + 1):
        chains.append((i,))
        chains.append((i, 1))
        chains.append((i, 1, 1))
    num_masks = 5
    batch_masks = [set(chains) for _ in range(num_masks)]
    dimensions = {1, 2, 3}
    codeflash_return_value = codeflash_wrap(get_masks_intersection_for_dimensions, 'tests.inference.unit_tests.test_get_masks_intersection_for_dimensions__unit_test_0', None, 'test_large_scale_many_chains_performance_and_correctness', 'get_masks_intersection_for_dimensions', '6', codeflash_loop_index, batch_masks, dimensions)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    sample_indices = [1, K // 2, K]
    for i in sample_indices:
        pass