import os
import shutil

from transformers import AutoModelForCausalLM
from peft import PeftModel

from dotenv import load_dotenv

import pickle
import torch
import json

load_dotenv()

DATA_SAVE_PATH = os.getenv("DATA_SAVE_PATH")
MODEL_PATH = os.getenv("MODEL_PATH")

def save_log_to_file(log_history, file_path, append_latest_only=False):
    """
    Saves the log history to a JSON file.
    If the file already exists, it appends to it.

    Parameters:
    - log_history: List of log entries (each entry is a dict).
    - file_path: Path to the file where logs will be saved.
    - append_latest_only: If True, only the latest log entry is appended.
    """
    # Initialize current_logs
    current_logs = []

    # If the file exists, load the current logs and append to them
    if os.path.exists(file_path):
        try:
            with open(file_path, "r") as f:
                content = f.read().strip()
                if content:
                    current_logs = json.loads(content)
                else:
                    current_logs = []
        except json.JSONDecodeError:
            print(f"Warning: {file_path} contains invalid JSON. Overwriting file.")
            current_logs = []
        except Exception as e:
            print(f"An error occurred while reading {file_path}: {e}")
            current_logs = []
    else:
        # File does not exist; current_logs remains an empty list
        pass

    # Decide whether to append the entire log history or just the latest entry
    if append_latest_only and log_history:
        # Append only the most recent epoch log
        current_logs.append(log_history[-1])
    else:
        # Append the entire log history
        current_logs.extend(log_history)

    # Save the updated log history
    try:
        with open(file_path, "w") as f:
            json.dump(current_logs, f, indent=4)
    except Exception as e:
        print(f"An error occurred while writing to {file_path}: {e}")

def clear_directory(directory, delete_directory=False):
    """
    Clears all files and subdirectories within a given directory.  Optionally deletes the directory itself.
    Creates the directory if it doesn't exist and delete_directory is False.

    Args:
        directory (str): The path to the directory to clear.
        delete_directory (bool): If True, delete the directory after clearing its contents. Defaults to False.

    Raises:
        OSError: If any error occurs during file or directory removal.  Provides details about the failure.
    Example:
        clear_directory('/path/to/my/directory')
        clear_directory('/path/to/my/directory', delete_directory=True)
    """
    if not os.path.exists(directory):
        if not delete_directory:
            os.makedirs(directory)
            print(f"Directory '{directory}' created.")
        else:
            raise ValueError("Directory does not exist and delete_directory is True. Cannot proceed.")
        return

    for item in os.listdir(directory):
        item_path = os.path.join(directory, item)
        try:
            if os.path.isdir(item_path):
                shutil.rmtree(item_path)
                print(f"Removed directory: {item_path}")
            else:
                os.remove(item_path)
                print(f"Removed file: {item_path}")
        except OSError as e:
            print(f"Failed to delete '{item_path}'. Reason: {e}")
            raise # Re-raise the exception to halt execution if a deletion fails

    if delete_directory:
        try:
            os.rmdir(directory)
            print(f"Removed directory: {directory}")
        except OSError as e:
            print(f"Failed to delete '{directory}'. Reason: {e}")
            raise # Re-raise the exception to halt execution if directory removal fails

def merge_lora_model(
    model_name="pythia-31M",
    base_model_repo_name="EleutherAI/",
    model_load_path=MODEL_PATH,
    model_save_path=MODEL_PATH,
):

    my_model_path = os.path.join(model_load_path, model_name)
    param_count = model_name.lower().split("m")[0].split("-")[1]
    base_model = f"pythia-{param_count}M"

    base_model = AutoModelForCausalLM.from_pretrained(
        os.path.join(base_model_repo_name, base_model)
    )
    model = PeftModel.from_pretrained(base_model, my_model_path)
    merged_model = model.merge_and_unload()
    my_model_save_path = os.path.join(model_save_path, f"{model_name}_merged")
    merged_model.save_pretrained(my_model_save_path)

def remove_repetition(question, answer):
    if question in answer:
        return answer.replace(question, "").strip()
    return answer

def load_model(
    model_type,
    model_path=None,
    blocks_str=None,
    vanilla_model_name=None,
    host_model_name=None,
):
    """
    Loads different types of models based on the model_type parameter.

    Parameters:
    model_type (str): The type of model to load. One of 'Tuned Model', 'Vanilla Model',
                      'Transformed Model', 'Final Model', or 'Host Model'.
    model_path (str): The base path where models are stored.
    blocks_str (str): A string representing the layers or blocks used in model naming.
    vanilla_model_name (str): The name or path of the vanilla (base) model.
    host_model_name (str): The name or path of the host model.

    Returns:
    model: The loaded model object.

    Raises:
    ValueError: If an unknown model_type is provided or required parameters are missing.
    IOError: If loading the model fails.

    Example:
    model = load_model(
        model_type="Tuned Model",
        model_path="/path/to/models",
        blocks_str="1-5",
        vanilla_model_name="EleutherAI/pythia-31M"
    )
    """
    if model_type == "Tuned Model":
        model_name = vanilla_model_name.split("/")[-1]

        # save_path = os.path.join(model_path)
        # model_save_name = f"{model_name}_trained_{footer}"
        # save_path = os.path.join(save_path, model_save_name)

        tuned_model_name = f"{model_name}_trained_layers_{blocks_str}_merged"
        tuned_model = AutoModelForCausalLM.from_pretrained(
            os.path.join(model_path, f"{tuned_model_name}")
        )
        return tuned_model

    elif model_type == "Vanilla Model":
        vanilla_model = AutoModelForCausalLM.from_pretrained(vanilla_model_name)
        return vanilla_model

    elif model_type == "Transformed Model":
        name = host_model_name.split("/")[-1]
        save_path = os.path.join(model_path, f"{name}_preGRAFTED_{blocks_str}.pkl")
        with open(save_path, "rb") as f:
            transformed_model = pickle.load(f)
        return transformed_model

    elif model_type == "Final Model":
        name = host_model_name.split("/")[-1]
        model_save_name = f"{name}_GRAFTED_{blocks_str}.pkl"
        save_path = os.path.join(model_path, model_save_name)
        with open(save_path, "rb") as f:
            final_model = pickle.load(f)
        return final_model
    elif model_type == "Host Model":
        host_model = AutoModelForCausalLM.from_pretrained(host_model_name)
        return host_model

    else:
        raise ValueError(f"Unknown model type: {model_type}")

def load_batch_losses(file_path):
    """
    Loads batch loss data from a checkpoint file.

    Parameters:
    file_path (str): The path to the checkpoint file.

    Returns:
    list or None: The batch losses if available, None otherwise.

    Logs:
    An error message if loading fails.

    Example:
    batch_losses = load_batch_losses('/path/to/checkpoint.pt')
    """
    try:
        checkpoint = torch.load(file_path, map_location=torch.device("cpu"))
        batch_losses = checkpoint.get("batch_losses", None)
        if batch_losses is not None:
            logging.info(f"Batch losses loaded from {file_path}")
        else:
            logging.warning(f"No 'batch_losses' key found in checkpoint at {file_path}")
        return batch_losses
    except (FileNotFoundError, IOError, RuntimeError) as e:
        logging.error(f"Error loading checkpoint from {file_path}: {e}")
        return None

Instruction: Here's a piece of code that needs optimization: Please suggest optimizations to improve its performance. For each suggestion, explain the expected improvement and any trade-offs.

import os
import tempfile
import shutil
import time
import logging
import pytest
import statistics
from typing import List, Dict, Tuple, Any

# Set up logging for tests
logging.basicConfig(level=logging.INFO)

def create_test_directory(
    base_dir: str,
    depth: int = 3,
    files_per_dir: int = 5,
    size_kb: int = 10,
    branching_factor: int = 3,
    long_filenames: bool = False,
) -> str:
    """Create a test directory structure with specified complexity parameters."""
    test_dir = os.path.join(base_dir, f"test_dir_{time.time()}")
    os.makedirs(test_dir)

    # Create a nested directory structure with files
    _create_nested_structure(
        test_dir, depth, files_per_dir, size_kb, branching_factor, long_filenames
    )

    return test_dir

def _create_nested_structure(
    current_dir: str,
    depth: int,
    files_per_dir: int,
    size_kb: int,
    branching_factor: int,
    long_filenames: bool,
):
    """Recursively create a nested directory structure with files."""
    # Create files in the current directory
    for i in range(files_per_dir):
        if long_filenames:
            # Create files with longer names to stress string operations
            filename = f"file_with_longer_name_to_stress_string_operations_{i:05d}.txt"
        else:
            filename = f"file_{i}.txt"

        file_path = os.path.join(current_dir, filename)
        with open(file_path, "wb") as f:
            # Create a file with specified size
            # Add some variability to file sizes to better simulate real-world scenarios
            actual_size = int(size_kb * (0.5 + i % 3)) * 1024
            f.write(b"0" * actual_size)

    # Create subdirectories if depth > 0
    if depth > 0:
        for i in range(branching_factor):
            if long_filenames:
                dirname = (
                    f"subdirectory_with_longer_name_for_performance_testing_{i:03d}"
                )
            else:
                dirname = f"subdir_{i}"

            subdir = os.path.join(current_dir, dirname)
            os.makedirs(subdir)
            _create_nested_structure(
                subdir,
                depth - 1,
                files_per_dir,
                size_kb,
                branching_factor,
                long_filenames,
            )

@pytest.fixture
def performance_test_directory(tmp_path):
    """Create a consistent test directory structure for performance testing."""
    # Create a more complex directory structure to amplify performance differences
    test_dir = create_test_directory(
        tmp_path,
        depth=4,  # Deeper directory structure
        files_per_dir=20,  # More files per directory
        size_kb=5,  # Keep file size moderate
        branching_factor=4,  # More subdirectories at each level
        long_filenames=True,  # Use longer filenames to stress string operations
    )
    yield test_dir
    # Cleanup is handled by the tmp_path fixture

class PerformanceResults:
    """Class to store and analyze performance test results."""

    def __init__(self):
        self.results = {}
        self.original_results = {}

    def add_result(self, impl_name: str, operation: str, times: List[float]):
        """Add a performance test result."""
        key = f"{impl_name}_{operation}"
        avg_time = statistics.mean(times)
        std_dev = statistics.stdev(times) if len(times) > 1 else 0

        self.results[key] = {"times": times, "avg_time": avg_time, "std_dev": std_dev}

        # Store original implementation results separately for comparison
        if impl_name == "original_code":
            self.original_results[operation] = avg_time

    def get_improvement(self, impl_name: str, operation: str) -> float:
        """Calculate percentage improvement compared to original implementation."""
        if operation not in self.original_results:
            return 0.0

        key = f"{impl_name}_{operation}"
        if key not in self.results:
            return 0.0

        original_time = self.original_results[operation]
        impl_time = self.results[key]["avg_time"]

        return ((original_time - impl_time) / original_time) * 100

    def print_summary(self):
        """Print a summary of performance test results."""
        # Find all unique implementations and operations
        implementations = set()
        operations = set()

        for key in self.results:
            impl_name, operation = key.rsplit("_", 1)
            implementations.add(impl_name)
            operations.add(operation)

        # Don't include original_code in the list of implementations to compare
        if "original_code" in implementations:
            implementations.remove("original_code")

        # Print summary header
        logging.info("\n=== Performance Comparison Summary ===")

        # Print results for each operation and implementation
        for operation in operations:
            logging.info(f"\n--- Operation: {operation} ---")

            # Get original implementation time for this operation
            if operation in self.original_results:
                original_time = self.original_results[operation]
                logging.info(f"original_code: {original_time:.6f} seconds (baseline)")

                # Compare each implementation to the original
                for impl_name in implementations:
                    key = f"{impl_name}_{operation}"
                    if key in self.results:
                        impl_time = self.results[key]["avg_time"]
                        std_dev = self.results[key]["std_dev"]
                        improvement = self.get_improvement(impl_name, operation)

                        faster_slower = "faster" if improvement > 0 else "slower"
                        logging.info(
                            f"{impl_name}: {impl_time:.6f} seconds (±{std_dev:.6f}) - "
                            f"{abs(improvement):.2f}% {faster_slower} than original"
                        )

# Global results collector
performance_results = PerformanceResults()

def load_original_code(sandbox_dir):
    """Load the original code module manually."""
    from test_utils import TestUtils

    original_path = os.path.join(sandbox_dir, "original_code.py")

    if os.path.exists(original_path):
        return TestUtils.load_module(original_path, "original_code")
    return None

@pytest.fixture(scope="function")
def ensure_original_code(all_implementations, sandbox_dir):
    """Ensure original_code is available in all_implementations."""
    if "original_code" not in all_implementations:
        # Load original code
        original_module = load_original_code(sandbox_dir)
        if original_module and not hasattr(original_module, "__error__"):
            all_implementations["original_code"] = original_module
            logging.info("Successfully loaded original_code.py")
        else:
            logging.error("Failed to load original_code.py")
            return None
    return all_implementations["original_code"]

def test_clear_directory_performance(
    implementation, performance_test_directory, tmp_path, ensure_original_code
):
    """Test the performance of clear_directory implementation."""
    impl_name, module = implementation

    # Skip performance assertions for original_code itself
    is_original = impl_name == "original_code"

    # Make sure original_code is available for comparison
    original_module = ensure_original_code
    if not is_original and original_module is None:
        pytest.skip("original_code implementation required for performance comparison")

    # Number of runs for each test (increased for more reliable results)
    runs = 5

    # Run both implementations on identical copies of the test directory
    # This provides a direct, controlled comparison
    if not is_original and original_module is not None:
        #
        # === TEST CLEARING DIRECTORY (KEEPING THE DIRECTORY) ===
        #
        keep_times_impl = []
        keep_times_orig = []

        for i in range(runs):
            # Create two identical test directories
            impl_dir = os.path.join(tmp_path, f"impl_keep_run_{i}")
            orig_dir = os.path.join(tmp_path, f"orig_keep_run_{i}")

            shutil.copytree(performance_test_directory, impl_dir)
            shutil.copytree(performance_test_directory, orig_dir)

            # Measure implementation performance
            start_time = time.time()
            module.clear_directory(impl_dir, delete_directory=False)
            end_time = time.time()
            impl_time = end_time - start_time
            keep_times_impl.append(impl_time)

            # Verify functionality for implementation
            assert os.path.exists(impl_dir)
            assert len(os.listdir(impl_dir)) == 0

            # Measure original implementation performance
            start_time = time.time()
            original_module.clear_directory(orig_dir, delete_directory=False)
            end_time = time.time()
            orig_time = end_time - start_time
            keep_times_orig.append(orig_time)

            # Verify functionality for original
            assert os.path.exists(orig_dir)
            assert len(os.listdir(orig_dir)) == 0

            # Log individual run times for debugging
            logging.info(
                f"Keep run {i}: {impl_name}={impl_time:.6f}s, original={orig_time:.6f}s, diff={(orig_time-impl_time)*1000:.2f}ms"
            )

        # Calculate statistics
        avg_keep_time_impl = statistics.mean(keep_times_impl)
        avg_keep_time_orig = statistics.mean(keep_times_orig)

        # Store results
        performance_results.add_result(impl_name, "keep", keep_times_impl)
        performance_results.add_result("original_code", "keep", keep_times_orig)

        # Log comparative results
        improvement_ms = (
            avg_keep_time_orig - avg_keep_time_impl
        ) * 1000  # Convert to milliseconds
        improvement_pct = (
            (avg_keep_time_orig - avg_keep_time_impl) / avg_keep_time_orig * 100
        )

        logging.info(f"\n=== KEEP DIRECTORY PERFORMANCE ===")
        logging.info(f"{impl_name}: {avg_keep_time_impl:.6f}s")
        logging.info(f"original_code: {avg_keep_time_orig:.6f}s")
        logging.info(f"Improvement: {improvement_ms:.2f}ms ({improvement_pct:.2f}%)")

        # Assert performance improvement
        # Add a small tolerance value (0.1%) to account for measurement noise
        assert avg_keep_time_impl < avg_keep_time_orig * 0.999, (
            f"Implementation {impl_name} (avg: {avg_keep_time_impl:.6f}s) is not faster than "
            f"original implementation ({avg_keep_time_orig:.6f}s) for keep operation"
        )

        #
        # === TEST DELETING DIRECTORY ===
        #
        delete_times_impl = []
        delete_times_orig = []

        for i in range(runs):
            # Create two identical test directories
            impl_dir = os.path.join(tmp_path, f"impl_delete_run_{i}")
            orig_dir = os.path.join(tmp_path, f"orig_delete_run_{i}")

            shutil.copytree(performance_test_directory, impl_dir)
            shutil.copytree(performance_test_directory, orig_dir)

            # Measure implementation performance
            start_time = time.time()
            module.clear_directory(impl_dir, delete_directory=True)
            end_time = time.time()
            impl_time = end_time - start_time
            delete_times_impl.append(impl_time)

            # Verify functionality for implementation
            assert not os.path.exists(impl_dir)

            # Measure original implementation performance
            start_time = time.time()
            original_module.clear_directory(orig_dir, delete_directory=True)
            end_time = time.time()
            orig_time = end_time - start_time
            delete_times_orig.append(orig_time)

            # Verify functionality for original
            assert not os.path.exists(orig_dir)

            # Log individual run times for debugging
            logging.info(
                f"Delete run {i}: {impl_name}={impl_time:.6f}s, original={orig_time:.6f}s, diff={(orig_time-impl_time)*1000:.2f}ms"
            )

        # Calculate statistics
        avg_delete_time_impl = statistics.mean(delete_times_impl)
        avg_delete_time_orig = statistics.mean(delete_times_orig)

        # Store results
        performance_results.add_result(impl_name, "delete", delete_times_impl)
        performance_results.add_result("original_code", "delete", delete_times_orig)

        # Log comparative results
        improvement_ms = (
            avg_delete_time_orig - avg_delete_time_impl
        ) * 1000  # Convert to milliseconds
        improvement_pct = (
            (avg_delete_time_orig - avg_delete_time_impl) / avg_delete_time_orig * 100
        )

        logging.info(f"\n=== DELETE DIRECTORY PERFORMANCE ===")
        logging.info(f"{impl_name}: {avg_delete_time_impl:.6f}s")
        logging.info(f"original_code: {avg_delete_time_orig:.6f}s")
        logging.info(f"Improvement: {improvement_ms:.2f}ms ({improvement_pct:.2f}%)")

        # Assert performance improvement
        # Add a small tolerance value (0.1%) to account for measurement noise
        assert avg_delete_time_impl < avg_delete_time_orig * 0.999, (
            f"Implementation {impl_name} (avg: {avg_delete_time_impl:.6f}s) is not faster than "
            f"original implementation ({avg_delete_time_orig:.6f}s) for delete operation"
        )

    # For original code or if original module is not available, just run the tests
    # without comparison to collect timing data
    elif is_original or original_module is None:
        # Test clearing directory (keeping the directory)
        keep_times = []
        for i in range(runs):
            run_dir = os.path.join(tmp_path, f"keep_run_{i}")
            shutil.copytree(performance_test_directory, run_dir)

            start_time = time.time()
            module.clear_directory(run_dir, delete_directory=False)
            end_time = time.time()

            elapsed = end_time - start_time
            keep_times.append(elapsed)

            assert os.path.exists(run_dir)
            assert len(os.listdir(run_dir)) == 0

        performance_results.add_result(impl_name, "keep", keep_times)
        avg_keep_time = statistics.mean(keep_times)
        logging.info(
            f"{impl_name} clear_directory (keep) took {avg_keep_time:.6f} seconds on average"
        )

        # Test deleting directory
        delete_times = []
        for i in range(runs):
            run_dir = os.path.join(tmp_path, f"delete_run_{i}")
            shutil.copytree(performance_test_directory, run_dir)

            start_time = time.time()
            module.clear_directory(run_dir, delete_directory=True)
            end_time = time.time()

            elapsed = end_time - start_time
            delete_times.append(elapsed)

            assert not os.path.exists(run_dir)

        performance_results.add_result(impl_name, "delete", delete_times)
        avg_delete_time = statistics.mean(delete_times)
        logging.info(
            f"{impl_name} clear_directory (delete) took {avg_delete_time:.6f} seconds on average"
        )

def test_clear_directory_large_scale_performance(
    implementation, tmp_path, ensure_original_code
):
    """Test the performance of clear_directory with an extremely large directory structure."""
    impl_name, module = implementation

    # Skip performance assertions for original_code itself
    is_original = impl_name == "original_code"

    # Make sure original_code is available for comparison
    original_module = ensure_original_code
    if not is_original and original_module is None:
        pytest.skip("original_code implementation required for performance comparison")

    # For the large scale test, create an extremely complex directory structure
    # This should make performance differences more pronounced
    logging.info(
        "Creating extremely large directory structure for performance testing..."
    )

    # Compare optimized implementation with original implementation
    if not is_original and original_module is not None:
        # Create two identical test directories with extreme complexity
        impl_dir = create_test_directory(
            tmp_path,
            depth=5,  # Very deep nesting
            files_per_dir=30,  # Many files per directory
            size_kb=2,  # Small files, but many of them
            branching_factor=5,  # High branching factor for more subdirectories
            long_filenames=True,  # Use long filenames to stress string operations
        )

        # Create an identical structure for the original code
        orig_dir = os.path.join(tmp_path, "orig_extreme_test")
        shutil.copytree(impl_dir, orig_dir)

        logging.info("Directory structure created. Running performance tests...")

        # Warm-up system (to reduce variability) with a small operation
        warm_up_dir = os.path.join(tmp_path, "warm_up")
        os.makedirs(warm_up_dir)
        with open(os.path.join(warm_up_dir, "test.txt"), "w") as f:
            f.write("test")
        shutil.rmtree(warm_up_dir)

        # Measure implementation performance
        impl_start_time = time.time()
        module.clear_directory(impl_dir, delete_directory=True)
        impl_end_time = time.time()
        impl_elapsed = impl_end_time - impl_start_time

        # Verify functionality for implementation
        assert not os.path.exists(impl_dir)

        # Measure original implementation performance
        orig_start_time = time.time()
        original_module.clear_directory(orig_dir, delete_directory=True)
        orig_end_time = time.time()
        orig_elapsed = orig_end_time - orig_start_time

        # Verify functionality for original
        assert not os.path.exists(orig_dir)

        # Store results
        performance_results.add_result(impl_name, "large_scale", [impl_elapsed])
        performance_results.add_result("original_code", "large_scale", [orig_elapsed])

        # Calculate improvement
        improvement_ms = (orig_elapsed - impl_elapsed) * 1000  # Convert to milliseconds
        improvement_pct = (orig_elapsed - impl_elapsed) / orig_elapsed * 100

        # Log detailed comparison
        logging.info(f"\n=== LARGE-SCALE PERFORMANCE TEST ===")
        logging.info(f"{impl_name}: {impl_elapsed:.6f} seconds")
        logging.info(f"original_code: {orig_elapsed:.6f} seconds")
        logging.info(f"Absolute improvement: {improvement_ms:.2f} milliseconds")
        logging.info(f"Relative improvement: {improvement_pct:.2f}%")

        # Assert that new implementation is faster than original
        # Using a stricter assertion for the large-scale test
        assert impl_elapsed < orig_elapsed * 0.999, (
            f"Implementation {impl_name} ({impl_elapsed:.6f}s) is not faster than "
            f"original implementation ({orig_elapsed:.6f}s) for large scale operation"
        )

    # For original code or if original module is not available, just run the test
    elif is_original or original_module is None:
        test_dir = create_test_directory(
            tmp_path,
            depth=5,
            files_per_dir=30,
            size_kb=2,
            branching_factor=5,
            long_filenames=True,
        )

        start_time = time.time()
        module.clear_directory(test_dir, delete_directory=True)
        end_time = time.time()

        elapsed = end_time - start_time

        # Add result for large scale test
        performance_results.add_result(impl_name, "large_scale", [elapsed])

        # Log time
        logging.info(
            f"{impl_name} large scale clear_directory took {elapsed:.6f} seconds"
        )

# Session-scope fixture to print performance summary at the end
@pytest.fixture(scope="session", autouse=True)
def print_performance_summary():
    """Print a summary of performance test results at the end of the session."""
    yield
    performance_results.print_summary()

pytest
pytest-mock
torch
transformers
peft
python-dotenv

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
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import skew

# def medcouple(data):
#     data = np.sort(data)  # Сортируем данные
#     n = len(data)
#     median = np.median(data)

#     # Разделяем данные на меньше медианы и больше медианы
#     left = data[data <= median]
#     right = data[data >= median]

#     # Функция ядра h(xi, xj)
#     def h(xi, xj):
#         if xi != xj:
#             return ((xj - median) - (median - xi)) / (xj - xi)
#         return 0  # Хотя xi != xj должно исключать этот случай

#     # Специальное ядро для случаев с повторениями медианы
#     def special_h(i, j, k):
#         if i + j - 1 < k:
#             return -1
#         elif i + j - 1 == k:
#             return 0
#         elif i + j - 1 > k:
#             return 1

#     # Генерация всех возможных h(xi, xj)
#     h_values = []
#     k = len(data[data == median])  # Количество повторяющихся значений медианы
#     if k > 1:  # Обработка случая с совпадающими медианами
#         for i, xi in enumerate(left):
#             for j, xj in enumerate(right):
#                 if xi == xj == median:
#                     h_values.append(special_h(i, j, k))
#                 else:
#                     h_values.append(h(xi, xj))
#     else:
#         for xi in left:
#             for xj in right:
#                 h_values.append(h(xi, xj))

#     # Возвращаем медиану всех значений h
#     return np.median(h_values)
# Нужно ускорить и переписать функцию medcouple

def medcouple(data):
    data = np.sort(data)
    n = len(data)
    median = np.median(data)

    # Split data into left and right of the median
    left = data[data <= median]
    right = data[data >= median]

# Kernel function h(xi, xj)
    def h(xi, xj):
        if xi != xj:
            return ((xj - median) - (median - xi)) / (xj - xi)
        return 0

    # Special kernel for cases with repeated medians
    def special_h(i, j, k):
        if i + j - 1 < k:
            return -1
        elif i + j - 1 == k:
            return 0
        elif i + j - 1 > k:
            return 1

    # Generate all possible h(xi, xj)
    h_values = []
    k = len(data[data == median])  # Count of repeated median values

    # Use numpy broadcasting for efficiency
    if k > 1:
        left_indices = np.arange(len(left))
        right_indices = np.arange(len(right))
        xi, xj = np.meshgrid(left, right, indexing='ij')
        i, j = np.meshgrid(left_indices, right_indices, indexing='ij')
        h_matrix = np.where((xi == median) & (xj == median), special_h(i, j, k), h(xi, xj))
    else:
        xi, xj = np.meshgrid(left, right, indexing='ij')
        h_matrix = h(xi, xj)

    # Flatten the matrix and calculate the median of h values
    return np.median(h_matrix.flatten())

def adjusted_boxplot_bounds(data):
    """
    Вычисляет границы adjusted boxplot с учетом skewness-adjusted fences.
    """
    q1 = np.percentile(data, 25)
    q3 = np.percentile(data, 75)
    iqr = q3 - q1
    _medcouple = medcouple(data)

    if _medcouple > 0:
        lower_fence = q1 - 1.5 * np.exp(-4 * _medcouple) * iqr
        upper_fence = q3 + 1.5 * np.exp(3 * _medcouple) * iqr
    else:
        lower_fence = q1 - 1.5 * np.exp(-3 * _medcouple) * iqr
        upper_fence = q3 + 1.5 * np.exp(4 * _medcouple) * iqr

    return lower_fence, upper_fence

def normalize_column(data):
    """
    Нормализация с использованием adjusted boxplot.
    """
    lower_fence, upper_fence = adjusted_boxplot_bounds(data)
    print(lower_fence)
    return (data - lower_fence) / (upper_fence - lower_fence)

# Генерация данных
np.random.seed(42)
data_normal = np.random.normal(loc=50, scale=10, size=10000)
data_skewed = np.random.exponential(scale=20, size=10000)
data_skewed = np.concatenate([data_skewed[5:], [200, 250, 300, -100, -50]])
data_with_outliers = np.concatenate([data_normal, [150, 160, 170]])

# Нормализация
df = pd.DataFrame({
    "Normal": data_normal,
    "Skewed": data_skewed,
    # "With_Outliers": data_with_outliers[3:],
})

normalized_df = df.apply(normalize_column)

plt.figure(figsize=(16, 4), dpi=250)

bins = np.linspace(-5, 200, 206)
bin_width = bins[1] - bins[0]  # Ширина одного бина

for col in df.columns:
    # plt.hist(df[col], bins=50, alpha=0.5, label=f'{col} - Original')
    # Вычисляем гистограммы без построения
    hist, _ = np.histogram(df[col], bins=bins)

    # Позиции столбиков для каждой гистограммы
    bin_centers = (bins[:-1] + bins[1:]) / 2

    # Смещение для каждого набора данных
    offset = bin_width / 4
    plt.bar(bin_centers - offset, hist, width=bin_width, align='center', alpha=0.2, label=f'{col}')
    plt.legend()
    plt.title(f"Histogram Before Normalization")
plt.xlim(-10, 200)
plt.show()

bins = np.linspace(-2, 2, 101)
bin_width = bins[1] - bins[0]  # Ширина одного бина

plt.figure(figsize=(16, 4), dpi=250)
for col in normalized_df.columns:
    # plt.hist(normalized_df[col], bins=50, alpha=0.5, label=f'{col} - Normalized')
    hist, _ = np.histogram(normalized_df[col], bins=bins)

    # Позиции столбиков для каждой гистограммы
    bin_centers = (bins[:-1] + bins[1:]) / 2
    # Смещение для каждого набора данных
    offset = bin_width / 2
    plt.bar(bin_centers - offset, hist, width=bin_width, align='center', label=f'{col}', alpha=0.2)
    plt.legend()
    plt.title(f"Histogram After Normalization")
plt.show()

Instruction: ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

import pytest
import numpy as np
import importlib.util
import inspect
import os
import re
import sys
import subprocess
import tempfile
from pathlib import Path

def test_medcouple_fixes_truth_value_error(implementation):
    """Test that implementation properly handles array truth value ambiguity"""
    impl_name, module = implementation

    # Test with an array having repeated median values
    np.random.seed(42)
    data_with_repeated_medians = np.array([1, 2, 3, 4, 4, 4, 5, 6, 7])

    # This should not raise ValueError about array truth value
    try:
        result = module.medcouple(data_with_repeated_medians)
        assert isinstance(
            result, (int, float)
        ), f"Expected numeric result, got {type(result)}"
    except ValueError as e:
        if "truth value of an array" in str(e):
            pytest.fail(f"Implementation {impl_name} still has truth value error: {e}")
        else:
            raise  # Other ValueErrors should be propagated

def test_special_h_scalar_operation(implementation):
    """Test that special_h function works with scalar operations"""
    impl_name, module = implementation

    # Get medcouple function source
    source = inspect.getsource(module.medcouple)

    # Find where special_h is used
    special_h_usage = re.search(r"special_h\((.*?)\)", source)

    if special_h_usage:
        # The special_h function is defined and used, so we can test its functionality
        # directly during medcouple run

        # Create a test array with repeated medians
        test_data = np.array([1, 2, 3, 4, 4, 4, 5, 6])
        result = module.medcouple(test_data)
        assert isinstance(
            result, (int, float)
        ), "medcouple should return a numeric value"

def test_adjusted_boxplot_bounds_after_fix(implementation):
    """Test that adjusted_boxplot_bounds function works correctly with the fixed medcouple"""
    impl_name, module = implementation

    # Test with normal distribution
    np.random.seed(42)
    normal_data = np.random.normal(0, 1, 100)

    lower_fence, upper_fence = module.adjusted_boxplot_bounds(normal_data)

    # For normal distribution, bounds should be roughly symmetric
    assert isinstance(lower_fence, (int, float)), "Lower fence should be a scalar"
    assert isinstance(upper_fence, (int, float)), "Upper fence should be a scalar"

    # Ensure the bounds are reasonable
    q1 = np.percentile(normal_data, 25)
    q3 = np.percentile(normal_data, 75)
    iqr = q3 - q1

    # Lower fence should be below q1 and upper fence should be above q3
    assert lower_fence < q1, "Lower fence should be below Q1"
    assert upper_fence > q3, "Upper fence should be above Q3"

def run_visualization_test(module_path, module_name):
    """Run visualization test in a subprocess to ensure proper cleanup"""
    with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp_file:
        temp_path = temp_file.name

        # Create a temporary script that imports the module and runs visualization
        script_content = f"""
import sys
import os
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt

# Add the parent directory to path to ensure the module can be imported
sys.path.insert(0, os.path.dirname('{module_path}'))

# Import the module
import importlib.util
spec = importlib.util.spec_from_file_location('{module_name}', '{module_path}')
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)

# Generate data
np.random.seed(42)
data_normal = np.random.normal(loc=50, scale=10, size=100)
data_skewed = np.random.exponential(scale=20, size=100)
data_skewed = np.concatenate([data_skewed[5:], [200, 250, 300, -100, -50]])

# Create DataFrame
import pandas as pd
df = pd.DataFrame({{
    "Normal": data_normal,
    "Skewed": data_skewed,
}})

# Apply normalization
normalized_df = df.apply(module.normalize_column)

# Test that normalized_df has expected structure - save results
assert isinstance(normalized_df, pd.DataFrame), "Expected DataFrame as result"
assert normalized_df.shape == df.shape, "Normalized DataFrame should have same shape as input"

# Generate a small plot and save to file instead of displaying
plt.figure(figsize=(8, 4))
for col in df.columns:
    plt.hist(df[col], bins=20, alpha=0.5, label=col)
plt.legend()
plt.savefig('test_plot.png')
plt.close('all')

# Create a second plot for normalized data
plt.figure(figsize=(8, 4))
for col in normalized_df.columns:
    plt.hist(normalized_df[col], bins=20, alpha=0.5, label=f'{{col}} (normalized)')
plt.legend()
plt.savefig('test_plot_normalized.png')
plt.close('all')

# Exit cleanly
plt.close('all')
"""

        temp_file.write(script_content.encode("utf-8"))

    try:
        # Run the script in a subprocess
        result = subprocess.run(
            [sys.executable, temp_path],
            capture_output=True,
            text=True,
            timeout=30,  # Set a timeout to avoid hanging
        )

        # Check for errors
        if result.returncode != 0:
            raise RuntimeError(f"Subprocess failed with error: {result.stderr}")

        # Clean up test plots
        for plot_file in ["test_plot.png", "test_plot_normalized.png"]:
            if os.path.exists(plot_file):
                os.remove(plot_file)

    finally:
        # Delete the temporary script
        if os.path.exists(temp_path):
            os.remove(temp_path)

def test_end_to_end_script_execution(implementation):
    """Test that the full script runs without any ValueError about array truth values"""
    impl_name, module = implementation
    module_path = module.__file__

    # Look for all functions in the module
    all_functions = [
        name
        for name, obj in inspect.getmembers(module)
        if inspect.isfunction(obj) and obj.__module__ == module.__name__
    ]

    # The script should have the key functions: medcouple, adjusted_boxplot_bounds, normalize_column
    expected_functions = ["medcouple", "adjusted_boxplot_bounds", "normalize_column"]
    for func_name in expected_functions:
        assert (
            func_name in all_functions
        ), f"Expected function {func_name} not found in {impl_name}"

    # Use the subprocess function to run the visualization test
    try:
        run_visualization_test(module_path, module.__name__)
    except Exception as e:
        if "truth value of an array" in str(e):
            pytest.fail(f"Implementation {impl_name} still has truth value error: {e}")
        else:
            raise  # Other errors should be propagated

numpy
pandas
matplotlib
scipy
pytest
pytest-mock

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

# Configure matplotlib to be non-interactive before all imports
try:
    import matplotlib

    matplotlib.use("Agg")  # Use non-interactive backend
except ImportError:
    pass

# 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:
                # Fix for handling different types of longrepr
                skip_reason = "Test skipped"
                if rep.longrepr:
                    if isinstance(rep.longrepr, tuple) and len(rep.longrepr) >= 3:
                        skip_reason = rep.longrepr[2]
                    else:
                        skip_reason = str(rep.longrepr)
                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()

# Hook to disable interactive matplotlib
@pytest.hookimpl(tryfirst=True)
def pytest_configure(config):
    """Configure the test environment before tests start."""
    # Ensure matplotlib is configured for non-interactive use
    try:
        import matplotlib

        matplotlib.use("Agg")  # Force non-interactive backend
        import matplotlib.pyplot as plt

        plt.ioff()  # Turn off interactive mode
    except ImportError:
        pass

# Add a fixture to close figures after tests
@pytest.fixture(autouse=True)
def close_figures():
    """Auto-use fixture to close matplotlib figures after each test."""
    yield
    # After the test function completes, close all matplotlib figures
    try:
        import matplotlib.pyplot as plt

        plt.close("all")
    except ImportError:
        pass  # If matplotlib isn't installed, do nothing

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

@staticmethod
def suppress_plots():
    """Configure matplotlib to not display plots (useful for test environments)"""
    try:
        import matplotlib

        matplotlib.use("Agg")  # Use non-interactive backend
        import matplotlib.pyplot as plt

        plt.ioff()  # Turn off interactive mode
    except ImportError:
        pass  # If matplotlib is not installed, do nothing

    return True


import requests
import json
import os
from huggingface_hub import InferenceClient
from datetime import datetime
from PIL import Image

class ImageGenerator:
    def __init__(self, openrouter_key, hf_token, output_folder):
        self.openrouter_key = openrouter_key
        self.hf_token = hf_token
        self.output_folder = output_folder

        # Create output folder if it doesn't exist
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

    def generate_prompt(self, base_prompt, model="openai/gpt-3.5-turbo"):
        response = requests.post(
            url="https://openrouter.ai/api/v1/chat/completions",
            headers={
                "Authorization": f"Bearer {self.openrouter_key}",
                "X-Title": "ImagePromptGenerator",
            },
            data=json.dumps({
                "model": model,
                "messages": [
                    {
                        "role": "user",
                        "content": base_prompt
                    }
                ],
                "temperature": 0.9,  # Higher temperature for more creativity
                "max_tokens": 150,
                "top_p": 0.9,
                "frequency_penalty": 0.5,
                "presence_penalty": 0.5
            })
        )

        return response.json()['choices'][0]['message']['content']

    def create_image(self, prompt, hf_model="black-forest-labs/FLUX.1-schnell"):
        client = InferenceClient(hf_model, token=self.hf_token)

        # Generate image with additional parameters for creativity
        image = client.text_to_image(
            prompt
        )

        return image

    def save_image(self, image, prompt):
        # Create timestamp for unique filename
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # Create sanitized filename from the first 30 chars of prompt
        sanitized_prompt = "".join(x for x in prompt[:30] if x.isalnum() or x in (' ','-','_')).strip()
        filename = f"{timestamp}_{sanitized_prompt}.png"

        # Save image
        filepath = os.path.join(self.output_folder, filename)
        image.save(filepath)

        return filepath

    def generate_batch(self, base_prompt, n_images=1, openrouter_model="openai/gpt-3.5-turbo", hf_model="black-forest-labs/FLUX.1-schnell"):
        generated_files = []

        for i in range(n_images):
            try:
                # Generate enhanced prompt
                enhanced_prompt = self.generate_prompt(base_prompt, model=openrouter_model)
                print(f"Generated prompt {i+1}: {enhanced_prompt}")

                # Create image
                image = self.create_image(enhanced_prompt, hf_model=hf_model)

                # Save image
                filepath = self.save_image(image, enhanced_prompt)
                generated_files.append(filepath)

                print(f"Successfully generated and saved image {i+1} to: {filepath}")

            except Exception as e:
                print(f"Error generating image {i+1}: {str(e)}")

        return generated_files

# Usage example
if __name__ == "__main__":
    # Configuration
    OPENROUTER_API_KEY = "MASK_1"
    HF_TOKEN = "MASK_2"
    OUTPUT_FOLDER = "kuvat/4"

    # Initialize generator
    generator = ImageGenerator(OPENROUTER_API_KEY, HF_TOKEN, OUTPUT_FOLDER)

    # Generate images
    base_prompt = "Make a unique and creative image prompt for a poster about \"BPR WIARD\" and billiards/pool. Do not say anything except for the prompt."
    n_images = 3
    openrouter_model = "qwen/qwen-2.5-72b-instruct"  # or any other available model
    hf_model = "black-forest-labs/FLUX.1-schnell"

    generated_files = generator.generate_batch(
        base_prompt=base_prompt,
        n_images=n_images,
        openrouter_model=openrouter_model,
        hf_model=hf_model
    )

    print("\nGenerated files:")
    for file in generated_files:
        print(file)

Instruction: on error try again in 61 seconds

import pytest
import inspect
from unittest.mock import MagicMock, patch, call
from PIL import Image
import re

class TestRetryOnError:

    @pytest.fixture
    def mock_image(self):
        """Create a mock image for testing"""
        mock_img = MagicMock(spec=Image.Image)
        mock_img.save = MagicMock()
        return mock_img

    def test_create_image_has_retry_mechanism(self, implementation):
        """Test that create_image method contains a retry mechanism"""
        impl_name, module = implementation

        # Get the source code of create_image method
        image_generator_class = module.ImageGenerator
        create_image_source = inspect.getsource(image_generator_class.create_image)

        # Check for retry-related code
        has_retry = False
        # Look for retry patterns without using exact string matching
        if (
            "try" in create_image_source
            and "except" in create_image_source
            and ("sleep" in create_image_source or "time.sleep" in create_image_source)
        ):
            has_retry = True

        assert (
            has_retry
        ), f"Implementation {impl_name} does not include a retry mechanism in create_image"

    def test_retry_time_delay_is_61_seconds(self, implementation):
        """Test that the retry time delay is 61 seconds as specified in the requirements"""
        impl_name, module = implementation

        # Get the source code of create_image method
        image_generator_class = module.ImageGenerator
        create_image_source = inspect.getsource(image_generator_class.create_image)

        # Check for time.sleep with 61 seconds, allowing for different ways it might be written
        has_sleep_61 = False
        if "sleep(61)" in create_image_source:
            has_sleep_61 = True
        elif "sleep" in create_image_source and "61" in create_image_source:
            # Find lines containing sleep and 61
            lines = create_image_source.split("\n")
            for line in lines:
                if "sleep" in line and "61" in line:
                    has_sleep_61 = True
                    break

        assert (
            has_sleep_61
        ), f"Implementation {impl_name} does not wait 61 seconds before retrying"

    def test_time_module_imported(self, implementation):
        """Test that the time module is imported"""
        impl_name, module = implementation

        # First try to check if time module is directly imported in the code
        module_source = inspect.getsource(module)

        # Check for time import in different formats
        has_time_import = False
        if re.search(r"import\s+time", module_source):
            has_time_import = True
        elif re.search(r"from\s+time\s+import", module_source):
            has_time_import = True

        # Even if not found in the pattern above, see if it's accessible in the module
        try:
            # First dynamically patch the module with time if missing
            if not has_time_import:
                import time

                if not hasattr(module, "time"):
                    setattr(module, "time", time)

            # Run a simple test that requires time module
            with patch.object(module, "time") as mock_time:
                # If this doesn't raise an exception, time is accessible
                has_time_import = True
        except:
            # Failed to patch or use time module
            has_time_import = False

        assert (
            has_time_import
        ), f"Implementation {impl_name} does not have access to the time module"

    def test_create_image_actually_retries(self, implementation, mock_image):
        """Test that create_image method actually retries on error"""
        impl_name, module = implementation

        # First ensure time module is available to the implementation
        import time

        if not hasattr(module, "time"):
            setattr(module, "time", time)

        # Patch time.sleep to avoid real delays
        with patch.object(module, "time") as mock_time_module:
            mock_time_module.sleep = MagicMock()

            # Setup retry testing infrastructure
            with patch.object(module, "InferenceClient") as mock_inference_client:
                # Setup client instance mock
                client_instance = MagicMock()
                mock_inference_client.return_value = client_instance

                # Configure the mock to fail once then succeed
                client_instance.text_to_image.side_effect = [
                    Exception("Simulated error"),  # First attempt fails
                    mock_image,  # Second attempt succeeds
                ]

                # Create the generator and execute the method to test
                generator = module.ImageGenerator(
                    "fake_key", "fake_token", "fake_folder"
                )

                # Call the method under test
                result = generator.create_image("test prompt")

                # Verify retry occurred
                assert (
                    client_instance.text_to_image.call_count == 2
                ), f"Implementation {impl_name} did not retry after error"
                assert (
                    mock_time_module.sleep.called
                ), f"Implementation {impl_name} did not sleep between retries"

                # Check that sleep was called with 61 seconds
                mock_time_module.sleep.assert_called_with(61)

                # Verify the result is the mock image
                assert (
                    result == mock_image
                ), f"Implementation {impl_name} did not return the image after successful retry"

pytest
pytest-mock
Pillow
requests
huggingface_hub

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