from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import Generic, Iterator, List, Optional, Set, Tuple, TypeVar

import dill as pickle
import pytest
from inference.core.workflows.execution_engine.entities.base import Batch


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

def test_remove_single_index_from_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing a single index from a batch with multiple elements.'
    batch = Batch(content=['a', 'b', 'c'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(1,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_single_index_from_batch', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_multiple_indices_from_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing multiple indices from a batch.'
    batch = Batch(content=[1, 2, 3, 4, 5], indices=[(0,), (1,), (2,), (3,), (4,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,), (4,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_multiple_indices_from_batch', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_no_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that removing no indices returns all original elements.'
    batch = Batch(content=['x', 'y', 'z', 'w'], indices=[(0,), (1,), (2,), (3,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=set())
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_no_indices', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_first_index():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing the first element from a batch.'
    batch = Batch(content=['first', 'second', 'third'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_first_index', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_last_index():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing the last element from a batch.'
    batch = Batch(content=[10, 20, 30], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(2,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_last_index', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_integer_content_elements():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with integer content elements.'
    batch = Batch(content=[100, 200, 300, 400], indices=[(0,), (1,), (2,), (3,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(1,), (3,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_integer_content_elements', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_nested_batch_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing indices from a batch containing nested batches.'
    inner_batch = Batch(content=['x', 'y', 'z'], indices=[(0,), (1,), (2,)])
    outer_batch = Batch(content=['a', inner_batch, 'b'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(outer_batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(outer_batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_nested_batch_removal', 'remove_by_indices', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_tuple_indices_with_multiple_dimensions():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with multi-dimensional tuple indices.'
    batch = Batch(content=[1, 2, 3, 4], indices=[(0, 0), (0, 1), (1, 0), (1, 1)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0, 0), (1, 1)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_tuple_indices_with_multiple_dimensions', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_all_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing all elements from a batch.'
    batch = Batch(content=['a', 'b', 'c'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,), (1,), (2,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_all_indices', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_from_single_element_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing the only element from a batch.'
    batch = Batch(content=['only'], indices=[(0,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_from_single_element_batch', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_nonexistent_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    "Test removing indices that don't exist in the batch."
    batch = Batch(content=['a', 'b', 'c'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(5,), (10,), (99,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_nonexistent_indices', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_empty_batch_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing indices from an already empty batch.'
    batch = Batch(content=[], indices=[])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_empty_batch_removal', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_with_none_values_in_content():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test batch containing None values.'
    batch = Batch(content=['a', None, 'c'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(1,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_with_none_values_in_content', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_with_mixed_type_content():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test batch with mixed data types.'
    batch = Batch(content=[1, 'string', 3.14, {'key': 'value'}, [1, 2, 3]], indices=[(0,), (1,), (2,), (3,), (4,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,), (3,), (4,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_with_mixed_type_content', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_remove_preserves_order():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that removal preserves the order of remaining elements.'
    batch = Batch(content=['z', 'y', 'x', 'w', 'v'], indices=[(0,), (1,), (2,), (3,), (4,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(1,), (3,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_remove_preserves_order', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_batch_with_three_dimensional_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test with 3D tuple indices.'
    batch = Batch(content=[1, 2, 3, 4, 5], indices=[(0, 0, 0), (0, 0, 1), (0, 1, 0), (1, 0, 0), (1, 1, 1)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0, 0, 0), (1, 1, 1)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_batch_with_three_dimensional_indices', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_duplicate_removal_specification():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test when indices_to_remove contains the same index multiple times.'
    batch = Batch(content=[1, 2, 3, 4], indices=[(0,), (1,), (2,), (3,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(1,), (1,), (2,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_duplicate_removal_specification', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_single_element_batch():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test operations on a batch with a single large element.'
    large_list = list(range(1000))
    batch = Batch(content=[large_list], indices=[(0,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_large_single_element_batch', 'remove_by_indices', '3', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_batch_with_zero_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test batch with index (0,) at various positions.'
    batch = Batch(content=['first', 'second', 'third'], indices=[(0,), (1,), (2,)])
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove={(0,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_batch_with_zero_indices', 'remove_by_indices', '2', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_batch_partial_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing elements from a large batch (500 elements).'
    large_content = list(range(500))
    large_indices = [(i,) for i in range(500)]
    batch = Batch(content=large_content, indices=large_indices)
    indices_to_remove = {(i,) for i in range(0, 500, 10)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_large_batch_partial_removal', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    content_list = list(result)

def test_large_batch_almost_complete_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing almost all elements from a large batch.'
    large_content = list(range(1000))
    large_indices = [(i,) for i in range(1000)]
    batch = Batch(content=large_content, indices=large_indices)
    indices_to_remove = {(i,) for i in range(5, 1000)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_large_batch_almost_complete_removal', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_batch_with_large_removal_set():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with a large set of indices to remove.'
    content = list(range(500))
    indices = [(i,) for i in range(500)]
    batch = Batch(content=content, indices=indices)
    indices_to_remove = {(i,) for i in range(250)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_large_batch_with_large_removal_set', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_large_batch_all_odd_indices_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test removing all odd indices from a large batch.'
    large_content = list(range(1000))
    large_indices = [(i,) for i in range(1000)]
    batch = Batch(content=large_content, indices=large_indices)
    indices_to_remove = {(i,) for i in range(1000) if i % 2 == 1}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_large_batch_all_odd_indices_removal', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    content_list = list(result)

def test_nested_batch_large_scale():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test nested batch removal with large internal batches.'
    inner_content = list(range(100))
    inner_indices = [(0, i) for i in range(100)]
    inner_batch = Batch(content=inner_content, indices=inner_indices)
    outer_content = [inner_batch, 'outer1', 'outer2']
    outer_indices = [(0,), (1,), (2,)]
    outer_batch = Batch(content=outer_content, indices=outer_indices)
    _call__bound__arguments = inspect.signature(outer_batch.remove_by_indices).bind(indices_to_remove={(1,)})
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(outer_batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_nested_batch_large_scale', 'remove_by_indices', '7', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_batch_with_many_multidimensional_indices():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test large batch with multi-dimensional indices.'
    content = list(range(100))
    indices = [(i // 25, i // 5 % 5, i % 5) for i in range(100)]
    batch = Batch(content=content, indices=indices)
    indices_to_remove = {idx for idx in indices if idx[0] == 0}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_batch_with_many_multidimensional_indices', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    content_list = list(result)

def test_batch_iteration_after_large_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that iteration works correctly after removing from a large batch.'
    content = list(range(500))
    indices = [(i,) for i in range(500)]
    batch = Batch(content=content, indices=indices)
    indices_to_remove = {(i,) for i in range(250)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_batch_iteration_after_large_removal', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output
    iterated_list = list(result)

def test_batch_indexing_after_large_removal():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test that indexing works correctly after removing from a large batch.'
    content = list(range(100))
    indices = [(i,) for i in range(100)]
    batch = Batch(content=content, indices=indices)
    indices_to_remove = {(i,) for i in range(0, 100, 2)}
    _call__bound__arguments = inspect.signature(batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_batch_indexing_after_large_removal', 'remove_by_indices', '5', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
    codeflash_output = codeflash_return_value
    result = codeflash_output

def test_many_sequential_removals():
    codeflash_loop_index = int(os.environ['CODEFLASH_LOOP_INDEX'])
    'Test performance with many sequential removal operations.'
    content = list(range(200))
    indices = [(i,) for i in range(200)]
    batch = Batch(content=content, indices=indices)
    current_batch = batch
    for removal_round in range(10):
        indices_to_remove = {(i,) for i in range(removal_round * 10, (removal_round + 1) * 10)}
        _call__bound__arguments = inspect.signature(current_batch.remove_by_indices).bind(indices_to_remove=indices_to_remove)
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(current_batch.remove_by_indices, 'tests.inference.unit_tests.test_remove_by_indices__unit_test_1', None, 'test_many_sequential_removals', 'remove_by_indices', '5_1', codeflash_loop_index, *_call__bound__arguments.args, **_call__bound__arguments.kwargs)
        codeflash_output = codeflash_return_value
        current_batch = codeflash_output