from __future__ import annotations

import gc
import inspect
import os
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union

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


def codeflash_wrap(
    wrapped: Callable[..., Any],
    test_module_name: str,
    test_class_name: str | None,
    test_name: str,
    function_name: str,
    line_id: str,
    loop_index: int,
    *args: Any,
    **kwargs: Any,
) -> Any:
    test_id = f"{test_module_name}:{test_class_name}:{test_name}:{line_id}:{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"{line_id}_{codeflash_test_index}"
    test_stdout_tag = f"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{loop_index}:{invocation_id}"
    print(f"!$######{test_stdout_tag}######$!")
    exception = None
    gc.disable()
    try:
        counter = time.perf_counter_ns()
        return_value = 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"{test_module_name}:{(test_class_name + '.' if test_class_name else '')}{test_name}:{function_name}:{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(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 ExecutionEngineRuntimeError(Exception):

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


class Batch:

    def __init__(self, data, indices):
        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):
        new_data = []
        new_indices = []
        for d, idx in zip(self.data, self.indices):
            if idx not in indices_to_remove:
                new_data.append(d)
                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})"


class ParameterSpecification:

    def __init__(self, parameter_name, nested_element_key=None):
        self.parameter_name = parameter_name
        self.nested_element_key = nested_element_key


class NodeInputCategory:
    pass


class StepInputDefinition:

    def __init__(
        self,
        parameter_specification,
        category,
        is_batch=False,
        points_to_input=True,
        points_to_step_output=False,
        value=None,
        selector=None,
        dimensionality=1,
        data_lineage=None,
    ):
        self.parameter_specification = parameter_specification
        self.category = category
        self._is_batch = is_batch
        self._points_to_input = points_to_input
        self._points_to_step_output = points_to_step_output
        self.value = value
        self.selector = selector or parameter_specification.parameter_name
        self._dimensionality = dimensionality
        self.data_lineage = data_lineage or tuple([0] * dimensionality)

    def is_compound_input(self):
        return False

    def get_dimensionality(self):
        return self._dimensionality

    def is_batch_oriented(self):
        return self._is_batch

    def points_to_input(self):
        return self._points_to_input

    def points_to_step_output(self):
        return self._points_to_step_output


class CompoundStepInputDefinition:

    def __init__(self, definitions, is_list=False):
        self._definitions = definitions
        self._is_list = is_list

    def is_compound_input(self):
        return True

    def iterate_through_definitions(self):
        return self._definitions

    def represents_list_of_inputs(self):
        return self._is_list

    def get_dimensionality(self):
        if not self._definitions:
            return 0
        return max([d.get_dimensionality() for d in self._definitions])


class StepNode:

    def __init__(
        self,
        name,
        input_data,
        step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=None,
        auto_batch_casting_lineage_supports=None,
    ):
        self.name = name
        self.input_data = input_data
        self.step_manifest = step_manifest
        self.step_execution_dimensionality = step_execution_dimensionality
        self.execution_branches_impacting_inputs = (
            execution_branches_impacting_inputs or []
        )
        self.auto_batch_casting_lineage_supports = (
            auto_batch_casting_lineage_supports or {}
        )


class WorkflowBlockManifest:

    def __init__(self, accepts_batch=True, accepts_empty=False):
        self._accepts_batch = accepts_batch
        self._accepts_empty = accepts_empty

    def accepts_batch_input(self):
        return self._accepts_batch

    def accepts_empty_values(self):
        return self._accepts_empty


class ExecutionCache:

    def __init__(self, batch_outputs=None, non_batch_outputs=None):
        self.batch_outputs = batch_outputs or {}
        self.non_batch_outputs = non_batch_outputs or {}

    def get_batch_output(self, selector, batch_elements_indices, mask=None):
        outputs = self.batch_outputs.get(selector, [])
        if mask is not None:
            return [
                out if idx in mask else None
                for out, idx in zip(outputs, batch_elements_indices)
            ]
        return [outputs[i] for i in range(len(batch_elements_indices))]

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


class DynamicBatchesManager:

    def __init__(self, lineage_to_indices):
        self.lineage_to_indices = lineage_to_indices

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


class BranchingManager:

    def __init__(self, masks):
        self._masks = masks
        self._batch_compatibility = {
            k: not isinstance(v, bool) for k, v in masks.items()
        }

    def get_mask(self, execution_branch):
        return self._masks[execution_branch]

    def is_execution_branch_batch_oriented(self, execution_branch):
        return self._batch_compatibility[execution_branch]

    def is_execution_branch_registered(self, execution_branch):
        return execution_branch in self._masks


def test_single_batch_input_simple():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input_param = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input_param},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": ["a", "b", "c"]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1, 2]})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_single_batch_input_simple",
        "construct_simd_step_input",
        "7",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_multiple_batch_inputs_same_indices():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    input2 = StepInputDefinition(
        ParameterSpecification("input2"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1, "input2": input2},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": [1, 2], "input2": [3, 4]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1]})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_multiple_batch_inputs_same_indices",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_batch_and_scalar_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    input2 = StepInputDefinition(
        ParameterSpecification("input2"),
        NodeInputCategory(),
        is_batch=False,
        points_to_input=True,
        dimensionality=0,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1, "input2": input2},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": [10, 20], "input2": 99}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1]})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_batch_and_scalar_input",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_compound_input_dict():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("x"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    input2 = StepInputDefinition(
        ParameterSpecification("y"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    compound = CompoundStepInputDefinition([input1, input2], is_list=False)
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"point": compound},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"x": [1, 2], "y": [3, 4]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1]})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_compound_input_dict",
        "construct_simd_step_input",
        "9",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_compound_input_list():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("a"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    input2 = StepInputDefinition(
        ParameterSpecification("b"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    compound = CompoundStepInputDefinition([input1, input2], is_list=True)
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"listparam": compound},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"a": [7, 8], "b": [9, 10]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1]})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_compound_input_list",
        "construct_simd_step_input",
        "9",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_empty_batch_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": []}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): []})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_empty_batch_input",
        "construct_simd_step_input",
        "7",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_scalar_input_only():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=False,
        points_to_input=True,
        dimensionality=0,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": 42}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({})
    branching_manager = BranchingManager({})
    with pytest.raises(ExecutionEngineRuntimeError):
        _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
            step_node,
            runtime_parameters,
            execution_cache,
            dynamic_batches_manager,
            branching_manager,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            construct_simd_step_input,
            "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
            None,
            "test_scalar_input_only",
            "construct_simd_step_input",
            "7_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_mismatched_batch_indices():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("a"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    input2 = StepInputDefinition(
        ParameterSpecification("b"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    compound = CompoundStepInputDefinition([input1, input2], is_list=True)
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"listparam": compound},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"a": [1, 2], "b": [3, 4, 5]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1, 2]})
    branching_manager = BranchingManager({})
    with pytest.raises(ExecutionEngineRuntimeError):
        _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
            step_node,
            runtime_parameters,
            execution_cache,
            dynamic_batches_manager,
            branching_manager,
        )
        _call__bound__arguments.apply_defaults()
        codeflash_return_value = codeflash_wrap(
            construct_simd_step_input,
            "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
            None,
            "test_mismatched_batch_indices",
            "construct_simd_step_input",
            "9_0",
            codeflash_loop_index,
            **_call__bound__arguments.arguments,
        )


def test_mask_applied_to_batch_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=["branch1"],
    )
    runtime_parameters = {"input1": ["x", "y", "z"]}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): [0, 1, 2]})
    branching_manager = BranchingManager({"branch1": {0, 2}})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_mask_applied_to_batch_input",
        "construct_simd_step_input",
        "7",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_accepts_empty_values_returns_empty():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True, accepts_empty=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": []}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): []})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_accepts_empty_values_returns_empty",
        "construct_simd_step_input",
        "7",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_batch_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    N = 1000
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    runtime_parameters = {"input1": list(range(N))}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): list(range(N))})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_large_batch_input",
        "construct_simd_step_input",
        "8",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output


def test_large_compound_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    N = 500
    input_defs = []
    runtime_parameters = {}
    for i in range(5):
        pname = f"p{i}"
        input_defs.append(
            StepInputDefinition(
                ParameterSpecification(pname),
                NodeInputCategory(),
                is_batch=True,
                points_to_input=True,
                dimensionality=1,
            )
        )
        runtime_parameters[pname] = list(range(N))
    compound = CompoundStepInputDefinition(input_defs, is_list=False)
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"comp": compound},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): list(range(N))})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_large_compound_input",
        "construct_simd_step_input",
        "10",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(5):
        pname = f"p{i}"


def test_large_masked_batch_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    N = 1000
    input1 = StepInputDefinition(
        ParameterSpecification("input1"),
        NodeInputCategory(),
        is_batch=True,
        points_to_input=True,
        dimensionality=1,
    )
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"input1": input1},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=["branch1"],
    )
    runtime_parameters = {"input1": list(range(N))}
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): list(range(N))})
    mask = set((i for i in range(N) if i % 2 == 0))
    branching_manager = BranchingManager({"branch1": mask})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_large_masked_batch_input",
        "construct_simd_step_input",
        "9",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i, v in enumerate(result.parameters["input1"].data):
        if i % 2 == 0:
            pass
        else:
            pass


def test_large_compound_list_input():
    codeflash_loop_index = int(os.environ["CODEFLASH_LOOP_INDEX"])
    N = 200
    input_defs = []
    runtime_parameters = {}
    for i in range(10):
        pname = f"arr{i}"
        input_defs.append(
            StepInputDefinition(
                ParameterSpecification(pname),
                NodeInputCategory(),
                is_batch=True,
                points_to_input=True,
                dimensionality=1,
            )
        )
        runtime_parameters[pname] = list(range(N))
    compound = CompoundStepInputDefinition(input_defs, is_list=True)
    step_manifest = WorkflowBlockManifest(accepts_batch=True)
    step_node = StepNode(
        name="test_step",
        input_data={"arrs": compound},
        step_manifest=step_manifest,
        step_execution_dimensionality=1,
        execution_branches_impacting_inputs=[],
    )
    execution_cache = ExecutionCache()
    dynamic_batches_manager = DynamicBatchesManager({(0,): list(range(N))})
    branching_manager = BranchingManager({})
    _call__bound__arguments = inspect.signature(construct_simd_step_input).bind(
        step_node,
        runtime_parameters,
        execution_cache,
        dynamic_batches_manager,
        branching_manager,
    )
    _call__bound__arguments.apply_defaults()
    codeflash_return_value = codeflash_wrap(
        construct_simd_step_input,
        "tests.inference.unit_tests.test_construct_simd_step_input__unit_test_1",
        None,
        "test_large_compound_list_input",
        "construct_simd_step_input",
        "10",
        codeflash_loop_index,
        **_call__bound__arguments.arguments,
    )
    codeflash_output = codeflash_return_value
    result = codeflash_output
    for i in range(10):
        pass
