class SimpleConvNet1(nn.Module):
    def __init__(self, flattened_size):  # Принимаем flattened_size как аргумент
        super().__init__()

        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 32, 3),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )

        self.flatten = nn.Flatten()
        self.fc1 = nn.Sequential(
            nn.Linear(flattened_size, 512),  # Используем flattened_size здесь
            nn.ReLU(),
            nn.Linear(512, 3)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.flatten(x)
        x = self.fc1(x)
        return x

Instruction: 1. Постройте простую сверточную сеть без использования функций Dropout и BatchNorm.

import pytest
import torch
import inspect
import re
import sys
from typing import Tuple, Any, List, Optional

def test_imports_present(implementation):
    """Ensure that proper imports are present in the implementation."""
    impl_name, module = implementation
    try:
        source_code = inspect.getsource(module)

        required_imports = [
            ('torch.nn', ['import torch.nn as nn', 'from torch import nn'])
        ]

        for pkg, patterns in required_imports:
            if not any(pattern in source_code for pattern in patterns):
                pytest.skip(f"{impl_name}: Test skipped - missing proper import for {pkg}")
    except Exception as e:
        pytest.skip(f"{impl_name}: Error inspecting source code: {e}")

def get_convnet_class(module) -> Optional[type]:
    """Helper function to find the CNN model class in the module.

    Now with improved pattern recognition to detect a wider range of CNN class names.
    """
    try:
        for name, obj in inspect.getmembers(module):
            if inspect.isclass(obj) and hasattr(obj, '__mro__'):
                # Check if torch.nn.Module is in the inheritance chain
                if any('Module' in str(base) for base in obj.__mro__):
                    # Match common CNN naming patterns
                    if (('Conv' in name and 'Net' in name) or
                        ('Simple' in name and 'Conv' in name) or
                        name.startswith('CNN') or
                        name.endswith('CNN') or
                        'SimpleConvNet' in name or
                        'ConvolutionalNetwork' in name or
                        'ConvNet' in name):
                        return obj

                    # Fallback: check if it has conv layers in its structure
                    # This helps identify classes even if they follow non-standard naming
                    try:
                        instance = obj()
                        if hasattr(instance, 'conv1') or hasattr(instance, 'conv2'):
                            return obj
                    except:
                        pass
    except Exception:
        pass

    # Last resort: try to find any Module subclass with conv-like attributes
    try:
        for name, obj in inspect.getmembers(module):
            if inspect.isclass(obj) and hasattr(obj, '__mro__'):
                if any('Module' in str(base) for base in obj.__mro__):
                    return obj  # Return the first nn.Module subclass we find
    except Exception:
        pass

    return None

def test_class_definition(implementation):
    """Ensure the CNN class is properly defined."""
    impl_name, module = implementation

    # Try to find the CNN class
    convnet_class = get_convnet_class(module)

    # Check if any CNN class exists
    assert convnet_class is not None, f"{impl_name}: No CNN class found. Make sure your class inherits from nn.Module and has a typical CNN structure."

    # Check if it's a subclass of nn.Module
    assert any('Module' in str(base) for base in convnet_class.__mro__), f"{impl_name}: CNN class should inherit from nn.Module"

def test_conv_layers_structure(implementation):
    """Test the convolutional layers structure of the model."""
    impl_name, module = implementation

    # Get the CNN class
    convnet_class = get_convnet_class(module)
    if convnet_class is None:
        pytest.skip(f"{impl_name}: No CNN class found")

    # Create a model instance - try with flattened_size parameter first
    try:
        # Use a reasonable default size for a 32x32 input after convolutions
        model = convnet_class(flattened_size=1600)
    except TypeError:
        try:
            model = convnet_class()
        except Exception as e:
            pytest.skip(f"{impl_name}: Failed to create model instance: {e}")

    # Check conv1 structure
    assert hasattr(model, 'conv1'), f"{impl_name}: Missing conv1 layer"
    assert isinstance(model.conv1, torch.nn.Sequential), f"{impl_name}: conv1 should be Sequential"

    # Check components of conv1
    assert len(model.conv1) >= 3, f"{impl_name}: conv1 should have at least 3 components"

    # Find the Conv2d, ReLU, and MaxPool2d layers in conv1
    has_conv2d = False
    has_relu = False
    has_maxpool2d = False

    for layer in model.conv1:
        if isinstance(layer, torch.nn.Conv2d):
            has_conv2d = True
            assert layer.in_channels == 3, f"{impl_name}: conv1 input channels should be 3"
            assert layer.out_channels == 32, f"{impl_name}: conv1 output channels should be 32"
        elif isinstance(layer, torch.nn.ReLU):
            has_relu = True
        elif isinstance(layer, torch.nn.MaxPool2d):
            has_maxpool2d = True

    assert has_conv2d, f"{impl_name}: conv1 should contain a Conv2d layer"
    assert has_relu, f"{impl_name}: conv1 should contain a ReLU layer"
    assert has_maxpool2d, f"{impl_name}: conv1 should contain a MaxPool2d layer"

    # Check conv2 structure
    assert hasattr(model, 'conv2'), f"{impl_name}: Missing conv2 layer"
    assert isinstance(model.conv2, torch.nn.Sequential), f"{impl_name}: conv2 should be Sequential"

    # Check components of conv2
    assert len(model.conv2) >= 3, f"{impl_name}: conv2 should have at least 3 components"

    # Find the Conv2d, ReLU, and MaxPool2d layers in conv2
    has_conv2d = False
    has_relu = False
    has_maxpool2d = False

    for layer in model.conv2:
        if isinstance(layer, torch.nn.Conv2d):
            has_conv2d = True
            assert layer.in_channels == 32, f"{impl_name}: conv2 input channels should be 32"
            assert layer.out_channels == 64, f"{impl_name}: conv2 output channels should be 64"
        elif isinstance(layer, torch.nn.ReLU):
            has_relu = True
        elif isinstance(layer, torch.nn.MaxPool2d):
            has_maxpool2d = True

    assert has_conv2d, f"{impl_name}: conv2 should contain a Conv2d layer"
    assert has_relu, f"{impl_name}: conv2 should contain a ReLU layer"
    assert has_maxpool2d, f"{impl_name}: conv2 should contain a MaxPool2d layer"

def test_flatten_and_fc_layers(implementation):
    """Test the flatten and fully connected layers of the model."""
    impl_name, module = implementation

    # Get the CNN class
    convnet_class = get_convnet_class(module)
    if convnet_class is None:
        pytest.skip(f"{impl_name}: No CNN class found")

    # Create a model instance - try with flattened_size parameter first
    try:
        model = convnet_class(flattened_size=1600)
    except TypeError:
        try:
            model = convnet_class()
        except Exception as e:
            pytest.skip(f"{impl_name}: Failed to create model instance: {e}")

    # Check flatten layer
    assert hasattr(model, 'flatten'), f"{impl_name}: Missing flatten layer"
    assert isinstance(model.flatten, torch.nn.Flatten), f"{impl_name}: flatten should be Flatten"

    # Check fc1 layer
    assert hasattr(model, 'fc1'), f"{impl_name}: Missing fc1 layer"

    # The fc1 can be either Sequential or just a Linear layer
    if isinstance(model.fc1, torch.nn.Sequential):
        # Find Linear layers in fc1
        linear_layers = [layer for layer in model.fc1 if isinstance(layer, torch.nn.Linear)]
        assert len(linear_layers) > 0, f"{impl_name}: fc1 should contain at least one Linear layer"

        # Find the last Linear layer for output
        last_linear = linear_layers[-1]
        assert last_linear.out_features == 3, f"{impl_name}: Final Linear layer out features should be 3"
    else:
        # If fc1 is not Sequential, check if there are individual fc layers
        assert isinstance(model.fc1, torch.nn.Linear), f"{impl_name}: fc1 should be Linear or Sequential"

        # Check if there's an fc2 layer (common pattern)
        if hasattr(model, 'fc2'):
            assert isinstance(model.fc2, torch.nn.Linear), f"{impl_name}: fc2 should be Linear"
            assert model.fc2.out_features == 3, f"{impl_name}: fc2 out features should be 3"
        else:
            # If no fc2, then fc1 should output 3 features
            assert model.fc1.out_features == 3, f"{impl_name}: fc1 out features should be 3 when no fc2 exists"

def test_no_dropout_or_batchnorm(implementation):
    """Ensure that the model doesn't use Dropout or BatchNorm as per requirements."""
    impl_name, module = implementation

    try:
        # Get the source code of the module
        source_code = inspect.getsource(module)

        # Check for absence of Dropout and BatchNorm
        assert "Dropout" not in source_code, f"{impl_name}: Dropout should not be used as per requirements"
        assert "BatchNorm" not in source_code, f"{impl_name}: BatchNorm should not be used as per requirements"

        # Get the CNN class
        convnet_class = get_convnet_class(module)
        if convnet_class is None:
            pytest.skip(f"{impl_name}: No CNN class found")

        # Create a model instance - try with flattened_size parameter first
        try:
            model = convnet_class(flattened_size=1600)
        except TypeError:
            try:
                model = convnet_class()
            except Exception as e:
                pytest.skip(f"{impl_name}: Failed to create model instance: {e}")

        # Check for the absence of Dropout and BatchNorm in the model components
        for name, module in model.named_modules():
            assert not isinstance(module, torch.nn.Dropout), f"{impl_name}: Dropout found in model at {name}"
            assert not isinstance(module, torch.nn.BatchNorm1d), f"{impl_name}: BatchNorm1d found in model at {name}"
            assert not isinstance(module, torch.nn.BatchNorm2d), f"{impl_name}: BatchNorm2d found in model at {name}"
            assert not isinstance(module, torch.nn.BatchNorm3d), f"{impl_name}: BatchNorm3d found in model at {name}"
    except Exception as e:
        pytest.skip(f"{impl_name}: Error checking for dropout or batch norm: {e}")

def test_forward_method(implementation):
    """Test the forward method of the model."""
    impl_name, module = implementation

    # Get the CNN class
    convnet_class = get_convnet_class(module)
    if convnet_class is None:
        pytest.skip(f"{impl_name}: No CNN class found")

    # Calculate appropriate flattened_size for a 32x32 input image with two conv+maxpool layers
    # For a 32x32 input, after 2 layers of Conv2d with kernel_size=3 and MaxPool2d with kernel_size=2,
    # the feature map size would be approximately 6x6
    flattened_size = 64 * 6 * 6  # 64 channels, 6x6 feature map

    try:
        # Try with flattened_size parameter
        model = convnet_class(flattened_size=flattened_size)
    except TypeError:
        # If that fails, try without parameters
        try:
            model = convnet_class()
        except Exception as e:
            pytest.skip(f"{impl_name}: Failed to create model instance: {e}")

    # Create a dummy input tensor (batch_size, channels, height, width)
    batch_size = 2
    input_tensor = torch.randn(batch_size, 3, 32, 32)

    try:
        # Call forward method
        output = model(input_tensor)

        # Check output shape
        assert output.shape[0] == batch_size, f"{impl_name}: Output batch size should be {batch_size}, got {output.shape[0]}"
        assert output.shape[1] == 3, f"{impl_name}: Output features should be 3, got {output.shape[1]}"
    except Exception as e:
        pytest.skip(f"{impl_name}: Forward pass failed with error: {str(e)}")

def test_model_flattened_size_parameter(implementation):
    """Test that the model correctly uses the flattened_size parameter if applicable."""
    impl_name, module = implementation

    # Get the CNN class
    convnet_class = get_convnet_class(module)
    if convnet_class is None:
        pytest.skip(f"{impl_name}: No CNN class found")

    # Check if the model accepts flattened_size parameter
    try:
        model = convnet_class(flattened_size=1600)

        # If we reach here, the model accepts flattened_size
        # Check if any linear layer has this size as input
        found_matching_linear = False

        for module in model.modules():
            if isinstance(module, torch.nn.Linear):
                if module.in_features == 1600:
                    found_matching_linear = True
                    break

        assert found_matching_linear, f"{impl_name}: No Linear layer with in_features=1600 found, flattened_size parameter may not be used correctly"

        # Try another value to ensure the parameter is actually being used
        model2 = convnet_class(flattened_size=2048)

        found_matching_linear = False
        for module in model2.modules():
            if isinstance(module, torch.nn.Linear):
                if module.in_features == 2048:
                    found_matching_linear = True
                    break

        assert found_matching_linear, f"{impl_name}: The flattened_size parameter doesn't seem to affect the model structure"

    except TypeError:
        # Model doesn't accept flattened_size, which is okay for some implementations
        pytest.skip(f"{impl_name}: Model doesn't accept flattened_size parameter")
    except Exception as e:
        pytest.skip(f"{impl_name}: Unexpected error in flattened_size test: {str(e)}")

# def test_end_to_end_execution(implementation):
#     """Test the end-to-end execution of the model with a small batch of data."""
#     impl_name, module = implementation

#     # Get the CNN class
#     convnet_class = get_convnet_class(module)
#     if convnet_class is None:
#         pytest.skip(f"{impl_name}: No CNN class found")

#     # For a 16x16 input with 2 conv+maxpool layers, the feature map size would be around 2x2
#     flattened_size = 64 * 2 * 2  # 64 channels, 2x2 feature map

#     try:
#         # Try with flattened_size parameter
#         model = convnet_class(flattened_size=flattened_size)
#     except TypeError:
#         # If that fails, try without parameters
#         try:
#             model = convnet_

def test_end_to_end_execution(implementation):
    """Test the end-to-end execution of the model with a small batch of data."""
    impl_name, module = implementation

    # Get the CNN class
    convnet_class = get_convnet_class(module)
    if convnet_class is None:
        pytest.skip(f"{impl_name}: No CNN class found")

    # For a 16x16 input with 2 conv+maxpool layers, the feature map size would be around 2x2
    flattened_size = 64 * 2 * 2  # 64 channels, 2x2 feature map

    try:
        # Try with flattened_size parameter
        model = convnet_class(flattened_size=flattened_size)
    except TypeError:
        # If that fails, try without parameters
        try:
            model = convnet_class()
        except Exception as e:
            pytest.skip(f"{impl_name}: Failed to create model instance: {e}")

    # Create a dummy input tensor (batch_size, channels, height, width)
    batch_size = 3
    input_tensor = torch.randn(batch_size, 3, 16, 16)

    try:
        # Set the model to training mode
        model.train()

        # Define loss function and optimizer
        criterion = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

        # Create dummy target labels (batch_size,)
        target_labels = torch.randint(0, 3, (batch_size,))

        # Forward pass
        outputs = model(input_tensor)

        # Check output shape
        assert outputs.shape == (batch_size, 3), f"{impl_name}: Expected output shape {(batch_size, 3)}, got {outputs.shape}"

        # Calculate loss
        loss = criterion(outputs, target_labels)

        # Check that loss is a scalar tensor
        assert loss.dim() == 0, f"{impl_name}: Loss should be a scalar tensor, got dimension {loss.dim()}"
        assert not torch.isnan(loss).any(), f"{impl_name}: Loss contains NaN values"

        # Backward pass
        optimizer.zero_grad()
        loss.backward()

        # Check that gradients are computed for parameters
        any_grad = False
        for name, param in model.named_parameters():
            if param.grad is not None and torch.sum(torch.abs(param.grad)) > 0:
                any_grad = True
                break

        assert any_grad, f"{impl_name}: No gradients were computed during backward pass"

        # Optimizer step
        optimizer.step()

        # Try a second forward pass to ensure model still works after weight update
        new_outputs = model(input_tensor)
        assert new_outputs.shape == (batch_size, 3), f"{impl_name}: Model failed after optimizer step"

        # Test evaluation mode
        model.eval()
        with torch.no_grad():
            eval_outputs = model(input_tensor)

        assert eval_outputs.shape == (batch_size, 3), f"{impl_name}: Model failed in evaluation mode"

    except Exception as e:
        pytest.skip(f"{impl_name}: End-to-end execution failed with error: {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


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, ImageDraw, ImageFont
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 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):
    """
    Computes similarity maps and embeddings from a batch of images and queries using the specified model and processor.

    Args:
        batch_images (dict): A dictionary of batched image inputs processed by the processor.
        batch_queries (dict): A dictionary of batched query inputs processed by the processor.
        model (nn.Module): The model used for computing embeddings.
        processor (Processor): The processor responsible for image and text preprocessing.

    Returns:
        tuple: A tuple containing:
            - original_maps (torch.Tensor): Similarity maps between images and queries
                with shape (num_queries, n_patches_x, n_patches_y).
            - original_image_embeddings (torch.Tensor): Embeddings of the input images.
            - original_query_embeddings (torch.Tensor): 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

def visualize_token_map(image, original_maps, token_list, token_index=2, cmap="Greens"):
    """
    Visualize a token's attention map in three ways: the original image, the raw attention map with numerical values,
    and an overlay of the attention map on the original image.
    Args:
        image (PIL.Image): The input image to visualize.
        original_maps (torch.Tensor or np.ndarray): Attention maps with shape (num_tokens, height, width).
        token_list (list[str]): List of token strings corresponding to each attention map.
        token_index (int, optional): Index of the token/map to visualize. Defaults to 2.
        cmap (str, optional): Matplotlib colormap name for visualizing the attention maps. Defaults to "Greens".

    The function creates a figure with three subplots:
    1. The original input image
    2. The raw attention map with numerical values annotated
    3. The attention map overlaid on the original image with a colorbar

    Returns:
        None. Displays the visualization using matplotlib.
    """
    # 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 with matching values to the raw map
    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=visual_map.min(), vmax=visual_map.max())), 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()

def create_single_patch_image(
    n_patches_x, n_patches_y, patch_size, main_color, special_color, special_patch, special_patch_width=2,
):
    """
    Creates an image composed of colored patches, with one special patch highlighted.

    The image is divided into a grid of n_patches_x by n_patches_y patches, each of size
    patch_size x patch_size pixels. All patches are filled with the main_color, except
    for the special_patch, which is filled with special_color.  The special patch can
    also have a width of more than one patch.
    Args:
        n_patches_x (int): Number of patches horizontally.
        n_patches_y (int): Number of patches vertically.
        patch_size (int): The size (in pixels) of each square patch.
        main_color (list): The [R, G, B] color for most patches.
        special_color (list): The [R, G, B] color for the special patch.
        special_patch (tuple): The (row, col) position of the top-left corner of the special patch (0-indexed).
        special_patch_width (int, optional): The width of the special patch in number of patches. Defaults to 2.

    Returns:
        PIL Image: The generated image.
    """

    # Create a 3D NumPy array for the image
    img_height = n_patches_y * patch_size
    img_width = n_patches_x * patch_size
    image_data = np.zeros((img_height, img_width, 3), dtype=np.uint8)

    # Fill the entire image with the main color
    image_data[:, :] = main_color

    # Assign the special color to the special patch
    special_row, special_col = special_patch
    image_data[
        special_row * patch_size : (special_row + special_patch_width) * patch_size,
        special_col * patch_size : (special_col + special_patch_width) * patch_size
    ] = special_color

    return Image.fromarray(image_data)

def extract_patch_mask(image, patch_size, special_color=[0, 0, 0]):
    """
    Extract a binary mask indicating the location of the special patch.

    Args:
        image (PIL.Image.Image): The input image.
        patch_size (int): The size of each square patch in pixels.
        special_color (list[int]): The RGB color of the special patch.

    Returns:
        np.ndarray: A binary mask of shape (n_patches_y, n_patches_x) indicating
                    the special patch location (1 for special patch, 0 otherwise).
    """
    # Convert the image to a NumPy array
    image_np = np.array(image)

    # Get image dimensions
    img_height, img_width, _ = image_np.shape

    # Compute the number of patches
    n_patches_y = img_height // patch_size
    n_patches_x = img_width // patch_size

    # Initialize the patch mask
    patch_mask = np.zeros((n_patches_y, n_patches_x), dtype=np.int32)

    # Iterate over all patches to locate the special patch
    for row in range(n_patches_y):
        for col in range(n_patches_x):
            # Extract the patch
            patch = image_np[
                row * patch_size : (row + 1) * patch_size,
                col * patch_size : (col + 1) * patch_size
            ]

            # Check if the patch matches the special color
            if np.allclose(patch.mean(axis=(0, 1)), special_color, atol=1e-6):
                patch_mask[row, col] = 1  # Mark this patch as special

    return patch_mask

def evaluate_map_quality(similarity_map, patch_mask):
    """
    Evaluate the quality of a similarity map with respect to a binary patch mask.

    Args:
        similarity_map (np.ndarray): The similarity map (height, width).
        patch_mask (np.ndarray): The binary mask for the patch (1 for black patch, 0 elsewhere).

    Returns:
        dict: Metrics including correlation, peak accuracy, and overlap score.
    """
    # Flatten the map and mask for easier computation
    sim_map_flat = similarity_map.flatten()
    patch_mask_flat = patch_mask.flatten()

    # (A) Correlation
    correlation = np.corrcoef(sim_map_flat, patch_mask_flat)[0, 1]

    # (B) Peak Signal Location
    max_location = np.unravel_index(np.argmax(similarity_map), similarity_map.shape)
    expected_location = np.unravel_index(np.argmax(patch_mask), patch_mask.shape)
    peak_accuracy = 1 if max_location == expected_location else 0

    # (C) Normalized Map Overlap
    black_patch_score = similarity_map[patch_mask == 1].mean()
    background_score = similarity_map[patch_mask == 0].mean()
    overlap_score = black_patch_score / (background_score + 1e-8)  # Avoid division by zero

    # Return all metrics
    return {
        "correlation": correlation,
        "peak_accuracy": peak_accuracy,
        "overlap_score": overlap_score,
    }

def create_single_patch_image_with_text(
    n_patches_x,
    n_patches_y,
    patch_size,
    main_color,
    special_color,
    special_patch,
    text="Hello",
    text_color=(255, 255, 255),
    special_patch_width=2,
    font_size=16,
):
    """
    Creates an image composed of colored patches, but places a single word (or text)
    inside the "special" patch area.
    """
    # Create a 3D NumPy array for the image
    img_height = n_patches_y * patch_size
    img_width = n_patches_x * patch_size
    image_data = np.zeros((img_height, img_width, 3), dtype=np.uint8)

    # Fill the entire image with the main color
    image_data[:, :] = main_color

    # Assign the special color to the special patch area
    special_row, special_col = special_patch
    image_data[
        special_row * patch_size : (special_row + special_patch_width) * patch_size,
        special_col * patch_size : (special_col + special_patch_width) * patch_size
    ] = special_color

    # Convert to a Pillow Image so we can draw on it
    img = Image.fromarray(image_data)
    draw = ImageDraw.Draw(img)

    # Load font with specified size
    try:
        font = ImageFont.truetype("arial.ttf", font_size)
    except IOError:
        font = ImageFont.load_default()

    # Calculate the center of the special patch in pixel coordinates
    patch_center_x = (
        special_col * patch_size
        + (special_patch_width * patch_size) // 2
    )
    patch_center_y = (
        special_row * patch_size
        + (special_patch_width * patch_size) // 2
    )

    # Calculate text bounding box to center the text
    text_bbox = draw.textbbox((0, 0), text, font=font)
    text_width = text_bbox[2] - text_bbox[0]
    text_height = text_bbox[3] - text_bbox[1]

    text_x = patch_center_x - text_width // 2
    text_y = patch_center_y - text_height // 2

    # Place text in the center of the special patch
    draw.text((text_x, text_y), text, fill=text_color, font=font)

    return img

Instruction: make the font_size argument work

import pytest
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import inspect
from unittest.mock import patch, MagicMock
import sys
import re

from PIL import Image

class MockFont:
    def __init__(self, size):
        self.size = size

    def getbbox(self, text, *args, **kwargs):
        w = len(text) * self.size
        return (0, 0, w, self.size)

    def getsize(self, text, *args, **kwargs):
        w = len(text) * self.size
        return (w, self.size)

    def getmask(self, text, *args, **kwargs):
        # create a tiny “L” (8‑bit) image and hand back its .im
        img = Image.new("L", (len(text) * self.size, self.size), color=255)
        return img.im

@pytest.fixture
def mock_pil_font():
    """Mock PIL.ImageFont to avoid file system dependencies"""
    with patch('PIL.ImageFont.truetype', return_value=MockFont(16)):
        yield

@pytest.fixture(scope="module")
def mock_colpali_engine():
    """Mock the external colpali_engine module that's not available"""
    colpali_mock = MagicMock()
    interpretability_mock = MagicMock()

    # Set up the necessary mocked functions or classes
    interpretability_mock.get_similarity_maps_from_embeddings = MagicMock(return_value=[MagicMock()])
    interpretability_mock.plot_all_similarity_maps = MagicMock()

    # Assign the mock to the module
    colpali_mock.interpretability = interpretability_mock

    # Add the mock to sys.modules
    with patch.dict('sys.modules', {
        'colpali_engine': colpali_mock,
        'colpali_engine.interpretability': interpretability_mock
    }):
        yield colpali_mock

def test_font_size_parameter_exists(implementation, mock_colpali_engine):
    """Test that the function has a font_size parameter."""
    impl_name, module = implementation

    # Access the function directly by name
    try:
        func = module.create_single_patch_image_with_text
    except AttributeError:
        pytest.fail(f"{impl_name} doesn't have a 'create_single_patch_image_with_text' function")

    # Inspect the function signature
    sig = inspect.signature(func)
    params = sig.parameters

    # Check for a font_size parameter
    has_font_size = any(param.lower() == 'font_size' for param in params)

    assert has_font_size, f"Function should have a font_size parameter (found: {list(params.keys())})"

def test_function_creates_image_with_text(implementation, mock_colpali_engine, mock_pil_font):
    """Test that the function actually creates a PIL image with text."""
    impl_name, module = implementation

    # Access the function directly by name
    try:
        func = module.create_single_patch_image_with_text
    except AttributeError:
        pytest.fail(f"{impl_name} doesn't have a 'create_single_patch_image_with_text' function")

    # Basic arguments to create an image
    args = {
        'n_patches_x': 5,
        'n_patches_y': 5,
        'patch_size': 50,
        'main_color': [200, 200, 200],
        'special_color': [0, 0, 0],
        'special_patch': (1, 1),
        'text': "Test",
        'font_size': 16,
        'text_color': (255, 255, 255),
        'special_patch_width': 2
    }

    # Call function with the arguments
    with patch('PIL.ImageFont.truetype', return_value=MockFont(16)):
        result = func(**args)

    # Verify the result is a PIL Image
    assert isinstance(result, Image.Image), "Function should return a PIL Image"

    # Verify the image has reasonable dimensions based on the input
    expected_width = args['n_patches_x'] * args['patch_size']
    expected_height = args['n_patches_y'] * args['patch_size']
    assert result.width == expected_width, f"Image width should be {expected_width}"
    assert result.height == expected_height, f"Image height should be {expected_height}"

def test_font_size_affects_image_creation(implementation, mock_colpali_engine):
    """Test that different font sizes result in different image outputs."""
    impl_name, module = implementation

    # Access the function directly by name
    try:
        func = module.create_single_patch_image_with_text
    except AttributeError:
        pytest.fail(f"{impl_name} doesn't have a 'create_single_patch_image_with_text' function")

    # Basic arguments to create an image
    base_args = {
        'n_patches_x': 5,
        'n_patches_y': 5,
        'patch_size': 50,
        'main_color': [200, 200, 200],
        'special_color': [0, 0, 0],
        'special_patch': (1, 1),
        'text': "Test",
        'text_color': (255, 255, 255),
        'special_patch_width': 2
    }

    # Store the font sizes used
    font_sizes_used = []

    def mock_truetype(font_path, size, *args, **kwargs):
        font_sizes_used.append(size)
        return MockFont(size)

    # Mock the fonts and create two images with different font sizes
    with patch('PIL.ImageFont.truetype', side_effect=mock_truetype):
        # Add small font size
        small_args = base_args.copy()
        small_args['font_size'] = 16
        img_small = func(**small_args)

        # Add large font size
        large_args = base_args.copy()
        large_args['font_size'] = 32
        img_large = func(**large_args)

    # Verify that both font sizes were used
    assert 16 in font_sizes_used, "Font size 16 should have been used"
    assert 32 in font_sizes_used, "Font size 32 should have been used"

    # Both should be PIL Images
    assert isinstance(img_small, Image.Image), "Function should return a PIL Image with small font"
    assert isinstance(img_large, Image.Image), "Function should return a PIL Image with large font"

def test_different_font_sizes_produce_different_results(implementation, mock_colpali_engine):
    """Test that using different font sizes produces visibly different results."""
    impl_name, module = implementation

    # Access the function directly by name
    try:
        func = module.create_single_patch_image_with_text
    except AttributeError:
        pytest.fail(f"{impl_name} doesn't have a 'create_single_patch_image_with_text' function")

    # Basic arguments to create an image
    base_args = {
        'n_patches_x': 5,
        'n_patches_y': 5,
        'patch_size': 50,
        'main_color': [200, 200, 200],
        'special_color': [0, 0, 0],
        'special_patch': (1, 1),
        'text': "Test",
        'text_color': (255, 255, 255),
        'special_patch_width': 2
    }

    # Create a small and large font mock object for comparison
    small_font = MockFont(16)
    large_font = MockFont(32)

    # Mock the truetype function to return our controlled font sizes
    mock_truetype_calls = []
    def mock_truetype(font_path, size, *args, **kwargs):
        mock_truetype_calls.append(size)
        return small_font if size == 16 else large_font

    # Create two images with different font sizes
    with patch('PIL.ImageFont.truetype', side_effect=mock_truetype):
        # Small font size
        small_args = base_args.copy()
        small_args['font_size'] = 16
        img_small = func(**small_args)

        # Large font size
        large_args = base_args.copy()
        large_args['font_size'] = 32
        img_large = func(**large_args)

    # Verify that both calls to truetype were made with different sizes
    assert 16 in mock_truetype_calls, "truetype was not called with font size 16"
    assert 32 in mock_truetype_calls, "truetype was not called with font size 32"

    # Convert images to numpy arrays for comparison
    img_small_np = np.array(img_small)
    img_large_np = np.array(img_large)

    # The images should be different (at least some pixels should differ)
    # If the font size is affecting the image, pixel differences would be expected
    # We convert to binary to avoid issues with anti-aliasing or other rendering differences
    diff = np.sum(img_small_np != img_large_np)

    assert diff > 0, "Images with different font sizes should look different"

pytest
pytest-mock
Pillow
numpy
matplotlib
seaborn
scikit-learn
torch
torchvision
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 os
import json
import sys
import re
from   tkinter import *
from   tkinter import filedialog
from   tkinter import messagebox as mb
import eel
import tasks

@eel.expose
def loadJSONFile(initialdir):
    root = Tk()
    root.withdraw()
    root.wm_attributes('-topmost', 1)
    file_path = filedialog.askopenfilename(title="Выберите JSON файл для загрузки",filetypes=[("JSON files", "*.json"), ("All files", "*.*")],initialdir=initialdir)
    if file_path:
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                data = json.load(file)
            return data
        except Exception as e:
            print(f"Ошибка при загрузке JSON файла: {e}")
            return None
    else:
        print("Файл не выбран.")
        return None

@eel.expose
def saveJSONFile(json_data):
    root = Tk()
    root.withdraw()
    root.wm_attributes('-topmost', 1)
    file_path = filedialog.asksaveasfilename(title="Сохранить JSON файл как",defaultextension=".json",filetypes=[("JSON files", "*.json"), ("All files", "*.*")])
    if file_path:
        try:
            with open(file_path, 'w', encoding='utf-8') as file:
                json.dump(json_data, file, ensure_ascii=False, indent=4)
            print(f"JSON файл успешно сохранен: {file_path}")
        except Exception as e:
            print(f"Ошибка при сохранении JSON файла: {e}")
    else:
        print("Файл не выбран.")

@eel.expose
def select_file(title,patterns,initialdir):
    patterns=patterns or ['*.*']
    filetypes = [ [p,p.split('/')[-1]] for p in patterns]
    regex     = [ p.split('/')[0] for p in patterns if len(p.split('/'))==2 ]
    root = Tk()
    root.withdraw()
    root.wm_attributes('-topmost', 1)
    while  True:
        file     = filedialog.askopenfilename(filetypes=filetypes,title=title,initialdir=initialdir)
        if not file or not regex: break
        folder=file.replace('\\','/').split('/')[-2]
        for r in regex:
            if re.match(r, folder):
                return file
        mb.showerror("Ошибка",f"Папка не соответствует паттерну {','.join(regex)}. \nПовторите выбор файла")
    return file

@eel.expose
def save_task(taskData,taskName,settings):
    try:
        return tasks.save_task(taskData,taskName,settings),0
    except Exception as e:
        return  0,str(e)

# Напиши get_task_list для получения списка имен *.xml файлов из папки settings['taskPath'].
# В первой строке каждого файла опционально записана строка : <!-- comment -->
# Ответ представь в формате списка объектов {fileName, comment, file_date_str}
@eel.expose

    path = settings['taskPath']
    try:
        # Get list of all .xml files in directory
        xml_files = [f for f in os.listdir(path) if f.endswith('.xml')]
        return xml_files
    except Exception as e:
        print(f"Error getting task list: {e}")
        return []

if __name__ == "__main__":
    options={'mode':'chrome'}
    for i in range(2,len(sys.argv),2):
        if sys.argv[i-1]=='mode':  options['mode']=sys.argv[i]

    eel.init('web')
    eel.start('index.html', **options) #, host='localhost', port=8000,size=(1280,800),cmdline_args=['--start-fullscreen'])
    #eel.start('index.html', mode='chrome-app', host='localhost', port=8000,cmdline_args=['--start-fullscreen'])

Instruction: # Напиши get_task_list для получения списка имен *.xml файлов из папки settings['taskPath'].  # В первой строке каждого файла опционально записана строка : <!-- comment --> # Ответ представь в формате списка объектов {fileName, comment, file_date_str}

import os
import re
import tempfile
import pytest
from datetime import datetime
from unittest.mock import patch, mock_open, MagicMock

@pytest.fixture
def mock_file_system():
    """Create a mock file system for testing."""
    with tempfile.TemporaryDirectory() as temp_dir:
        # Create sample XML files with different comment formats
        files = {
            'task1.xml': '<!-- This is a comment -->\n<task>Content</task>',
            'task2.xml': '<task>No comment</task>',
            'task3.xml': '<!--Multiple words comment-->\n<task>Content</task>',
            'nonxml.txt': '<!-- Not an XML file -->',
        }

        # Create the files in the temporary directory
        for filename, content in files.items():
            with open(os.path.join(temp_dir, filename), 'w', encoding='utf-8') as f:
                f.write(content)

        yield temp_dir

def extract_get_task_list(module):
    """Extract get_task_list function directly from module source code."""
    if hasattr(module, 'get_task_list'):
        return getattr(module, 'get_task_list')

    if hasattr(module, '__file__'):
        with open(module.__file__, 'r', encoding='utf-8') as f:
            source = f.read()

            # Check if the module has exposed get_task_list function via eel
            if '@eel.expose' in source and 'def get_task_list' in source:
                # Instead of trying to execute code we extract dynamically,
                # add a monkey-patched version to the module
                def get_task_list_wrapper(settings):
                    if not os.path.exists(settings.get('taskPath', '')):
                        return []

                    path = settings['taskPath']
                    result = []

                    try:
                        # Get list of all .xml files in directory
                        xml_files = [f for f in os.listdir(path) if f.endswith('.xml')]

                        for file_name in xml_files:
                            file_path = os.path.join(path, file_name)

                            # Get file modification date
                            mod_time = os.path.getmtime(file_path)

                            # Extract comment from first line if exists
                            comment = ""
                            try:
                                with open(file_path, 'r', encoding='utf-8') as file:
                                    first_line = file.readline().strip()
                                    comment_match = re.search(r'<!--\s*(.*?)\s*-->', first_line)
                                    if comment_match:
                                        comment = comment_match.group(1)
                            except Exception as e:
                                print(f"Error reading file {file_name}: {e}")

                            # Add file info to result
                            result.append({
                                "fileName": file_name,
                                "comment": comment,
                                "file_date_str": mod_time
                            })

                        return result
                    except Exception as e:
                        print(f"Error getting task list: {e}")
                        return []

                # Attach the function to the module
                setattr(module, 'get_task_list', get_task_list_wrapper)
                return get_task_list_wrapper

    return None

def test_get_task_list_function_exists(implementation):
    """Test that the get_task_list function exists in the implementation."""
    impl_name, module = implementation

    # Use the helper function to find the get_task_list function
    func = extract_get_task_list(module)

    # Check if we found the function
    assert func is not None, f"{impl_name}: get_task_list function is missing"

    # Make it available for other tests
    module.get_task_list = func

def test_get_task_list_is_exposed(implementation):
    """Test that the get_task_list function is exposed to the frontend."""
    impl_name, module = implementation

    # Test if we can find @eel.expose in the source code for get_task_list
    if hasattr(module, '__file__'):
        with open(module.__file__, 'r', encoding='utf-8') as f:
            source = f.read()
            assert '@eel.expose' in source and 'def get_task_list' in source, \
                f"{impl_name}: get_task_list function is not exposed with @eel.expose"

@patch('os.path.exists', return_value=True)
def test_get_task_list_returns_xml_files_only(mock_exists, implementation, mock_file_system):
    """Test that get_task_list only returns XML files."""
    impl_name, module = implementation

    # Make sure we have the function available
    if not hasattr(module, 'get_task_list'):
        test_get_task_list_function_exists(implementation)

    # Call the function with settings pointing to our mock file system
    settings = {'taskPath': mock_file_system}
    result = module.get_task_list(settings)

    # Check that the result is a list
    assert isinstance(result, list), f"{impl_name}: get_task_list should return a list"

    # Check the total count matches expected
    assert len(result) == 3, f"{impl_name}: Expected 3 XML files but got {len(result)}"

    # Check that only XML files are included
    filenames = [item.get('fileName', '') for item in result]
    assert 'task1.xml' in filenames, f"{impl_name}: task1.xml should be in the result"
    assert 'task2.xml' in filenames, f"{impl_name}: task2.xml should be in the result"
    assert 'task3.xml' in filenames, f"{impl_name}: task3.xml should be in the result"
    assert 'nonxml.txt' not in filenames, f"{impl_name}: nonxml.txt should not be in the result"

@patch('os.path.exists', return_value=True)
def test_get_task_list_extracts_comments(mock_exists, implementation, mock_file_system):
    """Test that get_task_list correctly extracts comments from the first line."""
    impl_name, module = implementation

    # Make sure we have the function available
    if not hasattr(module, 'get_task_list'):
        test_get_task_list_function_exists(implementation)

    # Call the function
    settings = {'taskPath': mock_file_system}
    result = module.get_task_list(settings)

    # Create a mapping of filename to result item for easier assertion
    result_map = {item.get('fileName', ''): item for item in result}

    # Check comments are correctly extracted
    assert 'This is a comment' in result_map.get('task1.xml', {}).get('comment', ''), \
        f"{impl_name}: Comment not correctly extracted for task1.xml"
    assert result_map.get('task2.xml', {}).get('comment', '') == '', \
        f"{impl_name}: File without comment should have empty comment field"
    assert 'Multiple words comment' in result_map.get('task3.xml', {}).get('comment', ''), \
        f"{impl_name}: Comment not correctly extracted for task3.xml"

@patch('os.path.exists', return_value=True)
def test_get_task_list_includes_date(mock_exists, implementation, mock_file_system):
    """Test that get_task_list includes a date string for each file."""
    impl_name, module = implementation

    # Make sure we have the function available
    if not hasattr(module, 'get_task_list'):
        test_get_task_list_function_exists(implementation)

    # Call the function
    settings = {'taskPath': mock_file_system}
    result = module.get_task_list(settings)

    # Check that each result has a file_date_str field
    for item in result:
        assert 'file_date_str' in item, f"{impl_name}: file_date_str missing from result item"

        # Accept either timestamp or formatted date string
        if isinstance(item['file_date_str'], (int, float)):
            # Valid timestamp
            assert item['file_date_str'] > 0, f"{impl_name}: file_date_str should be a positive number"
        else:
            # Should be a date string
            assert isinstance(item['file_date_str'], str), f"{impl_name}: file_date_str should be a string if not a timestamp"
            # Check if it has numbers and separators
            assert re.search(r'\d', item['file_date_str']), f"{impl_name}: file_date_str should contain numeric values"
            assert any(sep in item['file_date_str'] for sep in ['-', '/', '.', ' ', ':']), \
                f"{impl_name}: file_date_str should contain date/time separators"

@patch('os.path.exists', return_value=True)
def test_get_task_list_format(mock_exists, implementation, mock_file_system):
    """Test that get_task_list returns the correct object format."""
    impl_name, module = implementation

    # Make sure we have the function available
    if not hasattr(module, 'get_task_list'):
        test_get_task_list_function_exists(implementation)

    # Call the function
    settings = {'taskPath': mock_file_system}
    result = module.get_task_list(settings)

    # Check that result is not empty
    assert len(result) > 0, f"{impl_name}: get_task_list should return a non-empty list"

    # Check that each item has the required fields
    for item in result:
        assert 'fileName' in item, f"{impl_name}: Result items must have 'fileName' field"
        assert 'comment' in item, f"{impl_name}: Result items must have 'comment' field"
        assert 'file_date_str' in item, f"{impl_name}: Result items must have 'file_date_str' field"

        # Check types
        assert isinstance(item['fileName'], str), f"{impl_name}: 'fileName' must be a string"
        assert isinstance(item['comment'], str), f"{impl_name}: 'comment' must be a string"

@patch('os.path.exists', return_value=True)
def test_empty_directory_returns_empty_list(mock_exists, implementation, tmp_path):
    """When there are no XML files, get_task_list should return an empty list."""
    # point to an empty tmp_path
    settings = {'taskPath': str(tmp_path)}
    func = implementation[1].get_task_list
    result = func(settings)
    assert isinstance(result, list)
    assert result == []

@patch('os.path.exists', return_value=True)
def test_ignores_comments_not_on_first_line(mock_exists, implementation, tmp_path):
    """Only first‐line comments should be picked up, not ones further down."""
    p = tmp_path / "foo.xml"
    p.write_text("<task>no comment here</task>\n<!-- hidden -->\n<more/>")
    func = implementation[1].get_task_list
    out = func({'taskPath': str(tmp_path)})[0]
    assert out['comment'] == ""

@patch('os.path.exists', return_value=True)
def test_handles_empty_file_gracefully(mock_exists, implementation, tmp_path):
    """Zero‐byte XML files or files with no lines must not break."""
    p = tmp_path / "empty.xml"
    p.write_bytes(b"")
    func = implementation[1].get_task_list
    out = func({'taskPath': str(tmp_path)})[0]
    assert out['comment'] == ""

@patch('os.path.exists', return_value=True)
def test_strips_whitespace_in_comment(mock_exists, implementation, mock_file_system):
    """Comments with extra spaces inside the delimiters should be trimmed."""
    # create a file whose first line is “<!--   hello world   -->”
    comment_file = os.path.join(mock_file_system, "spaced.xml")
    with open(comment_file, 'w') as f:
        f.write("<!--   hello   world   -->\n<task/>")
    func = implementation[1].get_task_list
    out = next(item for item in func({'taskPath': mock_file_system})
               if item['fileName']=="spaced.xml")
    assert out['comment'] == "hello   world"

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