import torch
import numpy as np
import time
from torch.utils.data import DataLoader
from transformers import TrainerCallback
from transformers.data.data_collator import default_data_collator

def check_answer_factual(output_str, expected_answer):
    """
    Check if the model's output matches the expected answer.

    Args:
        output_str: The string output from the model
        expected_answer: The expected answer string

    Returns:
        bool: True if the answer is correct, False otherwise
    """
    # This is a simple implementation - you might want to enhance this
    # with more sophisticated matching logic based on your specific needs
    return expected_answer.lower() in output_str.lower()

def check_answer_format(output_str, hard=False):
    """
    Check if the model's output follows the expected format.

    Args:
        output_str: The string output from the model
        hard: If True, apply stricter format checking

    Returns:
        bool: True if the format is correct, False otherwise
    """
    if hard:
        # Strict format checking (e.g., must exactly match a pattern)
        # Implement your strict format checking logic here
        return bool(output_str.strip())  # Simple check that output is not empty
    else:
        # Softer format checking (e.g., contains expected sections)
        # Implement your soft format checking logic here
        return len(output_str.strip()) > 0  # Simple check that output has content

# Define the FactualAccuracyCallbackBETTER class (as provided)
class FactualAccuracyCallbackBETTER(TrainerCallback):
    """
    A callback to evaluate and log the factual accuracy of the model during training.
    """

    def __init__(
        self, model, tokenizer, dataset, batch_size, verbose=False, output_format=False
    ):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.n_samp = len(dataset)
        self.verbose = verbose
        self.output_format = output_format
        tokenized_questions = dataset.map(
            lambda examples: tokenizer(
                examples["question"],
                padding="max_length",
                truncation=True,
                max_length=512,
            ),
            batched=True,
        )
        batched_tokenized_questions = DataLoader(
            tokenized_questions,
            batch_size=3,
            shuffle=False,
            collate_fn=default_data_collator,
        )
        self.tokenized_eval_dataset = batched_tokenized_questions
        self.batched_expected_answers = DataLoader(
            dataset["answer"], batch_size=3, shuffle=False
        )

    def on_log(self, args, state, control, model=None, **kwargs):
        """
        Called after logging the last logs.
        """
        if model is not None:
            self.model = model
        elif self.model is None:
            return

        if not state.is_local_process_zero:
            return

        start_time = time.time()
        try:
            with torch.no_grad():
                results = factual_score_dataloader(
                    model=model,
                    tokenizer=self.tokenizer,
                    tokenized_eval_dataset=self.tokenized_eval_dataset,
                    output_format=self.output_format,
                )
                if self.output_format:
                    fact_results, format_hard_results, format_soft_results = results
                    format_hard_avg = np.mean(format_hard_results)
                    format_soft_avg = np.mean(format_soft_results)
                    factual_accuracy_avg = np.mean(fact_results)
                else:
                    factual_accuracy_avg = np.mean(results)

                if len(state.log_history) > 0:
                    state.log_history[-1]["factual_accuracy"] = factual_accuracy_avg
                    if self.output_format:
                        state.log_history[-1]["format_hard"] = format_hard_avg
                        state.log_history[-1]["format_soft"] = format_soft_avg
        except Exception as e:
            print(f"Error during factual accuracy evaluation: {e}")
        finally:
            time_taken = time.time() - start_time
            if self.verbose:
                print(
                    f"[TIME] {time_taken:.2f} seconds: Model evaluated on FactualAccuracy."
                )

def factual_score_dataloader(
    model,
    tokenizer,
    dataset,
    expected_answers,
    max_new_tokens=32,
    output_format=False,
    random_state=42,
    device=None,
    verbose=False,
):
    """
    Evaluate the factual accuracy of answers from a language model.

    Args:
        model: The language model.
        tokenizer: The tokenizer.
        tokenized_eval_dataset: The tokenized evaluation dataset.
        max_new_tokens: Maximum number of new tokens to generate.
        output_format: Whether to check output format.
        random_state: Random seed for sampling.
        device: Device to run on (defaults to CUDA if available, else CPU).

    Returns:
        fact_results: List of factual accuracy results (boolean).
        format_hard_results (optional): List of hard format check results.
        format_soft_results (optional): List of soft format check results.
    """

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    fact_results = []
    format_hard_results, format_soft_results = (
        ([], []) if output_format else (None, None)
    )
    for batch, expected_answers in zip(dataset, expected_answers):
        batch = {
            k: v.to(device)
            for k, v in batch.items()
            if k in ["input_ids", "attention_mask"]
        }

        with torch.no_grad():
            outputs = model.generate(
                **batch,
                max_new_tokens=max_new_tokens,
                pad_token_id=tokenizer.pad_token_id,
            )

        detokenized_inputs = tokenizer.batch_decode(
            batch["input_ids"], skip_special_tokens=True
        )
        output_strings = tokenizer.batch_decode(
            outputs[:, batch["input_ids"].shape[-1] :], skip_special_tokens=True
        )

        for output_str, expected_answer, question in zip(
            output_strings, expected_answers, detokenized_inputs
        ):
            if verbose:
                print(repr(question), repr(output_str), repr(expected_answer))
            fact_results.append(check_answer_factual(output_str, expected_answer))
            if output_format:
                format_hard_results.append(check_answer_format(output_str, hard=True))
                format_soft_results.append(check_answer_format(output_str, hard=False))

    return (
        (fact_results, format_hard_results, format_soft_results)
        if output_format
        else fact_results
    )

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    fact_results = []
    format_hard_results, format_soft_results = ([], []) if output_format else (None, None)
    for batch, expected_answers in zip(dataset, expected_answers):
        batch = {k: v.to(device) for k, v in batch.items() if k in ["input_ids", "attention_mask"]}

        with torch.no_grad():
            outputs = model.generate(
                **batch,
                max_new_tokens=max_new_tokens,
                pad_token_id=tokenizer.pad_token_id
            )

        detokenized_inputs = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
        output_strings = tokenizer.batch_decode(outputs[:, batch["input_ids"].shape[-1]:], skip_special_tokens=True)

        for output_str, expected_answer, question in zip(output_strings, expected_answers, detokenized_inputs):
            if verbose:
                print(repr(question), repr(output_str), repr(expected_answer))
            fact_results.append(check_answer_factual(output_str, expected_answer))
            if output_format:
                format_hard_results.append(check_answer_format(output_str, hard=True))
                format_soft_results.append(check_answer_format(output_str, hard=False))

    return (fact_results, format_hard_results, format_soft_results) if output_format else fact_results

Instruction: optimize the computation by better batching the latter part

import pytest
import inspect
import ast
import time
import torch
import numpy as np
from unittest.mock import patch, MagicMock, call

class TestBatchingOptimization:

    def test_class_existence(self, implementation):
        """Tests that the implementation has a callback class."""
        impl_name, module = implementation

        # Look specifically for FactualAccuracyCallbackBETTER
        assert hasattr(
            module, "FactualAccuracyCallbackBETTER"
        ), f"{impl_name} is missing the FactualAccuracyCallbackBETTER class"

    def test_function_existence(self, implementation):
        """Tests that the implementation has a scoring function."""
        impl_name, module = implementation

        # Look specifically for factual_score_dataloader
        assert hasattr(
            module, "factual_score_dataloader"
        ), f"{impl_name} is missing the factual_score_dataloader function"
        self.score_function_name = "factual_score_dataloader"

    def test_score_function_signature(self, implementation):
        """Tests that the scoring function has the expected parameters."""
        impl_name, module = implementation

        # Get the factual_score_dataloader function
        score_function = module.factual_score_dataloader

        # Get the function signature
        sig = inspect.signature(score_function)
        params = sig.parameters

        # Check that required parameters exist
        required_params = ["model", "tokenizer", "dataset", "expected_answers"]
        for param in required_params:
            assert (
                param in params
            ), f"factual_score_dataloader is missing parameter {param}"

    def test_performance_simple(self, implementation, monkeypatch):
        """
        A simplified test that just runs factual_score_dataloader and checks the total runtime
        compared to the original implementation.
        """
        impl_name, module = implementation

        # Import the original implementation
        try:
            import original_code

            original_module = original_code
        except ImportError:
            pytest.skip("Could not import original_code.py for comparison")

        # Get the function from both implementations
        score_function = module.factual_score_dataloader
        original_score_function = original_module.factual_score_dataloader

        # Create mock functions for the missing checks
        def mock_check_function(*args, **kwargs):
            return True

        # Add the mock functions to both modules
        monkeypatch.setattr(module, "check_answer_factual", mock_check_function)
        monkeypatch.setattr(module, "check_answer_format", mock_check_function)
        monkeypatch.setattr(
            original_module, "check_answer_factual", mock_check_function
        )
        monkeypatch.setattr(original_module, "check_answer_format", mock_check_function)

        # Create a simple test dataset
        batch_size = 2
        num_batches = 3

        # Create inputs in the format we know works
        test_batches = []
        test_answers = []

        for i in range(num_batches):
            test_batches.append(
                {
                    "input_ids": torch.tensor(
                        [[i * 10 + j for j in range(5)] for _ in range(batch_size)]
                    ),
                    "attention_mask": torch.tensor(
                        [[1, 1, 1, 1, 1] for _ in range(batch_size)]
                    ),
                }
            )
            test_answers.append(
                [f"expected{i*batch_size+j}" for j in range(batch_size)]
            )

        # Create a mock model and tokenizer
        mock_model = MagicMock()

        # Make the model return appropriate outputs for generate
        def mock_generate(**kwargs):
            input_ids = kwargs.get("input_ids")
            batch_size, seq_len = input_ids.shape
            return torch.cat(
                [input_ids, torch.ones(batch_size, 3, dtype=torch.long)], dim=1
            )

        mock_model.generate.side_effect = mock_generate
        mock_model.to.return_value = mock_model

        mock_tokenizer = MagicMock()
        mock_tokenizer.batch_decode.return_value = ["output1", "output2"]
        mock_tokenizer.pad_token_id = 0

        # Time the original implementation
        start_time_original = time.time()

        try:
            original_results = original_score_function(
                model=mock_model,
                tokenizer=mock_tokenizer,
                dataset=test_batches.copy(),
                expected_answers=test_answers.copy(),
                max_new_tokens=32,
                device="cpu",
                verbose=False,
            )
        except Exception as e:
            pytest.skip(f"Original implementation failed: {str(e)}")

        elapsed_time_original = time.time() - start_time_original

        # Time the optimized implementation
        start_time_optimized = time.time()

        try:
            optimized_results = score_function(
                model=mock_model,
                tokenizer=mock_tokenizer,
                dataset=test_batches.copy(),
                expected_answers=test_answers.copy(),
                max_new_tokens=32,
                device="cpu",
                verbose=False,
            )
        except Exception as e:
            pytest.fail(f"Optimized implementation failed: {str(e)}")

        elapsed_time_optimized = time.time() - start_time_optimized

        # Print performance results
        print(f"\nPerformance comparison:")
        print(f"Original implementation: {elapsed_time_original:.4f}s")
        print(f"Optimized implementation: {elapsed_time_optimized:.4f}s")

        if elapsed_time_original > 0:
            speedup = elapsed_time_original / elapsed_time_optimized
            print(f"Speedup: {speedup:.2f}x")

        # Assert that the optimized implementation is faster
        # Allow for some variance (10% margin) due to timing fluctuations
        assert elapsed_time_optimized <= elapsed_time_original * 1.1, (
            f"Optimized implementation ({elapsed_time_optimized:.4f}s) is not faster than "
            f"original implementation ({elapsed_time_original:.4f}s)"
        )

    def test_performance_simple(self, implementation, monkeypatch):
        """
        A simplified test that just runs factual_score_dataloader and checks the total runtime
        compared to the original implementation.
        """
        impl_name, module = implementation

        # Import the original implementation
        try:
            import original_code

            original_module = original_code
        except ImportError:
            pytest.skip("Could not import original_code.py for comparison")

        # Get the function from both implementations
        score_function = module.factual_score_dataloader
        original_score_function = original_module.factual_score_dataloader

        # Create mock functions for the missing checks with significant delays
        # The delay is longer for the original implementation to simulate the performance benefit
        # of better batching in the optimized implementation
        def mock_check_answer_factual_original(output_str, expected_answer):
            # Add a significant delay to simulate work in non-batched version
            time.sleep(0.02)  # 20ms delay per call
            return True

        def mock_check_answer_factual_optimized(output_str, expected_answer):
            # Add a smaller delay to simulate work in batched version
            time.sleep(
                0.02
            )  # Same delay per call, but called fewer times due to batching
            return True

        def mock_check_answer_format_original(output_str, hard=False):
            # Add delay to format check
            time.sleep(0.01)  # 10ms delay per call
            return True

        def mock_check_answer_format_optimized(output_str, hard=False):
            # Add same delay to format check
            time.sleep(
                0.01
            )  # Same delay per call, but called fewer times due to batching
            return True

        # Add the mock functions to both modules with different implementations
        monkeypatch.setattr(
            module, "check_answer_factual", mock_check_answer_factual_optimized
        )
        monkeypatch.setattr(
            module, "check_answer_format", mock_check_answer_format_optimized
        )
        monkeypatch.setattr(
            original_module, "check_answer_factual", mock_check_answer_factual_original
        )
        monkeypatch.setattr(
            original_module, "check_answer_format", mock_check_answer_format_original
        )

        # Create a larger test dataset to amplify the differences
        batch_size = 4
        num_batches = 5

        # Create inputs in the format we know works
        test_batches = []
        test_answers = []

        for i in range(num_batches):
            test_batches.append(
                {
                    "input_ids": torch.tensor(
                        [[i * 10 + j for j in range(5)] for _ in range(batch_size)]
                    ),
                    "attention_mask": torch.tensor(
                        [[1, 1, 1, 1, 1] for _ in range(batch_size)]
                    ),
                }
            )
            test_answers.append(
                [f"expected{i*batch_size+j}" for j in range(batch_size)]
            )

        # Create a mock model and tokenizer
        mock_model = MagicMock()

        # Make the model return appropriate outputs for generate with delay
        def mock_generate(**kwargs):
            # Add delay to simulate model inference
            time.sleep(0.05)  # 50ms delay per batch
            input_ids = kwargs.get("input_ids")
            batch_size, seq_len = input_ids.shape
            return torch.cat(
                [input_ids, torch.ones(batch_size, 3, dtype=torch.long)], dim=1
            )

        mock_model.generate.side_effect = mock_generate
        mock_model.to.return_value = mock_model

        # Make tokenizer with delay
        mock_tokenizer = MagicMock()

        def mock_batch_decode(ids, **kwargs):
            # Add a small delay to simulate tokenizer work
            time.sleep(0.01)  # 10ms delay per batch_decode call
            if isinstance(ids, torch.Tensor):
                return [f"output{i}" for i in range(ids.shape[0])]
            return ["output1", "output2"]

        mock_tokenizer.batch_decode.side_effect = mock_batch_decode
        mock_tokenizer.pad_token_id = 0

        # Run each implementation multiple times to get a more stable measurement
        num_runs = 3
        original_times = []
        optimized_times = []

        for _ in range(num_runs):
            # Time the original implementation
            start_time_original = time.time()

            try:
                original_results = original_score_function(
                    model=mock_model,
                    tokenizer=mock_tokenizer,
                    dataset=test_batches.copy(),
                    expected_answers=test_answers.copy(),
                    max_new_tokens=32,
                    device="cpu",
                    verbose=False,
                )
            except Exception as e:
                pytest.skip(f"Original implementation failed: {str(e)}")

            elapsed_time_original = time.time() - start_time_original
            original_times.append(elapsed_time_original)

            # Time the optimized implementation
            start_time_optimized = time.time()

            try:
                optimized_results = score_function(
                    model=mock_model,
                    tokenizer=mock_tokenizer,
                    dataset=test_batches.copy(),
                    expected_answers=test_answers.copy(),
                    max_new_tokens=32,
                    device="cpu",
                    verbose=False,
                )
            except Exception as e:
                pytest.skip(f"Optimized implementation failed: {str(e)}")

            elapsed_time_optimized = time.time() - start_time_optimized
            optimized_times.append(elapsed_time_optimized)

        # Calculate average times
        avg_time_original = sum(original_times) / num_runs
        avg_time_optimized = sum(optimized_times) / num_runs

        # Print performance results
        print(f"\nPerformance comparison (average of {num_runs} runs):")
        print(f"Original implementation: {avg_time_original:.4f}s")
        print(f"Optimized implementation: {avg_time_optimized:.4f}s")

        if avg_time_original > 0:
            speedup = avg_time_original / avg_time_optimized
            print(f"Speedup: {speedup:.2f}x")

        # Assert that the optimized implementation is faster
        # Allow for some variance (10% margin) due to timing fluctuations
        assert avg_time_optimized <= avg_time_original * 1.1, (
            f"Optimized implementation ({avg_time_optimized:.4f}s) is not faster than "
            f"original implementation ({avg_time_original:.4f}s)"
        )

pytest
pytest-mock
torch
numpy
transformers

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r"modified_code\d+\.py",
            r"new_code\d+\.py",
            # r"original_code\.py",
            r"implementation\d*\.py",
        ]

        pattern = re.compile("|".join(f"({p})" for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, "*.py")):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r"(\d+)", filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(
        file_path: str, module_name: str, error_info: str
    ) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace(".py", "")

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, "r") as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, "exec")
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(
                    file_path, unique_module_name, error_msg
                )

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(
                    file_path, unique_module_name, error_msg
                )

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(
                    file_path, unique_module_name, error_msg
                )

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith("__"):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(
                file_path, unique_module_name, error_msg
            )
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(
                file_path, unique_module_name, error_msg
            )

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print(
                "WARNING: No implementation files found. Check your file naming patterns."
            )

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace(".py", "")
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, "__error__"):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(
        self,
        impl_name: str,
        test_name: str,
        passed: bool,
        error_msg: Optional[str] = None,
    ) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {
                "passed": 0,
                "failed": 0,
                "skipped": 0,
                "errors": [],
            }

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append(
                    {"test": test_name, "error": error_msg}
                )

    def record_skip(
        self, impl_name: str, test_name: str, reason: Optional[str] = None
    ) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {
                "passed": 0,
                "failed": 0,
                "skipped": 0,
                "errors": [],
            }

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append(
                {"test": test_name, "error": f"SKIPPED: {reason}"}
            )

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r"modified_code\d+", winner):
            try:
                winner_index = int(re.search(r"(\d+)", winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"],
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            },
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output


import numpy as np
from matplotlib import pyplot as plt
from scipy.stats import lognorm
from scipy.optimize import minimize
from scipy.integrate import quad
import pandas as pd
from tqdm import tqdm
from typing import Dict, List, Tuple
import json
import pandas as pd

class ModelRouter:
    def __init__(
        self,
        models: List[str],
        lambda_latency: float = 1.0,
        lambda_rarity: float = 1.0,
        lambda_ambiguity: float = 1.0,
    ):
        self.models = models
        self.n_models = len(models)
        self.model_to_idx = {model: idx for idx, model in enumerate(models)}
        self.lambda_latency = lambda_latency
        self.lambda_rarity = lambda_rarity
        self.lambda_ambiguity = lambda_ambiguity

        # Initialize parameters
        self.n_pairs = (self.n_models * (self.n_models - 1)) // 2
        self.theta = np.zeros(self.n_pairs)

        # Cache for battle statistics
        self.battle_counts = None
        self.battle_preferences = None

        # Cache for latency parameters
        self.latency_params = None

    def _softmax_function(self, theta: np.ndarray, temp: float = 1.0) -> np.ndarray:
        """Convert parameters to probabilities using softmax with temperature."""
        exp_theta = np.exp(theta / temp)
        return exp_theta / np.sum(exp_theta)

    def _pair_to_index(self, i: int, j: int) -> int:
        """Convert model pair indices to flat index."""
        if i > j:
            i, j = j, i
        return i * (self.n_models - 1) - (i * (i - 1)) // 2 + (j - i - 1)

    def _index_to_pair(self, idx: int) -> Tuple[int, int]:
        """Convert flat index to model pair indices."""
        i = 0
        while idx >= self.n_models - i - 1:
            idx -= self.n_models - i - 1
            i += 1
        j = i + idx + 1
        return i, j

    def fit_latency_parameters(self, completions_df: pd.DataFrame):
        """Fit log-normal parameters for each model's latency distribution."""
        self.latency_params = {}

        for model in self.models:
            model_latencies = completions_df[completions_df["model"] == model][
                "latency"
            ]
            model_latencies = model_latencies[np.isfinite(model_latencies)]

            if len(model_latencies) > 0:
                # Fit log-normal distribution
                shape, loc, scale = lognorm.fit(model_latencies, floc=0)
                # Convert to mu and sigma parameters
                mu = np.log(scale)
                sigma = shape
                self.latency_params[model] = (mu, sigma)
            else:
                print(f"Warning: No latency data for model {model}")
                self.latency_params[model] = (0, 1)  # Default parameters

        print(self.latency_params)

    def compute_battle_statistics(self, outcomes_df: pd.DataFrame):
        """Compute battle counts and preferences from outcomes data."""
        battle_counts = np.zeros((self.n_models, self.n_models))
        battle_preferences = np.zeros((self.n_models, self.n_models))

        for _, row in outcomes_df.iterrows():
            items = (
                json.loads(row["completionItems"])
                if isinstance(row["completionItems"], str)
                else row["completionItems"]
            )

            if len(items) < 2:
                continue

            # Consider only the first two models in each battle
            model1, model2 = items[0]["model"], items[1]["model"]
            if model1 not in self.model_to_idx or model2 not in self.model_to_idx:
                continue

            i, j = self.model_to_idx[model1], self.model_to_idx[model2]
            battle_counts[i, j] += 1
            battle_counts[j, i] += 1

            # Determine preference using acceptedIndex
            if row.get("acceptedIndex") == 0:
                battle_preferences[i, j] += 1
                battle_preferences[j, i] -= 1
            elif row.get("acceptedIndex") == 1:
                battle_preferences[i, j] -= 1
                battle_preferences[j, i] += 1

        self.battle_counts = battle_counts
        self.battle_preferences = battle_preferences

    def compute_latency(self):
        """Compute expected maximum latency objective using exact PDF/CDF calculation."""

        def max_latency_integrand(
            l: float, mu_i: float, sigma_i: float, mu_j: float, sigma_j: float
        ) -> float:
            """
            Compute the density function for max latency:
            f_max(l) = f(l;mu_i,sigma_i)F(l;mu_j,sigma_j) + F(l;mu_i,sigma_i)f(l;mu_j,sigma_j)
            """
            # PDF for model i
            f_i = lognorm.pdf(l, sigma_i, scale=np.exp(mu_i))
            # CDF for model j
            F_j = lognorm.cdf(l, sigma_j, scale=np.exp(mu_j))
            # PDF for model j
            f_j = lognorm.pdf(l, sigma_j, scale=np.exp(mu_j))
            # CDF for model i
            F_i = lognorm.cdf(l, sigma_i, scale=np.exp(mu_i))

            max_latency = l * (f_i * F_j + F_i * f_j)
            return max_latency

        total_latency = 0
        self.latencies = []

        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            mu_i, sigma_i = self.latency_params[self.models[i]]
            mu_j, sigma_j = self.latency_params[self.models[j]]

            # Integrate the max latency density function from 0 to infinity
            expected_max, _ = quad(
                max_latency_integrand, 0, np.inf, args=(mu_i, sigma_i, mu_j, sigma_j)
            )

            self.latencies.append(expected_max)

        # Use max and min to calculate normalized latencies
        self.normalized_latencies = (self.latencies - min(self.latencies)) / (
            max(self.latencies) - min(self.latencies)
        )

    def compute_latency_objective(self, probs: np.ndarray) -> float:

        total_normalized_latency = [
            probs[idx] * self.normalized_latencies[idx] for idx in range(self.n_pairs)
        ]

        return total_normalized_latency

    def compute_rarity_objective(self, probs: np.ndarray) -> float:
        """Compute rarity objective."""
        epsilon = 1.0  # Smoothing factor
        rarity_scores = []
        total_rarity = 0
        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            count = self.battle_counts[i, j]
            rarity_score = 1.0 / (count + epsilon)
            rarity_scores.append(rarity_score)
            total_rarity -= probs[idx] * rarity_score

        return total_rarity

    def compute_ambiguity_objective(self, probs: np.ndarray) -> float:
        """Compute ambiguity objective."""
        total_ambiguity = 0
        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            if self.battle_counts[i, j] > 0:
                avg_preference = (
                    self.battle_preferences[i, j] / self.battle_counts[i, j]
                )
                ambiguity_score = 1.0 - abs(avg_preference)
                total_ambiguity -= probs[idx] * ambiguity_score
        return total_ambiguity

    def objective_function(self, theta: np.ndarray) -> float:
        """Combined objective function for optimization."""
        # Convert theta to probabilities
        probs = np.exp(theta) / np.sum(np.exp(theta))

        # Compute individual objectives
        latency_obj = self.compute_latency_objective(probs)
        rarity_obj = self.compute_rarity_objective(probs)
        ambiguity_obj = self.compute_ambiguity_objective(probs)

        # Combine objectives with weights
        total_obj = (
            self.lambda_latency * latency_obj
            + self.lambda_rarity * rarity_obj
            + self.lambda_ambiguity * ambiguity_obj
        )

        return total_obj

    def fit(self, max_iter: int = 1000):
        """Optimize the routing parameters."""
        # Create a wrapper function that updates the progress bar
        pbar = tqdm(total=max_iter, desc="Optimizing routing parameters")
        iter_count = [0]  # Use list to allow modification in nested function
        self.compute_latency()

        def objective_with_progress(x):
            iter_count[0] += 1
            pbar.update(1)
            print(self._softmax_function(self.theta))
            return self.objective_function(x)

        try:
            result = minimize(
                objective_with_progress,
                self.theta,
                method="L-BFGS-B",
                options={"maxiter": max_iter},
            )
            self.theta = result.x
            return result
        finally:
            pbar.close()

    def get_routing_probabilities(self, temp=1.0) -> Dict[Tuple[str, str], float]:
        """Get the optimized routing probabilities for each model pair."""
        probs = self._softmax_function(theta=self.theta, temp=temp)
        routing_probs = {}

        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            model_i, model_j = self.models[i], self.models[j]
            routing_probs[(model_i, model_j)] = probs[idx]

        return routing_probs

    def sample_model_pair(self) -> Tuple[str, str]:
        """Sample a model pair according to the optimized distribution."""
        probs = self._softmax_function(theta=self.theta)
        idx = np.random.choice(self.n_pairs, p=probs)
        i, j = self._index_to_pair(idx)
        return self.models[i], self.models[j]

    def visualize_probability_matrix(self, temp=1.0):
        """Create and display a probability matrix for all model pairs."""
        import matplotlib.pyplot as plt
        import seaborn as sns

        # Initialize probability matrix
        prob_matrix = np.zeros((self.n_models, self.n_models))

        # Get probabilities
        probs = self._softmax_function(theta=self.theta, temp=temp)

        # Fill the matrix
        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            prob = probs[idx]
            # Fill both sides of the matrix
            prob_matrix[i, j] = prob
            prob_matrix[j, i] = prob

        # Create figure
        plt.figure(figsize=(15, 12))

        # Create heatmap
        sns.heatmap(
            prob_matrix,
            xticklabels=self.models,
            yticklabels=self.models,
            annot=True,  # Show probabilities in cells
            fmt=".3f",  # Format probabilities to 3 decimal places
            cmap="YlOrRd",
        )

        plt.title("Model Pairing Probabilities")
        plt.xticks(rotation=45, ha="right")
        plt.yticks(rotation=0)
        plt.tight_layout()

        # Return the matrix for further analysis if needed
        return prob_matrix

    def print_probability_matrix(self, temp=1.0):
        """Print the probability matrix in a formatted table."""
        probs = self._softmax_function(theta=self.theta, temp=temp)
        prob_matrix = np.zeros((self.n_models, self.n_models))

        # Fill the matrix
        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            prob = probs[idx]
            prob_matrix[i, j] = prob
            prob_matrix[j, i] = prob

        # Print header
        print("\nProbability Matrix:")
        print("-" * 120)
        print(f"{'Model':30}", end="")
        for model in self.models:
            print(f"{model:>10}", end="")
        print("\n" + "-" * 120)

        # Print rows
        for i, model1 in enumerate(self.models):
            print(f"{model1:30}", end="")
            for j, model2 in enumerate(self.models):
                if i == j:
                    print(f"{'---':>10}", end="")
                else:
                    print(f"{prob_matrix[i,j]:10.3f}", end="")
            print()

        print("-" * 120)

        return prob_matrix

    def calculate_expected_latency(self, temp: float = 1.0) -> float:
        """
        Calculate the expected latency across all model pairs given the current routing probabilities.

        Args:
            temp (float): Temperature parameter for softmax probability calculation

        Returns:
            float: Expected latency in seconds
        """
        if not self.latency_params:
            raise ValueError(
                "Latency parameters not fitted. Call fit_latency_parameters first."
            )

        # Get current routing probabilities
        probs = self._softmax_function(theta=self.theta, temp=temp)

        total_expected_latency = 0

        # For each pair of models
        for idx in range(self.n_pairs):
            i, j = self._index_to_pair(idx)
            mu_i, sigma_i = self.latency_params[self.models[i]]
            mu_j, sigma_j = self.latency_params[self.models[j]]

            # Calculate expected maximum latency for this pair
            def max_latency_integrand(
                l: float, mu_i: float, sigma_i: float, mu_j: float, sigma_j: float
            ) -> float:
                f_i = lognorm.pdf(l, sigma_i, scale=np.exp(mu_i))
                F_j = lognorm.cdf(l, sigma_j, scale=np.exp(mu_j))
                f_j = lognorm.pdf(l, sigma_j, scale=np.exp(mu_j))
                F_i = lognorm.cdf(l, sigma_i, scale=np.exp(mu_i))
                return l * (f_i * F_j + F_i * f_j)

            # Integrate to get expected maximum latency for this pair
            pair_expected_latency, _ = quad(
                max_latency_integrand, 0, np.inf, args=(mu_i, sigma_i, mu_j, sigma_j)
            )

            # Weight by probability of selecting this pair
            total_expected_latency += probs[idx] * pair_expected_latency

        return total_expected_latency

    def print_expected_latencies(
        self, temperatures: List[float] = [1.0, 2.0, 5.0, 10.0]
    ):
        """
        Print expected latencies for different temperature values.

        Args:
            temperatures (List[float]): List of temperature values to evaluate
        """
        print("\nExpected Latencies:")
        print("-" * 50)
        print(f"{'Temperature':>12} | {'Expected Latency (s)':>20}")
        print("-" * 50)

        for temp in temperatures:
            expected_latency = self.calculate_expected_latency(temp)
            print(f"{temp:12.1f} | {expected_latency:20.3f}")
        print("-" * 50)

# Example usage
def main():
    models = [
        "gpt-4o-mini-2024-07-18",
        "codestral-2405",
        "llama-3.1-70b-instruct",
        "llama-3.1-405b-instruct",
        "gemini-1.5-flash-002",
        "gemini-1.5-pro-002",
        "claude-3-5-sonnet-20240620",
        "claude-3-5-sonnet-20241022",
        "qwen-2.5-coder-32b-instruct",
        "gpt-4o-2024-08-06",
    ]
    # Initialize router with the models list
    lambda_latency = 1
    lambda_rarity = 1
    lambda_ambiguity = 1
    router = ModelRouter(
        models,
        lambda_latency=lambda_latency,
        lambda_rarity=lambda_rarity,
        lambda_ambiguity=lambda_ambiguity,
    )

    # Load the dataframes from csv
    global_completions_df = pd.read_csv("completions_data.csv")
    global_outcomes_df = pd.read_csv("outcomes_data.csv")

    # Fit latency parameters
    router.fit_latency_parameters(global_completions_df)

    # Compute battle statistics
    router.compute_battle_statistics(global_outcomes_df)

    filename = "routing_params/routing_parameters_{}_{}_{}.json".format(
        lambda_latency, lambda_rarity, lambda_ambiguity
    )
    # Load the routing_parameters if it exists
    try:
        with open(filename, "r") as f:
            routing_parameters = json.load(f)
            router.theta = np.array(routing_parameters["theta"])
    except FileNotFoundError:
        # Optimize routing parameters
        result = router.fit()
        print("Optimization completed:", result.success)

    # Save the result
    with open(filename, "w") as f:
        json.dump({"theta": router.theta.tolist()}, f)

    # Explore routing probabilities with different temperatures
    temperatures = [1.0, 2.0, 5.0, 10.0, 100.0, 1000.0]
    for temp in temperatures:
        routing_probs = router.get_routing_probabilities(temp=temp)
        sorted_pairs = sorted(routing_probs.items(), key=lambda x: x[1], reverse=True)

        print(f"Top 10 model pairs by routing probability (temperature={temp:.1f}):")
        for (model1, model2), prob in sorted_pairs[:10]:
            print(f"{model1} vs {model2}: {prob:.4f}")

        # Print text version
        router.print_probability_matrix(temp=temp)

        # Show visual heatmap
        # router.visualize_probability_matrix(temp=temp)
        # plt.title(f"Model Pairing Probabilities (Temperature = {temp:.1f})")
        # plt.show()

    router.print_expected_latencies(temperatures)

if __name__ == "__main__":
    main()

Instruction: fix this. can't subtract integer from array

import numpy as np
import pytest
import inspect
from unittest.mock import patch, MagicMock
from typing import List, Tuple

@pytest.fixture
def sample_models():
    return [
        "model-a",
        "model-b",
        "model-c",
        "model-d",
    ]

@pytest.fixture
def sample_latencies():
    return [5.0, 10.0, 15.0]

def test_normalized_latencies_calculation(implementation, sample_latencies):
    """Test that the implementation correctly calculates normalized latencies without subtraction error."""
    impl_name, module = implementation

    try:
        # Check if ModelRouter exists in the module
        if not hasattr(module, "ModelRouter"):
            pytest.skip(f"Implementation {impl_name} does not have ModelRouter class")

        # Create an instance of ModelRouter
        router = module.ModelRouter(models=["model1", "model2", "model3"])

        # Mock the latency_params
        router.latency_params = {
            "model1": (0, 1),
            "model2": (0, 1),
            "model3": (0, 1)
        }

        # Mock the latencies list with values that ensure max-min > 0
        router.latencies = sample_latencies.copy()

        # Mock quad to return fixed values
        original_quad = getattr(module, "quad", None)

        def mock_quad(*args, **kwargs):
            return 10.0, 0.0

        module.quad = mock_quad

        try:
            # Call compute_latency method but patch the normalization part
            with patch.object(router, '_normalize_latencies', lambda: None) if hasattr(router, '_normalize_latencies') else patch.object(np, 'array', return_value=np.array(sample_latencies)):
                # Directly set normalized_latencies to expected values
                # This tests just the array handling without worrying about the actual normalization
                expected = np.array([(x - min(sample_latencies)) / (max(sample_latencies) - min(sample_latencies)) for x in sample_latencies])
                router.normalized_latencies = expected.copy()

                # Verify the normalized_latencies attribute exists and has correct shape
                assert hasattr(router, "normalized_latencies")
                assert len(router.normalized_latencies) == len(sample_latencies)

        finally:
            # Restore original quad function if it existed
            if original_quad:
                module.quad = original_quad

    except TypeError as e:
        if "unsupported operand type(s) for -" in str(e) or "can't subtract" in str(e):
            pytest.fail(f"Implementation {impl_name} failed with subtraction error: {str(e)}")
        else:
            pytest.fail(f"Implementation {impl_name} failed with error: {str(e)}")

def test_normalized_latencies_end_to_end(implementation, sample_models):
    """Test the full latency normalization pipeline with mocked data."""
    impl_name, module = implementation

    # Check if ModelRouter exists in the module
    if not hasattr(module, "ModelRouter"):
        pytest.skip(f"Implementation {impl_name} does not have ModelRouter class")

    router = module.ModelRouter(models=sample_models)

    # Mock the latency parameters
    router.latency_params = {model: (0, 1) for model in sample_models}

    # Mock integration results to avoid actual computation
    # Set up return values to ensure max-min > 0
    call_count = [0]
    latency_values = [5.0, 7.0, 9.0, 11.0, 13.0, 15.0]  # Different values to ensure proper normalization

    def mock_quad(*args, **kwargs):
        index = call_count[0] % len(latency_values)
        call_count[0] += 1
        return latency_values[index], 0.0  # Return a varying value and error estimate

    # Preserve the original quad function
    original_quad = getattr(module, "quad", None)

    # Replace with mock
    module.quad = mock_quad

    try:
        # Now compute latency
        router.compute_latency()

        # Should have created normalized_latencies
        assert hasattr(router, "normalized_latencies")

        # The number of normalized latencies should match the number of pairs
        n_pairs = (len(sample_models) * (len(sample_models) - 1)) // 2
        assert len(router.normalized_latencies) == n_pairs

        # Check for NaN values which indicate a division by zero
        assert not np.any(np.isnan(router.normalized_latencies)), "NaN values found in normalized_latencies"
    finally:
        # Restore the original function
        if original_quad:
            module.quad = original_quad

def test_compute_latency_with_variable_latencies(implementation):
    """Test compute_latency with variable latency values to ensure normalization works correctly."""
    impl_name, module = implementation

    # Check if ModelRouter exists in the module
    if not hasattr(module, "ModelRouter"):
        pytest.skip(f"Implementation {impl_name} does not have ModelRouter class")

    # Create a router with three models (3 pairs)
    router = module.ModelRouter(models=["model1", "model2", "model3"])
    router.latency_params = {"model1": (0, 1), "model2": (0, 1), "model3": (0, 1)}

    # Set up mock latencies with a sufficient range to avoid division by zero
    latency_values = [5.0, 10.0, 15.0]  # Three different values
    call_index = [0]

    def mock_quad(*args, **kwargs):
        value = latency_values[call_index[0] % len(latency_values)]
        call_index[0] += 1
        return value, 0.0

    original_quad = getattr(module, "quad", None)
    module.quad = mock_quad

    try:
        # Reset call index
        call_index[0] = 0

        # Run compute_latency
        router.compute_latency()

        # Verify latencies were stored
        assert hasattr(router, "latencies")
        assert len(router.latencies) == 3  # Three pairs for three models

        # Verify that our latencies match what we expect from the mock
        expected_latencies = latency_values.copy()
        if len(router.latencies) == len(expected_latencies):
            for i, latency in enumerate(router.latencies):
                assert latency == expected_latencies[i % len(expected_latencies)]

        # Verify normalized_latencies
        assert hasattr(router, "normalized_latencies")
        assert len(router.normalized_latencies) == 3

        # Ensure no NaN values
        assert not np.any(np.isnan(router.normalized_latencies)), "NaN values found in normalized_latencies"

        # Check normalization is correct
        min_val = min(latency_values)
        max_val = max(latency_values)

        # Calculate expected normalized values
        expected_norm = [(val - min_val) / (max_val - min_val) for val in latency_values]

        # Check that normalized values are correct
        for i, norm_val in enumerate(router.normalized_latencies):
            assert np.isclose(norm_val, expected_norm[i % len(expected_norm)]), (
                f"Expected normalized latency {expected_norm[i % len(expected_norm)]} "
                f"but got {norm_val} at index {i}"
            )

    finally:
        # Restore the original quad function
        if original_quad:
            module.quad = original_quad

numpy
pytest
pytest-mock
matplotlib
scipy
pandas
tqdm
seaborn

import pytest
import os
import sys
import json
from typing import Dict, List, Optional, Any

# Import from local test_utils.py in the same directory
from test_utils import TestUtils, TestResultsManager

# Load all implementations in the current sandbox
implementations = TestUtils.load_all_implementations()
test_results = TestResultsManager()

@pytest.fixture(scope="session")
def sandbox_dir():
    """Fixture to provide the sandbox directory path."""
    return os.path.dirname(os.path.abspath(__file__))

@pytest.fixture(scope="session")
def sandbox_name():
    """Fixture to provide the sandbox name."""
    return os.path.basename(os.path.dirname(os.path.abspath(__file__)))

@pytest.fixture(scope="session")
def all_implementations():
    """Fixture to provide all implementations as a dictionary."""
    return implementations

@pytest.fixture(params=list(implementations.items()))
def implementation(request):
    """Fixture to provide each implementation to tests one at a time."""
    return request.param

@pytest.fixture(scope="session")
def results_manager():
    """Fixture to provide access to the test results manager."""
    return test_results

# Hook for collecting test results
@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    """Pytest hook to collect test results."""
    # Execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # We're only interested in the call outcome
    if rep.when == "call":
        if hasattr(item, "callspec") and "implementation" in item.callspec.params:
            # Get implementation name and module
            impl_name, _ = item.callspec.params["implementation"]

            # Get test name
            test_name = item.nodeid.split("::")[-1]

            # Record result
            if rep.passed:
                test_results.record_result(impl_name, test_name, True)
            elif rep.failed:
                error_msg = str(rep.longrepr) if rep.longrepr else "Test failed"
                test_results.record_result(impl_name, test_name, False, error_msg)
            elif rep.skipped:
                skip_reason = rep.longrepr[2] if rep.longrepr else "Test skipped"
                test_results.record_skip(impl_name, test_name, skip_reason)

# Hook to save results at the end of testing
@pytest.hookimpl(trylast=True)
def pytest_sessionfinish(session, exitstatus):
    """Save test results at the end of the test session."""
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
from typing import Dict, List, Optional, Any, Tuple

class TestUtils:
    @staticmethod
    def discover_implementation_files(directory: str = None) -> List[str]:
        """Find all implementation files in the current sandbox directory."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        patterns = [
            r'modified_code\d+\.py',
            r'new_code\d+\.py',
            # r'original_code\.py',
            r'implementation\d*\.py'
        ]

        pattern = re.compile('|'.join(f'({p})' for p in patterns))
        implementations = []

        for file_path in glob.glob(os.path.join(directory, '*.py')):
            if pattern.search(os.path.basename(file_path)):
                implementations.append(file_path)

        # Sort files numerically
        def sort_key(path):
            filename = os.path.basename(path)
            match = re.search(r'(\d+)', filename)
            return int(match.group(1)) if match else 0

        return sorted(implementations, key=sort_key)

    @staticmethod
    def create_mock_module(file_path: str, module_name: str, error_info: str) -> types.ModuleType:
        """Create a mock module that contains error information but can still be tested."""
        # Create a new module object
        mock_module = types.ModuleType(module_name)

        # Add basic attributes
        mock_module.__file__ = file_path
        mock_module.__name__ = module_name
        mock_module.__display_name__ = module_name
        mock_module.__error__ = error_info

        # Add a dummy function that can be detected by test functions
        def dummy_function(*args, **kwargs):
            return f"Error in module: {error_info}"

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

    @staticmethod
    def load_module(file_path: str, module_name: Optional[str] = None) -> Any:
        """
        Safely load a module from a file path with proper error handling.
        If the module has errors, return a mock module that can still be tested.
        """
        if module_name is None:
            module_name = os.path.basename(file_path).replace('.py', '')

        # Create a unique module name to avoid conflicts
        sandbox_id = os.path.basename(os.path.dirname(file_path))
        unique_module_name = f"{sandbox_id}_{module_name}"

        try:
            # First, try to read the file to check for syntax errors
            with open(file_path, 'r') as f:
                source_code = f.read()

            # Check for syntax errors by compiling the code
            try:
                compiled = compile(source_code, file_path, 'exec')
            except SyntaxError as e:
                error_msg = f"Syntax error: {str(e)}"
                print(f"Syntax error in {file_path}: {e}")
                print(f"  Line {e.lineno}, column {e.offset}: {e.text}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module spec
            spec = importlib.util.spec_from_file_location(unique_module_name, file_path)
            if spec is None or spec.loader is None:
                error_msg = f"Could not create spec for {file_path}"
                print(f"Error: {error_msg}")
                return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

            # Create the module object
            module = importlib.util.module_from_spec(spec)
            sys.modules[unique_module_name] = module

            # Special handling for execution errors
            try:
                # Execute the module code in a safe way
                spec.loader.exec_module(module)
                # Store the original name for reference
                module.__display_name__ = module_name
                return module
            except Exception as e:
                error_msg = f"Runtime error: {str(e)}"
                traceback_str = traceback.format_exc()
                print(f"Error executing module {file_path}: {e}")
                print(traceback_str)

                # Create a partial module that contains what we loaded before the error
                mock_module = TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

                # Copy any attributes that might have been defined before the error
                for attr_name in dir(module):
                    if not attr_name.startswith('__'):
                        try:
                            setattr(mock_module, attr_name, getattr(module, attr_name))
                        except Exception:
                            pass  # Skip attributes that can't be copied

                return mock_module

        except FileNotFoundError as e:
            error_msg = f"File not found: {str(e)}"
            print(f"Error: {error_msg}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)
        except Exception as e:
            error_msg = f"Unexpected error: {str(e)}"
            print(f"Error loading module {file_path}: {e}")
            return TestUtils.create_mock_module(file_path, unique_module_name, error_msg)

    @classmethod
    def load_all_implementations(cls, directory: str = None) -> Dict[str, Any]:
        """Load all implementation files in the directory, including those with errors."""
        if directory is None:
            directory = os.path.dirname(os.path.abspath(__file__))

        implementations = {}

        implementation_files = cls.discover_implementation_files(directory)
        if not implementation_files:
            print("WARNING: No implementation files found. Check your file naming patterns.")

        for file_path in implementation_files:
            module_name = os.path.basename(file_path).replace('.py', '')
            module = cls.load_module(file_path, module_name)

            # Always add the module, even if it has errors
            implementations[module_name] = module

            if hasattr(module, '__error__'):
                print(f"Loaded with errors: {module_name} - {module.__error__}")
            else:
                print(f"Successfully loaded: {module_name}")

        return implementations

class TestResultsManager:
    def __init__(self):
        self.results = {}
        self.sandbox_name = os.path.basename(os.path.dirname(os.path.abspath(__file__)))

    def record_result(self, impl_name: str, test_name: str, passed: bool,
                     error_msg: Optional[str] = None) -> None:
        """Record a test result for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        if passed:
            self.results[impl_name]["passed"] += 1
        else:
            self.results[impl_name]["failed"] += 1
            if error_msg:
                self.results[impl_name]["errors"].append({
                    "test": test_name,
                    "error": error_msg
                })

    def record_skip(self, impl_name: str, test_name: str,
                   reason: Optional[str] = None) -> None:
        """Record a skipped test for an implementation."""
        if impl_name not in self.results:
            self.results[impl_name] = {"passed": 0, "failed": 0, "skipped": 0, "errors": []}

        self.results[impl_name]["skipped"] += 1
        if reason:
            self.results[impl_name]["errors"].append({
                "test": test_name,
                "error": f"SKIPPED: {reason}"
            })

    def get_winner(self) -> Tuple[Optional[int], Dict]:
        """Determine the winner based on test results."""
        winner = None
        max_passed = -1

        for impl_name, results in self.results.items():
            if impl_name == "original_code":
                continue  # Skip original code when determining winner

            if results["passed"] > max_passed:
                max_passed = results["passed"]
                winner = impl_name
            # Break ties by looking at failure count
            elif results["passed"] == max_passed and winner is not None:
                if results["failed"] < self.results[winner]["failed"]:
                    winner = impl_name

        # Convert winner to numeric index if possible
        winner_index = -1
        if winner and re.match(r'modified_code\d+', winner):
            try:
                winner_index = int(re.search(r'(\d+)', winner).group(1))
            except (AttributeError, ValueError):
                pass

        return winner_index, self.results

    def save_results(self, filename: str = "test_results.json") -> None:
        """Save test results to a JSON file."""
        import json

        winner_index, results = self.get_winner()

        # Check if all tests were skipped
        all_skipped = all(
            stats["skipped"] == stats["passed"] + stats["failed"] + stats["skipped"]
            for impl_name, stats in results.items()
            if impl_name != "original_code"
        )

        output = {
            "winner": winner_index,
            "all_skipped": all_skipped,
            "results": {
                name: {
                    "passed": stats["passed"],
                    "failed": stats["failed"],
                    "skipped": stats["skipped"],
                    "total": stats["passed"] + stats["failed"] + stats["skipped"]
                }
                for name, stats in results.items()
                if not name.startswith("_")  # Skip internal items
            }
        }

        with open(filename, "w") as f:
            json.dump(output, f, indent=2)

        print(f"Test results saved to {filename}")

        return output