from __future__ import annotations

import gc
import os
import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

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

def codeflash_wrap(
    codeflash_wrapped: Callable[..., Any],
    codeflash_test_module_name: str,
    codeflash_test_class_name: Optional[str],
    codeflash_test_name: str,
    codeflash_function_name: str,
    codeflash_line_id: str,
    codeflash_loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> str:
    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}"
    exception = None
    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}######$!")
    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

class ExecutionEngineRuntimeError(Exception):

    def __init__(self, public_message, context=None):
        super().__init__(public_message)
        self.public_message = public_message
        self.context = context


DynamicBatchIndex = Tuple[int, ...]


class Batch:

    def __init__(self, data: List[Any], indices: List[DynamicBatchIndex]):
        self.data = data
        self.indices = indices

    def iter_with_indices(self):
        return zip(self.indices, self.data)

    def remove_by_indices(self, indices_to_remove: Set[DynamicBatchIndex]):
        new_data = []
        new_indices = []
        for idx, dat in zip(self.indices, self.data):
            if idx not in indices_to_remove:
                new_data.append(dat)
                new_indices.append(idx)
        return Batch(new_data, new_indices)

    def __eq__(self, other):
        if not isinstance(other, Batch):
            return False
        return self.data == other.data and self.indices == other.indices

    def __repr__(self):
        return f"Batch(data={self.data}, indices={self.indices})"


@dataclass(frozen=True)
class ParameterSpecification:
    parameter_name: str
    nested_element_key: Optional[str] = None


class NodeInputCategory:
    pass


@dataclass(frozen=True)
class StepInputDefinition:
    parameter_specification: ParameterSpecification
    category: NodeInputCategory
    _batch_oriented: bool = False
    _dimensionality: int = 0
    _points_to_input: bool = True
    _points_to_step_output: bool = False
    _value: Any = None
    _selector: Optional[str] = None
    _data_lineage: Optional[List[int]] = None

    @classmethod
    def is_compound_input(cls) -> bool:
        return False

    def is_batch_oriented(self) -> bool:
        return self._batch_oriented

    def get_dimensionality(self) -> int:
        return self._dimensionality

    def points_to_input(self) -> bool:
        return self._points_to_input

    def points_to_step_output(self) -> bool:
        return self._points_to_step_output

    @property
    def selector(self):
        return self._selector

    @property
    def data_lineage(self):
        return self._data_lineage

    @property
    def value(self):
        return self._value


@dataclass(frozen=True)
class StaticStepInputDefinition(StepInputDefinition):
    pass


@dataclass(frozen=True)
class DynamicStepInputDefinition(StepInputDefinition):
    pass


@dataclass(frozen=True)
class CompoundStepInputDefinition:
    _definitions: Union[List[StepInputDefinition], Dict[str, StepInputDefinition]]
    _list: bool = False

    @classmethod
    def is_compound_input(cls) -> bool:
        return True

    def represents_list_of_inputs(self) -> bool:
        return self._list

    def iterate_through_definitions(self):
        if self._list:
            return iter(self._definitions)
        else:
            return iter(self._definitions.values())

    def get_dimensionality(self) -> int:
        dims = []
        for d in self.iterate_through_definitions():
            dims.append(d.get_dimensionality())
        dims.append(0)
        return max(dims)


@dataclass
class StepNode:
    name: str
    input_data: Dict[str, Any]
    execution_branches_impacting_inputs: List[str]
    step_execution_dimensionality: int
    step_manifest: Any
    auto_batch_casting_lineage_supports: Dict[str, Any]


class BatchModeSIMDStepInput:

    def __init__(self, indices, parameters):
        self.indices = indices
        self.parameters = parameters

    def __eq__(self, other):
        if not isinstance(other, BatchModeSIMDStepInput):
            return False
        return self.indices == other.indices and self.parameters == other.parameters

    def __repr__(self):
        return f"BatchModeSIMDStepInput(indices={self.indices}, parameters={self.parameters})"


class DummyManifest:

    @classmethod
    def accepts_batch_input(cls):
        return True

    @classmethod
    def accepts_empty_values(cls):
        return False


class DynamicBatchesManager:

    def __init__(
        self, lineage_to_indices: Dict[Tuple[int, ...], List[DynamicBatchIndex]]
    ):
        self.lineage_to_indices = lineage_to_indices

    def get_indices_for_data_lineage(self, lineage):
        return self.lineage_to_indices.get(tuple(lineage), [])


class ExecutionCache:

    def __init__(
        self, step_outputs: Dict[str, Any], batch_outputs: Dict[str, List[Any]]
    ):
        self.step_outputs = step_outputs
        self.batch_outputs = batch_outputs

    def get_non_batch_output(self, selector):
        return self.step_outputs.get(selector, None)

    def get_batch_output(self, selector, batch_elements_indices, mask):
        data = self.batch_outputs.get(selector, [])
        if mask is not None:
            return [
                val if idx in mask else None
                for idx, val in zip(batch_elements_indices, data)
            ]
        else:
            return [val for idx, val in zip(batch_elements_indices, data)]


class BranchingManager:

    def __init__(self, masks: Dict[str, Union[Set[DynamicBatchIndex], bool]]):
        self._masks = masks
        self._batch_compatibility = {
            branch_name: not isinstance(mask, bool)
            for branch_name, mask in masks.items()
        }

    def get_mask(self, execution_branch: str) -> Union[Set[DynamicBatchIndex], bool]:
        return self._masks[execution_branch]

    def is_execution_branch_batch_oriented(self, execution_branch: str) -> bool:
        return self._batch_compatibility[execution_branch]

    def is_execution_branch_registered(self, execution_branch: str) -> bool:
        return execution_branch in self._masks


def test_single_batch_input_basic():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a single batch input with one execution branch and straightforward mask.\n    "
    parameter = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _points_to_step_output=False,
        _selector="input1",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step1",
        input_data={"input1": parameter},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": [10, 20, 30]}
    masks = {"branch1": {(0,), (1,), (2,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,), (2,)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_single_batch_input_basic",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_batch = Batch([10, 20, 30], [(0,), (1,), (2,)])
    expected = BatchModeSIMDStepInput(
        indices=[(0,), (1,), (2,)], parameters={"input1": expected_batch}
    )


def test_multiple_inputs_basic():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test multiple batch inputs with matching indices and masks.\n    "
    param1 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input1",
        _data_lineage=[0],
    )
    param2 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input2"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input2",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step2",
        input_data={"input1": param1, "input2": param2},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": [1, 2], "input2": [3, 4]}
    masks = {"branch1": {(0,), (1,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_multiple_inputs_basic",
        "construct_simd_step_input",
        "9",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(
        indices=[(0,), (1,)],
        parameters={
            "input1": Batch([1, 2], [(0,), (1,)]),
            "input2": Batch([3, 4], [(0,), (1,)]),
        },
    )


def test_static_input_basic():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a static (non-batch) input.\n    "
    param = StaticStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="static1"),
        category=NodeInputCategory(),
        _batch_oriented=False,
        _dimensionality=0,
        _points_to_input=False,
        _points_to_step_output=False,
        _value=42,
    )
    step_node = StepNode(
        name="step_static",
        input_data={"static1": param},
        execution_branches_impacting_inputs=[],
        step_execution_dimensionality=0,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {}
    branching_manager = BranchingManager({})
    dynamic_batches_manager = DynamicBatchesManager({})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_static_input_basic",
        "construct_simd_step_input",
        "7",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(indices=[], parameters={"static1": 42})


def test_empty_input_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test with an empty input list, should result in empty indices and values.\n    "
    param = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input1",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step_empty",
        input_data={"input1": param},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": []}
    masks = {"branch1": set()}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): []})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_empty_input_list",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(indices=[], parameters={"input1": Batch([], [])})


def test_mask_removes_some_elements():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test where the mask removes some elements from the batch.\n    "
    param = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input1",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step_mask",
        input_data={"input1": param},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": [5, 6, 7, 8]}
    masks = {"branch1": {(1,), (3,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,), (2,), (3,)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_mask_removes_some_elements",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected_batch = Batch(
        [None, 6, None, 8], [(0,), (1,), (2,), (3,)]
    ).remove_by_indices({(0,), (2,)})
    expected = BatchModeSIMDStepInput(
        indices=[(1,), (3,)], parameters={"input1": expected_batch}
    )


def test_compound_input_dict():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a compound input as a dict of two batch parameters.\n    "
    param1 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(
            parameter_name="inputA", nested_element_key="A"
        ),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputA",
        _data_lineage=[0],
    )
    param2 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(
            parameter_name="inputB", nested_element_key="B"
        ),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputB",
        _data_lineage=[0],
    )
    compound = CompoundStepInputDefinition({"A": param1, "B": param2}, _list=False)
    step_node = StepNode(
        name="step_compound",
        input_data={"compound": compound},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"inputA": [1, 2], "inputB": [10, 20]}
    masks = {"branch1": {(0,), (1,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_compound_input_dict",
        "construct_simd_step_input",
        "10",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(
        indices=[(0,), (1,)],
        parameters={
            "compound": {
                "A": Batch([1, 2], [(0,), (1,)]),
                "B": Batch([10, 20], [(0,), (1,)]),
            }
        },
    )


def test_compound_input_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a compound input as a list of batch parameters.\n    "
    param1 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="inputA"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputA",
        _data_lineage=[0],
    )
    param2 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="inputB"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputB",
        _data_lineage=[0],
    )
    compound = CompoundStepInputDefinition([param1, param2], _list=True)
    step_node = StepNode(
        name="step_compound_list",
        input_data={"compound": compound},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"inputA": [1, 2], "inputB": [10, 20]}
    masks = {"branch1": {(0,), (1,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_compound_input_list",
        "construct_simd_step_input",
        "10",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(
        indices=[(0,), (1,)],
        parameters={
            "compound": [Batch([1, 2], [(0,), (1,)]), Batch([10, 20], [(0,), (1,)])]
        },
    )


def test_mismatched_indices_in_compound_raises():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test that mismatched indices in compound input raises an error.\n    "
    param1 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="inputA"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputA",
        _data_lineage=[0],
    )
    param2 = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="inputB"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="inputB",
        _data_lineage=[0],
    )
    compound = CompoundStepInputDefinition([param1, param2], _list=True)
    step_node = StepNode(
        name="step_compound_list",
        input_data={"compound": compound},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"inputA": [1, 2], "inputB": [10, 20, 30]}
    masks = {"branch1": {(0,), (1,), (2,)}}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(0,), (1,), (2,)]})
    execution_cache = ExecutionCache({}, {})
    with pytest.raises(ExecutionEngineRuntimeError):
        codeflash_return_value = codeflash_wrap(
            construct_simd_step_input,
            "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
            None,
            "test_mismatched_indices_in_compound_raises",
            "construct_simd_step_input",
            "10_0",
            codeflash_loop_index,
            step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
        )


def test_no_batch_inputs_raises():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test that if there are no batch inputs, an error is raised.\n    "
    param = StaticStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="static1"),
        category=NodeInputCategory(),
        _batch_oriented=False,
        _dimensionality=0,
        _points_to_input=False,
        _points_to_step_output=False,
        _value=123,
    )
    step_node = StepNode(
        name="step_no_batch",
        input_data={"static1": param},
        execution_branches_impacting_inputs=[],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {}
    branching_manager = BranchingManager({})
    dynamic_batches_manager = DynamicBatchesManager({})
    execution_cache = ExecutionCache({}, {})
    with pytest.raises(ExecutionEngineRuntimeError):
        codeflash_return_value = codeflash_wrap(
            construct_simd_step_input,
            "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
            None,
            "test_no_batch_inputs_raises",
            "construct_simd_step_input",
            "7_0",
            codeflash_loop_index,
            step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
        )


def test_large_batch_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test with a large batch input (size 1000).\n    "
    N = 1000
    param = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input1",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step_large",
        input_data={"input1": param},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": list(range(N))}
    masks = {"branch1": set(((i,) for i in range(N)))}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(i,) for i in range(N)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_large_batch_input",
        "construct_simd_step_input",
        "9",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(
        indices=[(i,) for i in range(N)],
        parameters={"input1": Batch(list(range(N)), [(i,) for i in range(N)])},
    )


def test_large_compound_input_dict():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a large compound input dict with 10 keys, each with a batch of size 100.\n    "
    N = 100
    keys = [f"k{i}" for i in range(10)]
    params = {
        k: DynamicStepInputDefinition(
            parameter_specification=ParameterSpecification(
                parameter_name=k, nested_element_key=k
            ),
            category=NodeInputCategory(),
            _batch_oriented=True,
            _dimensionality=1,
            _points_to_input=True,
            _selector=k,
            _data_lineage=[0],
        )
        for k in keys
    }
    compound = CompoundStepInputDefinition(params, _list=False)
    step_node = StepNode(
        name="step_large_compound",
        input_data={"compound": compound},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {k: list(range(N)) for k in keys}
    masks = {"branch1": set(((i,) for i in range(N)))}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(i,) for i in range(N)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_large_compound_input_dict",
        "construct_simd_step_input",
        "11",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    expected = BatchModeSIMDStepInput(
        indices=[(i,) for i in range(N)],
        parameters={
            "compound": {
                k: Batch(list(range(N)), [(i,) for i in range(N)]) for k in keys
            }
        },
    )


def test_large_mask_removal():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    "\n    Test a large batch input with a mask that removes half the elements.\n    "
    N = 1000
    param = DynamicStepInputDefinition(
        parameter_specification=ParameterSpecification(parameter_name="input1"),
        category=NodeInputCategory(),
        _batch_oriented=True,
        _dimensionality=1,
        _points_to_input=True,
        _selector="input1",
        _data_lineage=[0],
    )
    step_node = StepNode(
        name="step_large_mask",
        input_data={"input1": param},
        execution_branches_impacting_inputs=["branch1"],
        step_execution_dimensionality=1,
        step_manifest=DummyManifest,
        auto_batch_casting_lineage_supports={},
    )
    runtime_parameters = {"input1": list(range(N))}
    mask_indices = set(((i,) for i in range(0, N, 2)))
    masks = {"branch1": mask_indices}
    branching_manager = BranchingManager(masks)
    dynamic_batches_manager = DynamicBatchesManager({(0,): [(i,) for i in range(N)]})
    execution_cache = ExecutionCache({}, {})
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_0",
        None,
        "test_large_mask_removal",
        "construct_simd_step_input",
        "10",
        codeflash_loop_index,
        step_node,runtime_parameters,execution_cache,dynamic_batches_manager,branching_manager,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    kept_indices = [(i,) for i in range(0, N, 2)]
    kept_values = [i for i in range(0, N, 2)]
    expected_batch = Batch(
        [None if i % 2 == 1 else i for i in range(N)], [(i,) for i in range(N)]
    ).remove_by_indices(set(((i,) for i in range(1, N, 2))))
    expected = BatchModeSIMDStepInput(
        indices=kept_indices, parameters={"input1": expected_batch}
    )
