import tkinter as tk
from tkinter import ttk
from tkinter import colorchooser

def create_paint_canvas():
    root = tk.Tk()
    root.title("Painting Canvas")

    # Create toolbar frame
    toolbar = tk.Frame(root, bd=1, relief=tk.RAISED)
    toolbar.pack(side=tk.TOP, fill=tk.X)

    # Brush size control
    size_label = tk.Label(toolbar, text="Brush Size:")
    size_label.pack(side=tk.LEFT, padx=5)
    brush_size = tk.Scale(toolbar, from_=1, to=20, orient=tk.HORIZONTAL)
    brush_size.set(2)  # Default size
    brush_size.pack(side=tk.LEFT, padx=5)

    # Current color display and color picker
    current_color = tk.StringVar(value="black")
    color_display = tk.Label(toolbar, bg=current_color.get(), width=3)
    color_display.pack(side=tk.LEFT, padx=5)

    def choose_color():
        color = colorchooser.askcolor(title="Choose brush color")[1]
        if color:  # If a color was chosen (not cancelled)
            current_color.set(color)
            color_display.config(bg=color)

    color_btn = tk.Button(toolbar, text="Choose Color", command=choose_color)
    color_btn.pack(side=tk.LEFT, padx=5)

    # Add this after the existing color picker button
    bg_color = tk.StringVar(value="white")  # Store current background color
    bg_display = tk.Label(toolbar, bg=bg_color.get(), width=3)
    bg_display.pack(side=tk.LEFT, padx=5)

    def ereaser():
        current_color.set(bg_color.get())  # Set brush color to background color
        color_display.config(bg=bg_color.get())

    #ereaser button
    eraser_btn = tk.Button(toolbar, text="Eraser", command=ereaser)
    eraser_btn.pack(side=tk.LEFT, padx=5)

    def choose_background():
        color = colorchooser.askcolor(title="Choose background color")[1]
        if color:
            bg_color.set(color)
            bg_display.config(bg=color)
            canvas.config(bg=color)

    bg_btn = tk.Button(toolbar, text="Background Color", command=choose_background)
    bg_btn.pack(side=tk.LEFT, padx=5)

    # Create canvas
    canvas = tk.Canvas(root, bg="white", width=800, height=600)
    canvas.pack(expand=tk.YES, fill=tk.BOTH)

    def clear_canvas():
        canvas.delete("all")  # Removes all drawings from the canvas

    # Clear canvas button
    clear_btn = tk.Button(toolbar, text="Clear Canvas", command=clear_canvas)
    clear_btn.pack(side=tk.LEFT, padx=5)

    def paint(event):
        size = brush_size.get()  # Get current brush size
        x1, y1 = (event.x - size), (event.y - size)  # Calculate top-left corner of oval
        x2, y2 = (event.x + size), (event.y + size)  # Calculate bottom-right corner of oval
        canvas.create_oval(x1, y1, x2, y2, fill=current_color.get(), outline=current_color.get())  # Draw oval on canvas with current color

    canvas.bind("<B1-Motion>", paint)

    root.mainloop()

if __name__ == "__main__":
    create_paint_canvas()

Instruction: add a button that changes the background to a random color

import pytest
import inspect
import random
import re

def test_imports_random(implementation):
    """Test that the implementation imports the random module."""
    impl_name, module = implementation

    # Get the source of the implementation
    module_source = inspect.getsource(module)

    # Check if random is imported - expanded patterns to catch more variations
    imports_random = any(pattern in module_source for pattern in [
        "import random",
        "from random import",
        "import random as"
    ])

    # Skip checking the original code since it's not supposed to have this feature
    if impl_name == "original_code":
        pytest.skip("Original code doesn't need the random module")
    else:
        assert imports_random, f"{impl_name} should import the random module for generating random colors."

def test_random_background_button_exists(implementation):
    """Test that the implementation has a random background button."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't need a random background button")
        return

    # Get the source of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Look for a button with a name that suggests it's for random background
    has_random_bg_button = False
    button_patterns = [
        r'Button\(.*[Rr]andom.*[Bb]ackground',
        r'Button\(.*text=["\']\s*Random\s*Background\s*["\']',
        r'Button\(.*text=["\']\s*Random\s*BG\s*["\']',
        r'rand.*_btn\s*=\s*.*Button\(',
        r'random_bg.*btn\s*=\s*.*Button\('
    ]

    for pattern in button_patterns:
        if re.search(pattern, func_source):
            has_random_bg_button = True
            break

    assert has_random_bg_button, f"{impl_name} should have a button for random background color."

def test_random_background_function_exists(implementation):
    """Test that the implementation has a function to set random background."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't need a random background function")
        return

    # Get the source of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Look for a function that generates random background colors
    has_random_bg_function = False
    function_patterns = [
        r'def\s+random_background',
        r'def\s+random_background_color',
        r'def\s+set_random_background',
        r'def\s+random_bg',
        r'def\s+rand.*_background'
    ]

    # Also look for lambda functions or anonymous functions that might be directly assigned to a command
    lambda_patterns = [
        r'command\s*=\s*lambda.*random\.randint',
        r'command\s*=\s*lambda.*random.*color'
    ]

    for pattern in function_patterns:
        if re.search(pattern, func_source):
            has_random_bg_function = True
            break

    # If named function not found, check for lambda implementation
    if not has_random_bg_function:
        for pattern in lambda_patterns:
            if re.search(pattern, func_source):
                has_random_bg_function = True
                break

    assert has_random_bg_function, f"{impl_name} should have a function to set the background to a random color."

def test_random_color_generation(implementation):
    """Test that the implementation generates random colors correctly."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't implement random color generation")
        return

    # Get the source code of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Check for correct random color generation patterns - expanded for more variations
    valid_random_color_patterns = [
        r'random\.randint\(0,\s*0xFFFFFF\)',
        r'random\.randint\(0,\s*16777215\)',
        r'f"#{random\.randint\(0,\s*0xFFFFFF\):06x}"',
        r'f"#{random\.randint\(0,\s*16777215\):06x}"',
        r'"#{:06x}"\.format\(random\.randint\(0,\s*0xFFFFFF\)\)',
        r'"#{:06x}"\.format\(random\.randint\(0,\s*16777215\)\)',
        r'random\.randint\(0,\s*255\).*random\.randint\(0,\s*255\).*random\.randint\(0,\s*255\)',  # RGB approach
        r'"#%06x".*random\.randint',  # Alternative string formatting
        r'random\.choice\(\[\s*["\'](#[0-9A-Fa-f]{6})["\']',  # Predefined color list approach
        r'random\.random\(\).*255'  # Using random.random() * 255 approach
    ]

    has_valid_color_gen = False
    for pattern in valid_random_color_patterns:
        if re.search(pattern, func_source):
            has_valid_color_gen = True
            break

    # If no specific pattern found, look for any random color generation attempt
    if not has_valid_color_gen:
        # Look for any usage of random in the context of color generation
        general_random_color_pattern = r'random\.(?:randint|random|choice).*(?:color|bg|background)'
        has_valid_color_gen = re.search(general_random_color_pattern, func_source) is not None

    assert has_valid_color_gen, f"{impl_name} should generate random colors in a valid hex format."

def test_button_updates_bg_display(implementation):
    """Test that the random background button updates the background color display."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't implement random background feature")
        return

    # Get the source code of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Check for code that updates the bg_display in the random background function
    updates_display_patterns = [
        r'bg_display\.config\(bg=.*\)',
        r'bg_display\.configure\(bg=.*\)',
        r'bg_display\[["\'](background|bg)["\'].*='
    ]

    updates_display = any(re.search(pattern, func_source) for pattern in updates_display_patterns)

    assert updates_display, f"{impl_name} should update the background color display when random color is selected."

def test_button_updates_canvas_bg(implementation):
    """Test that the random background button updates the canvas background."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't implement random background feature")
        return

    # Get the source code of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Check for code that updates the canvas background in the random background function
    updates_canvas_patterns = [
        r'canvas\.config\(bg=.*\)',
        r'canvas\.configure\(bg=.*\)',
        r'canvas\[["\'](background|bg)["\'].*='
    ]

    updates_canvas = any(re.search(pattern, func_source) for pattern in updates_canvas_patterns)

    assert updates_canvas, f"{impl_name} should update the canvas background when random color is selected."

def test_button_updates_bg_color_var(implementation):
    """Test that the random background button updates the bg_color StringVar."""
    impl_name, module = implementation

    # Skip checking the original code
    if impl_name == "original_code":
        pytest.skip("Original code doesn't implement random background feature")
        return

    # Get the source code of the create_paint_canvas function
    create_paint_canvas_func = module.create_paint_canvas
    func_source = inspect.getsource(create_paint_canvas_func)

    # Check for code that updates the bg_color StringVar in the random background function
    updates_var_patterns = [
        r'bg_color\.set\(.*\)',
        r'bg_color\.delete\(0,\s*tk\.END\).*insert',  # For Entry widgets
        r'bg_color\s*=\s*.*random'  # Direct assignment
    ]

    updates_var = any(re.search(pattern, func_source) for pattern in updates_var_patterns)

    assert updates_var, f"{impl_name} should update the bg_color StringVar when random color is selected."

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


from beem.discussions import Discussions, Query
from beem.comment import Comment

n_respuestas_minimas = 5
diccionario = {}

def procesar (texto: str):
    return "count me " in texto
def is_own_author (autor: str):
    return author == 'subidu'
def is_banned (autor: str):
    list_banned = []
    return autor in list_banned
def generar_permlink_unico () -> str:
    return "".join(random.choices(string.digits, k=10))
def procesar_replies (replies: Comment):
    pass
def preparar_comentario (parent_author: str, parent_permlink: str, permlink: str, title: str = '', author: str = 'subidu' , body: str = 'Count me in ^^ @subidu') -> dict[str:str]:
    return {
        "parent_author": parent_author,
        "parent_permlink": parent_permlink,
        "author": author,
        "permlink": permlink,
        "title": title,
        "body": body,
    }

q = Query()
d = Discussions()
posts_generator = d.get_discussions("created", q, limit=6000)
X = 0

for post in posts_generator:
    post_author = post['author']
    post_permlink = post['permlink']
    post_replies = post['children']
    cnt = 0
    X += 1
    if post_replies > n_respuestas_minimas:
        comment = Comment(authorperm=f"{post_author}/{post_permlink}")
        post_replies :list = comment.get_replies()

        cnt = 0
        for replies in post_replies:

            author = replies['author']
            text = replies['body']
            if is_own_author(author):
                # Reevaluar el comentario
                break
            if is_banned(author):
                break
            if procesar(text):
                cnt+= 1
    if cnt > 3:
        print("Iterador: ",X)
        print(replies['author'],'/',replies['permlink'])

Instruction: Quiero añadir una funcion para crear un diccionario jerarquico por posts_generatos y replies

import pytest
import inspect
from unittest.mock import patch

@patch('beem.discussions.Discussions')
@patch('beem.comment.Comment')
def test_function_returns_dictionary(mock_comment, mock_discussions, implementation):
    """Test if the function returns a dictionary"""
    impl_name, module = implementation

    # Set up mocks
    mock_instance = mock_discussions.return_value
    mock_instance.get_discussions.return_value = []

    mock_input = []
    # Search for any function that returns a dictionary
    for name, func in inspect.getmembers(module, inspect.isfunction):
        try:
            sig = inspect.signature(func)
            if len(sig.parameters) == 0:
                result = func()
            else:
                result = func(mock_input)
        except Exception:
            continue  # Skip functions that raise errors

        if isinstance(result, dict):
            # ✅ Found a function that returns a dictionary
            return

    # ❌ No function returned a dictionary
    assert False, f"{impl_name} has no function that returns a dictionary given mock discussion input"

@patch('beem.discussions.Discussions')
@patch('beem.comment.Comment')
def test_hierarchical_structure(mock_comment, mock_discussions, implementation):
    """Test if the function creates a hierarchical structure with posts and replies"""
    impl_name, module = implementation

    # Create mock post data
    mock_post = {
        'author': 'author1',
        'permlink': 'permlink1',
        'children': 10  # More than n_respuestas_minimas
    }

    # Create mock replies
    mock_replies = [
        {'author': 'user1', 'permlink': 'reply1', 'body': 'test reply'},
        {'author': 'user2', 'permlink': 'reply2', 'body': 'count me in test'}
    ]

    # Set up mock for Discussions and get_discussions
    mock_discussions_instance = mock_discussions.return_value
    mock_discussions_instance.get_discussions.return_value = [mock_post]

    # Set up mock for Comment
    mock_comment_instance = mock_comment.return_value
    mock_comment_instance.get_replies.return_value = mock_replies

    # Try each function in the module
    for name, func in inspect.getmembers(module, inspect.isfunction):
        try:
            sig = inspect.signature(func)
            if len(sig.parameters) == 0:
                result = func()
            else:
                result = func(mock_discussions_instance.get_discussions.return_value)
        except Exception:
            continue  # Skip functions that raise

        # --- Validate structure ---
        if isinstance(result, dict) and len(result) > 0:
            for key, value in result.items():
                if isinstance(value, dict):
                    # ✅ Found nested dictionary — implies hierarchy
                    return

    # ❌ No valid function found
    assert False, f"{impl_name} has no function that creates a hierarchical dictionary"

pytest
pytest-mock
pytest-cov
cryptography
beem

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 os
import random
import torch
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.metrics import precision_score, recall_score
from torch.nn import functional as F
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
from colpali_engine.interpretability import (
    get_similarity_maps_from_embeddings,
    plot_all_similarity_maps,
)

# Path to extracted Flickr8k dataset
FLICKR8K_IMAGES_PATH = "flickr8k/Images"
FLICKR8K_CAPTIONS_PATH = "flickr8k/captions.txt"

# Function to load image-text pairs from Flickr8k
def load_flickr8k_data(images_path, captions_path, fraction=0.1):
    # Read captions file
    with open(captions_path, "r") as f:
        captions_data = f.readlines()[1:]  # Skip header

    # Parse captions
    image_text_pairs = {}
    for line in captions_data:
        image_name, caption = line.strip().split(",", 1)
        if image_name not in image_text_pairs:
            image_text_pairs[image_name] = []
        image_text_pairs[image_name].append(caption)

    # Load only a fraction of the dataset
    selected_images = random.sample(list(image_text_pairs.keys()), int(len(image_text_pairs) * fraction))
    image_text_pairs = {k: image_text_pairs[k] for k in selected_images}

    # Create pairs of images and captions
    pairs = []
    for image_name, captions in image_text_pairs.items():
        image_path = os.path.join(images_path, image_name)
        if os.path.exists(image_path):
            pairs.append((Image.open(image_path), random.choice(captions)))
    return pairs

# Function to create unrelated pairs
def create_unrelated_pairs(image_text_pairs):
    """
    Creates unrelated pairs of images and texts by randomly shuffling the texts.

    Args:
        image_text_pairs (list): A list of tuples containing images and their corresponding texts.

    Returns:
        list: A list of tuples containing images and unrelated texts.
    """
    images, texts = zip(*image_text_pairs)
    unrelated_texts = random.sample(texts, len(texts))
    return list(zip(images, unrelated_texts))

def create_visual_pairs(image_text_pairs):
    """
    Creates pairs of original and augmented images from image-text pairs.

    This function takes a list of image-text pairs and creates new pairs consisting
    of the original images and their augmented versions. The augmentation used
    in this implementation is a horizontal flip.

    Args:
        image_text_pairs (list): A list of tuples containing (image, text) pairs,
            where images are PIL Image objects and texts are strings.

    Returns:
        list: A list of tuples containing (original_image, augmented_image) pairs,
            where both elements are PIL Image objects.
    """
    from torchvision.transforms import ToTensor
    images, _ = zip(*image_text_pairs)
    augmented_images = [ToTensor()(image).flip(-1) for image in images]  # Example augmentation: horizontal flip
    return list(zip(images, augmented_images))

def get_embeddings(images, texts, model_id="google/siglip-base-patch16-224"):
    """
    Given lists of images and texts, returns normalized embeddings for both.
    """
    # Ensure texts is a list of strings
    if not all(isinstance(t, str) for t in texts):
        raise ValueError("All text inputs must be strings.")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = AutoModel.from_pretrained(model_id, ignore_mismatched_sizes=True).to(device)
    processor = AutoProcessor.from_pretrained(model_id)

    # Preprocess images and texts
    image_inputs = processor(images=images, return_tensors="pt").to(device)
    text_inputs = processor(text=texts, return_tensors="pt", padding="max_length").to(device)

    with torch.no_grad():
        image_embeds = model.get_image_features(**image_inputs)
        text_embeds = model.get_text_features(**text_inputs)

    # Normalize embeddings
    image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
    text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)

    return image_embeds, text_embeds

def cosine_similarity_analysis(embeddings1, embeddings2, title):
    """
    Computes cosine similarity for matching and unrelated pairs and compares distributions.
    """
    similarities = cosine_similarity(embeddings1.cpu().numpy(), embeddings2.cpu().numpy())

    # Matching pairs: Diagonal of the similarity matrix
    matching_similarities = np.diag(similarities)

    # Unrelated pairs: Off-diagonal similarities
    unrelated_similarities = similarities[~np.eye(similarities.shape[0], dtype=bool)]

    print(f"### {title} ###")
    print(f"Mean Matching Similarity: {np.mean(matching_similarities):.4f}")
    print(f"Mean Unrelated Similarity: {np.mean(unrelated_similarities):.4f}")
    print()

    # Plot distributions
    plt.figure(figsize=(10, 6))
    sns.histplot(matching_similarities, kde=True, label="Matching Pairs", color="blue", bins=30)
    sns.histplot(unrelated_similarities, kde=True, label="Unrelated Pairs", color="red", bins=30)
    plt.title(f"{title}: Cosine Similarity Distributions")
    plt.xlabel("Cosine Similarity")
    plt.ylabel("Frequency")
    plt.legend()
    plt.show()

### b. Nearest-Neighbor Retrieval
def retrieval_metrics(query_embeds, target_embeds, ground_truth_indices, k=5):
    """
    Computes Precision@k and Recall@k for nearest-neighbor retrieval.

    This function evaluates the effectiveness of retrieval by calculating Precision@k and Recall@k.
    Precision@k measures the accuracy of the top-k retrieved items, while Recall@k measures the ability
    to find the relevant item within the top-k retrieved items.  It assumes there's only one true
    match per query.

    Args:
        query_embeds (torch.Tensor): Embeddings of the query data.
        target_embeds (torch.Tensor): Embeddings of the target data (database).
        ground_truth_indices (list): List of indices in the target data representing the true matches for each query.
        k (int): The number of top results to consider.

    Returns:
        tuple: A tuple containing mean Precision@k and mean Recall@k.
    """
    similarities = cosine_similarity(query_embeds.cpu().numpy(), target_embeds.cpu().numpy())
    sorted_indices = np.argsort(-similarities, axis=1)[:, :k]  # Top-k indices

    # Compute metrics
    precisions = []
    recalls = []
    for i, true_idx in enumerate(ground_truth_indices):
        retrieved_indices = sorted_indices[i]
        true_positives = int(true_idx in retrieved_indices)
        precisions.append(true_positives / k)
        recalls.append(true_positives / 1)  # Only one true match per query

    mean_precision = np.mean(precisions)
    mean_recall = np.mean(recalls)

    return mean_precision, mean_recall

def plot_query_token_importance(
    pil_image,
    similarity_maps,
    query_tokens,
    alpha: float = 0.5
) -> None:
    """
    Plot a separate heatmap for each query token in the similarity_maps.

    Args:
        pil_image (PIL.Image.Image): The original image (e.g., loaded via Image.open(...)).
        similarity_maps (torch.Tensor):
            Shape = (num_query_tokens, n_patches_x, n_patches_y).
        query_tokens (List[str]): A list of strings for each token in the query.
        alpha (float): Transparency for the heatmap overlays (0=transparent, 1=opaque).
    """
    # Convert PIL to numpy
    image_np = np.array(pil_image)
    H, W = image_np.shape[:2]

    num_tokens = similarity_maps.size(0)
    assert num_tokens == len(query_tokens), (
        f"The number of query tokens in similarity_maps ({num_tokens}) "
        f"doesn't match the length of query_tokens list ({len(query_tokens)})."
    )

    fig, axs = plt.subplots(1, num_tokens, figsize=(5 * num_tokens, 5))
    if num_tokens == 1:
        # If there's only one token, axs won't be an iterable
        axs = [axs]

    for idx in range(num_tokens):
        # Each similarity_map for a single query token: shape = (n_patches_x, n_patches_y)
        single_map = similarity_maps[idx]  # (n_patches_x, n_patches_y)

        # Upsample to full image size
        single_map_4d = single_map.unsqueeze(0).unsqueeze(0)  # (1,1,n_patches_x, n_patches_y)
        upsampled = F.interpolate(
            single_map_4d,
            size=(H, W),
            mode='bilinear',
            align_corners=False
        )

        # .to(torch.float32) fix if your map is bfloat16
        heatmap = upsampled.squeeze().to(torch.float32).cpu().numpy()  # (H, W)

        # Optionally normalize heatmap (uncomment if desired)
        # heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)

        # Plot
        axs[idx].imshow(image_np, cmap=None if image_np.ndim == 3 else 'gray')
        axs[idx].imshow(heatmap, cmap='jet', alpha=alpha)
        axs[idx].set_title(f"Query: {query_tokens[idx]}")
        axs[idx].axis('off')

    plt.tight_layout()
    plt.show()

def get_maps_and_embeds(batch_images, batch_queries, model, processor, image, use_qwen=False):
    """
    Gets similarity maps and embeddings from batched images and queries using a given model and processor.

    This function processes batched images and queries through a model to obtain embeddings and
    similarity maps between them. It handles the computation of image masks and patch-based
    similarity calculations.

    Args:
        batch_images: Batched image inputs processed by the processor
        batch_queries: Batched query inputs processed by the processor
        model: The model to use for computing embeddings
        processor: The processor used for image/text preprocessing

    Returns:
        tuple: A tuple containing:
            - original_maps (torch.Tensor): Similarity maps between images and queries
                with shape (query_length, n_patches_x, n_patches_y)
            - original_image_embeddings: Embeddings of the input images
            - original_query_embeddings: Embeddings of the input queries
    """
    with torch.no_grad():
        original_image_embeddings = model.forward(**batch_images)
        original_query_embeddings = model.forward(**batch_queries)
    if use_qwen:
        n_patches = processor.get_n_patches(image_size=image.size, patch_size=model.patch_size, spatial_merge_size=model.spatial_merge_size)
    else:
        n_patches = processor.get_n_patches(image_size=image.size, patch_size=model.patch_size)
    image_mask = processor.get_image_mask(batch_images)

    # Compute original similarity maps
    original_batched_maps = get_similarity_maps_from_embeddings(
        image_embeddings=original_image_embeddings,
        query_embeddings=original_query_embeddings,
        n_patches=n_patches,
        image_mask=image_mask,
    )
    original_maps = original_batched_maps[0]  # (query_length, n_patches_x, n_patches_y)
    return original_maps, original_image_embeddings, original_query_embeddings

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import torch

def visualize_token_map(image, original_maps, token_list, token_index=2, cmap="Greens"):
    """
    Visualize the raw image, raw map, and an overlay of the image with the resized map
    for a specific token.

    Args:
        image (PIL.Image): The input image.
        original_maps (list or tensor): A collection of maps to select from.
        token_list (list): A list of tokens corresponding to the maps.
        token_index (int, optional): The index of the token to visualize. Default is 2.
        cmap (str, optional): The colormap to use for visualizing the map. Default is "Greens".
    """
    # Convert the image to a NumPy array
    image_np = np.array(image)

    # Select the map corresponding to the token
    visual_map = original_maps[token_index]

    # Convert visual_map to NumPy array if it's a tensor
    if isinstance(visual_map, torch.Tensor):
        visual_map = visual_map.cpu().to(dtype=torch.float32).numpy()
    elif not isinstance(visual_map, np.ndarray):
        visual_map = np.array(visual_map)

    # Convert map to a PIL image
    visual_map_pil = Image.fromarray(visual_map)

    # Resize using NEAREST to keep "big pixels"
    visual_map_pil = visual_map_pil.resize(
        (image_np.shape[1], image_np.shape[0]),  # (width, height)
        resample=Image.NEAREST
    )

    # Convert back to NumPy
    resized_map = np.array(visual_map_pil)

    # Create a figure with subplots
    fig, axes = plt.subplots(1, 3, figsize=(15, 6))

    # Display the raw image
    axes[0].imshow(image_np)
    axes[0].set_title("Raw Image")
    axes[0].axis("off")
    # Display the raw map with annotations
    im = axes[1].imshow(visual_map, cmap=cmap)
    axes[1].set_title("Raw Map")
    axes[1].axis("off")

    # Annotate the heatmap
    for i in range(visual_map.shape[0]):
        for j in range(visual_map.shape[1]):
            text = axes[1].text(j, i, f"{visual_map[i, j]:.2f}",
                           ha="center", va="center", color="w" if visual_map[i, j] > visual_map.max() / 2 else "black")

    # Display the overlay plot
    axes[2].imshow(image_np, alpha=1)
    axes[2].imshow(resized_map, cmap=cmap, alpha=0.6)
    axes[2].set_title("Overlay: Image + Map")
    axes[2].axis("off")

    # Add a colorbar for the overlay
    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap), ax=axes[2], shrink=0.8, orientation="vertical")
    cbar.set_label("Map Intensity")

    # Add a title with the token name
    plt.suptitle(f"Token: {token_list[token_index]}")

    # Adjust layout and show
    plt.tight_layout()
    plt.show()

Instruction: modify the cmap so the displayed values are the same as the text displayed on the raw map.

# test_visualize_token_map_no_gui.py

import pytest
import numpy as np
import torch
from PIL import Image
import matplotlib
# Use a non-interactive backend to prevent GUI windows during tests
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from unittest.mock import MagicMock, patch

def get_simple_test_input():
    """Tiny 2×2 map and dummy image/tokens for testing."""
    img = Image.new("RGB", (4, 4), color="white")
    # single‑token map: shape (1, 2, 2)
    maps = torch.tensor([[[0.10, 0.20],
                          [0.30, 0.40]]], dtype=torch.float32)
    tokens = ["only_token"]
    idx = 0
    return img, maps, tokens, idx

@pytest.fixture
def viz_fn(implementation):
    """Grab visualize_token_map from the tested module or skip."""
    impl_name, module = implementation
    if not hasattr(module, "visualize_token_map"):
        pytest.skip(f"{impl_name}: no visualize_token_map found")
    return getattr(module, "visualize_token_map")

@patch("matplotlib.pyplot.show")          # prevent any show() calls
@patch("matplotlib.pyplot.subplots")
def test_colorbar_attached_to_raw_map_mappable(mock_subplots, mock_show, viz_fn):
    """
    The colorbar must be created from the mappable returned by the raw‑map imshow,
    without spinning up any GUI.
    """
    # Arrange: stub out subplots
    fig = MagicMock()
    axes = [MagicMock(), MagicMock(), MagicMock()]
    mock_subplots.return_value = (fig, axes)

    img, maps, tokens, idx = get_simple_test_input()

    # Act
    viz_fn(img, maps, tokens, token_index=idx, cmap="plasma")

    # The raw‑map imshow returns an AxesImage
    im_obj = axes[1].imshow.return_value

    # Assert: colorbar called with that mappable on axes[2]
    fig.colorbar.assert_called_once_with(
        im_obj,
        ax=axes[2],
        shrink=0.8,
        orientation="vertical"
    )

@patch("matplotlib.pyplot.show")          # prevent any show() calls
@patch("matplotlib.pyplot.subplots")
def test_annotation_text_matches_data(mock_subplots, mock_show, viz_fn):
    """
    Each cell in the raw map must be annotated with its exact value (.2f),
    and no GUI window should pop up.
    """
    # Arrange: stub out subplots
    fig = MagicMock()
    axes = [MagicMock(), MagicMock(), MagicMock()]
    mock_subplots.return_value = (fig, axes)

    img, maps, tokens, idx = get_simple_test_input()
    vm = maps[idx].cpu().numpy().flatten()
    n_cells = vm.size

    # Act
    viz_fn(img, maps, tokens, token_index=idx, cmap="Greens")

    # Gather all text() calls on axes[1]
    calls = axes[1].text.call_args_list
    assert len(calls) == n_cells, f"Expected {n_cells} annotations, got {len(calls)}"

    # Verify each annotation string matches data
    expected = [f"{val:.2f}" for val in vm]
    actual = [call.args[2] for call in calls]  # text(x, y, string, ...)
    assert actual == expected, f"Annotations {actual} do not match expected {expected}"

numpy
matplotlib
torch
pytest
pytest-mock
Pillow
seaborn
scikit-learn
colpali_engine
einops

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 math
class Value:
    """ stores a single scalar value and its gradient """

    def __init__(self, data, _children=(), _op=''):
        self.data = data
        self.grad = 0
        # internal variables used for autograd graph construction
        self._backward = lambda: None
        self._prev = set(_children)
        self._op = _op # the op that produced this node, for graphviz / debugging / etc

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data + other.data, (self, other), '+')

        def _backward():
            self.grad += out.grad
            other.grad += out.grad
        out._backward = _backward

        return out

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        out = Value(self.data * other.data, (self, other), '*')

        def _backward():
            self.grad += other.data * out.grad
            other.grad += self.data * out.grad
        out._backward = _backward

        return out

    def __pow__(self, other):
        assert isinstance(other, (int, float)), "only supporting int/float powers for now"
        out = Value(self.data**other, (self,), f'**{other}')

        def _backward():
            self.grad += (other * self.data**(other-1)) * out.grad
        out._backward = _backward

        return out

        print(f'** EXP {self}')
        out = Value(math.exp(self.data), (self,), "e")
        def _backward():
            self.grad += math.exp(self.data) * out.grad
        out._backward = _backward

        return out

    def log(self):
        out = Value(math.log(self.data), (self,), f'log{self.data}')
        def _backward():
            self.grad += 1 / self.data
        out._backward = _backward

        return out

    def relu(self):
        assert isinstance(self.data, (int, float)), "only supporting int/float powers for now"
        out = Value(0 if self.data < 0 else self.data, (self,), 'ReLU')

        def _backward():
            self.grad += (out.data > 0) * out.grad
        out._backward = _backward

        return out

    def softmax(x):
        e_x = np.exp(x - np.max(x))

    def backward(self):

        # topological order all of the children in the graph
        topo = []
        visited = set()
        def build_topo(v):
            if v not in visited:
                visited.add(v)
                for child in v._prev:
                    build_topo(child)
                topo.append(v)
        build_topo(self)

        # go one variable at a time and apply the chain rule to get its gradient
        self.grad = 1
        for v in reversed(topo):
            v._backward()

    def __neg__(self): # -self
        return self * -1
    def __gt__(self, other):
        return self.data > other.data

    def __radd__(self, other): # other + self
        return self + other

    def __sub__(self, other): # self - other
        return self + (-other)

    def __rsub__(self, other): # other - self
        return other + (-self)

    def __rmul__(self, other): # other * self
        return self * other

    def __truediv__(self, other): # self / other
        return self * other**-1

    def __rtruediv__(self, other): # other / self
        return other * self**-1

    def __repr__(self):
        return f"Value(data={self.data}, grad={self.grad})"

Instruction: add a function to differentiate the softmax function

import pytest
import inspect
import math
from typing import List, Any, Callable

# Import numpy safely with fallback
try:
    import numpy as np
except ImportError:
    # Create minimal mock for numpy if not available
    class MockNumpy:
        def exp(self, x):
            if hasattr(x, "__iter__"):
                return [math.exp(v) for v in x]
            return math.exp(x)

        def max(self, x):
            if hasattr(x, "__iter__"):
                return max(x)
            return x

        def sum(self, x, axis=None):
            if hasattr(x, "__iter__"):
                return sum(x)
            return x
    np = MockNumpy()

def extract_value_class(module):
    """
    Extract Value class from module, handling various implementations.
    """
    if hasattr(module, "Value"):
        return module.Value

    # If the module doesn't directly expose Value, try to find it
    for attr_name in dir(module):
        attr = getattr(module, attr_name)
        if isinstance(attr, type) and "Value" in attr.__name__:
            return attr

    # Changed from skip to fail
    pytest.fail(f"Module {module.__name__} doesn't contain a Value class")

def to_list_of_values(module, values):
    """Convert a list of numbers to a list of Value objects for the given module"""
    Value = extract_value_class(module)
    return [Value(v) if not hasattr(v, "data") else v for v in values]

def test_softmax_function_exists(implementation):
    """
    Test that a softmax function is added to the Value class.
    """
    impl_name, module = implementation

    try:
        Value = extract_value_class(module)
        # Check if softmax method exists in the class
        assert hasattr(Value, 'softmax'), f"Implementation {impl_name} does not have a softmax method"

        # Verify it's callable
        assert callable(getattr(Value, 'softmax')), f"Implementation {impl_name} softmax is not callable"
    except (AttributeError, TypeError) as e:
        pytest.fail(f"Implementation {impl_name} test failed: {str(e)}")

def test_softmax_basic_computation(implementation):
    """
    Test that the softmax function performs basic computation correctly.
    """
    impl_name, module = implementation

    try:
        Value = extract_value_class(module)

        # Create a value object
        v = Value(0.0)

        # Make sure numpy is available to the module if it needs it
        if "np" not in dir(module) and "numpy" not in dir(module):
            # Add numpy to the module
            setattr(module, "np", np)

        # Try calling with a simple array
        inputs = [1.0, 2.0, 3.0]
        value_inputs = to_list_of_values(module, inputs)

        # Try different calling conventions
        result = None

        # Approach 1: Static method
        try:
            if hasattr(Value.softmax, '__self__') and Value.softmax.__self__ is Value:
                # It's a class method
                result = Value.softmax(inputs)
        except (TypeError, ValueError, AttributeError):
            try:
                result = Value.softmax(value_inputs)
            except (TypeError, ValueError, AttributeError):
                pass

        # Approach 2: Instance method
        if result is None:
            try:
                result = v.softmax(inputs)
            except (TypeError, ValueError, AttributeError):
                try:
                    result = v.softmax(value_inputs)
                except (TypeError, ValueError, AttributeError):
                    try:
                        result = value_inputs[0].softmax(value_inputs)
                    except (TypeError, ValueError, AttributeError):
                        pytest.fail(f"Implementation {impl_name}: Could not call softmax with any approach")

        # Result validation
        if isinstance(result, list):
            # Check the sum is close to 1
            sum_prob = sum(val.data for val in result)
            assert abs(sum_prob - 1.0) < 1e-6, f"Implementation {impl_name}: Softmax outputs should sum to 1"

            # Check values are in expected order (highest input -> highest output)
            assert result[-1].data > result[0].data, f"Implementation {impl_name}: Softmax should preserve order"
        else:
            # If a single value is returned, check if it's a numpy array
            if hasattr(result.data, "shape") and hasattr(result.data, "sum"):
                # Numpy array result
                assert abs(result.data.sum() - 1.0) < 1e-6, f"Implementation {impl_name}: Softmax outputs should sum to 1"
            else:
                # Single scalar value
                assert 0 <= result.data <= 1, f"Implementation {impl_name}: Softmax output should be a probability"

    except (TypeError, ValueError, AttributeError) as e:
        pytest.fail(f"Implementation {impl_name} failed with error: {str(e)}")

def test_softmax_numerical_stability(implementation):
    """
    Test that the softmax handles large values without numerical overflow.
    """
    impl_name, module = implementation

    try:
        Value = extract_value_class(module)

        # Make sure numpy is available to the module if it needs it
        if "np" not in dir(module) and "numpy" not in dir(module):
            # Add numpy to the module
            setattr(module, "np", np)

        # Create a value object
        v = Value(0.0)

        # Large values that would cause exp overflow if not handled properly
        large_inputs = [100.0, 200.0, 300.0]
        value_inputs = to_list_of_values(module, large_inputs)

        # Try different calling conventions
        result = None

        # Approach 1: Static method
        try:
            if hasattr(Value.softmax, '__self__') and Value.softmax.__self__ is Value:
                result = Value.softmax(large_inputs)
        except (TypeError, ValueError, AttributeError):
            try:
                result = Value.softmax(value_inputs)
            except (TypeError, ValueError, AttributeError):
                pass

        # Approach 2: Instance method
        if result is None:
            try:
                result = v.softmax(large_inputs)
            except (TypeError, ValueError, AttributeError):
                try:
                    result = v.softmax(value_inputs)
                except (TypeError, ValueError, AttributeError):
                    try:
                        result = value_inputs[0].softmax(value_inputs)
                    except (TypeError, ValueError, AttributeError):
                        pytest.fail(f"Implementation {impl_name}: Could not call softmax with any approach")

        # Check if we got a result without overflow errors
        if isinstance(result, list):
            # The largest input should dominate (be close to 1)
            assert abs(result[-1].data - 1.0) < 1e-3, \
                f"Implementation {impl_name}: Largest value should dominate in softmax"
        else:
            # If we got a single Value with numpy array data
            if hasattr(result.data, "__iter__"):
                result_data = result.data
                if hasattr(result_data, "tolist"):  # Handle numpy arrays
                    result_data = result_data.tolist()
                assert abs(result_data[-1] - 1.0) < 1e-3, \
                    f"Implementation {impl_name}: Largest value should dominate in softmax"

    except (TypeError, ValueError, AttributeError) as e:
        pytest.fail(f"Implementation {impl_name} numerical stability test failed: {str(e)}")
    except OverflowError:
        pytest.fail(f"Implementation {impl_name} failed with numerical overflow - not handling large values correctly")

def test_softmax_gradient_computation(implementation):
    """
    Test that the softmax function correctly sets up the backward pass.
    """
    impl_name, module = implementation

    try:
        Value = extract_value_class(module)

        # Make sure numpy is available to the module if it needs it
        if "np" not in dir(module) and "numpy" not in dir(module):
            # Add numpy to the module
            setattr(module, "np", np)

        # Create a value object
        v = Value(0.0)

        # Large values that would cause exp overflow if not handled properly
        inputs = [100.0, 200.0, 300.0]
        value_inputs = to_list_of_values(module, inputs)

        # Try different calling conventions
        result = None

        # Approach 1: Static method
        try:
            if hasattr(Value.softmax, '__self__') and Value.softmax.__self__ is Value:
                result = Value.softmax(inputs)
        except (TypeError, ValueError, AttributeError):
            try:
                result = Value.softmax(value_inputs)
            except (TypeError, ValueError, AttributeError):
                pass

        # Approach 2: Instance method
        if result is None:
            try:
                result = v.softmax(inputs)
            except (TypeError, ValueError, AttributeError):
                try:
                    result = v.softmax(value_inputs)
                except (TypeError, ValueError, AttributeError):
                    try:
                        result = value_inputs[0].softmax(value_inputs)
                    except (TypeError, ValueError, AttributeError):
                        pytest.fail(f"Implementation {impl_name}: Could not call softmax with any approach")

        # Different implementations may return different structures
        if isinstance(result, list):
            # Check that backward is set for each output
            for r in result:
                assert hasattr(r, "_backward"), \
                    f"Implementation {impl_name}: _backward function missing from softmax outputs"

            # Run backward on one of the outputs
            result[0].grad = 1.0
            if hasattr(result[0], "_backward") and callable(result[0]._backward):
                result[0]._backward()

            # If backward propagation is set up but not working yet, don't skip but fail
            grad_sum = sum(abs(v.grad) if hasattr(v, 'grad') else 0 for v in inputs)
            if grad_sum == 0:
                # Changed from skip to fail
                pytest.fail(f"Implementation {impl_name}: Gradient propagation not implemented")

        else:
            # Single output case - less common
            assert hasattr(result, "_backward"), \
                f"Implementation {impl_name}: _backward function not properly set"

            result.grad = 1.0
            if hasattr(result, "_backward") and callable(result._backward):
                result._backward()

    except (TypeError, ValueError, AttributeError) as e:
        pytest.fail(f"Implementation {impl_name} gradient test failed: {str(e)}")

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