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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return image_embeds, text_embeds

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

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

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

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

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

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

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

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

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

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

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

    return mean_precision, mean_recall

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

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

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

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

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

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

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

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

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

    plt.tight_layout()
    plt.show()

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

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

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

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

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

Instruction: do not interpolate, just upscale linearly

import pytest
import inspect
import numpy as np
import torch
from unittest.mock import MagicMock, Mock, patch
from PIL import Image
import matplotlib.pyplot as plt
import re

@pytest.fixture
def mock_image():
    """Create a mock PIL image."""
    return Image.new('RGB', (100, 80))

@pytest.fixture
def mock_similarity_maps():
    """Create mock similarity maps tensor."""
    return torch.rand(3, 10, 8)  # 3 tokens, 10x8 map size

@pytest.fixture
def mock_query_tokens():
    """Create mock query tokens."""
    return ["token1", "token2", "token3"]

def find_function_by_signature(module, signature_pattern):
    """Find a function in a module that matches the expected signature."""
    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if obj.__module__ == module.__name__:
            try:
                source = inspect.getsource(obj)
                if re.search(signature_pattern, source):
                    return obj
            except (TypeError, OSError):
                continue
    return None

def get_plot_function(module):
    """Get the plotting function from the module that matches our expected signature."""
    # Look for a function that takes image, similarity maps, and query tokens
    # More flexible pattern to find various function signatures
    patterns = [
        r"def\s+(\w+)\s*\(\s*(?:.*?pil_image|.*?image).*?similarity_maps.*?query_tokens",
        r"def\s+plot_query_token_importance\s*\(",  # Explicit function name
        r"def\s+(\w+)\s*\(.*?image.*?maps.*?tokens"  # More generic pattern
    ]

    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if obj.__module__ == module.__name__:
            try:
                source = inspect.getsource(obj)
                for pattern in patterns:
                    if re.search(pattern, source, re.DOTALL):
                        return obj
            except (TypeError, OSError):
                continue

    # If not found by pattern, try to find by inspecting function arguments
    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if obj.__module__ == module.__name__:
            try:
                sig = inspect.signature(obj)
                param_names = set(sig.parameters.keys())
                # Check for common parameter names in the plot functions
                if len(param_names.intersection({"image", "pil_image", "img", "similarity_maps", "maps", "query_tokens", "tokens"})) >= 3:
                    return obj
            except (TypeError, ValueError):
                continue

    # Fall back to checking if the function name contains expected terms
    for name, obj in inspect.getmembers(module, inspect.isfunction):
        if obj.__module__ == module.__name__:
            if "plot" in name.lower() and any(term in name.lower() for term in ["token", "importance", "heatmap", "similarity"]):
                return obj

    return None

def add_plot_function(module):
    """
    Add a minimal plot_query_token_importance function to a module if it doesn't exist.
    This allows tests to run against implementations without the function.
    """
    if get_plot_function(module) is None:
        # Define a minimal implementation
        def plot_query_token_importance(pil_image, similarity_maps, query_tokens, alpha=0.5):
            """
            Plot a separate heatmap for each query token in the similarity_maps.

            Args:
                pil_image (PIL.Image.Image): The original image
                similarity_maps (torch.Tensor): Maps of shape (num_tokens, height, width)
                query_tokens (list): A list of query token strings
                alpha (float): Transparency value for heatmap overlay
            """
            # Convert PIL to numpy
            image_np = np.array(pil_image)
            H, W = image_np.shape[:2]

            num_tokens = similarity_maps.size(0)

            fig, axs = plt.subplots(1, num_tokens, figsize=(5 * num_tokens, 5))
            if num_tokens == 1:
                axs = [axs]

            for idx in range(num_tokens):
                # Get similarity map for current token
                single_map = similarity_maps[idx]

                # Upscale without interpolation
                H_map, W_map = single_map.shape
                scale_x, scale_y = W // W_map, H // H_map
                heatmap = np.kron(single_map.cpu().numpy(), np.ones((scale_y, scale_x)))

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

            plt.tight_layout()

        # Add the function to the module
        setattr(module, "plot_query_token_importance", plot_query_token_importance)
        return plot_query_token_importance

    return get_plot_function(module)

def test_interpolation_not_used(implementation):
    """Test that bilinear interpolation is not used for upscaling the similarity maps."""
    impl_name, module = implementation

    plot_function = get_plot_function(module)
    if not plot_function:
        plot_function = add_plot_function(module)

    # Get the source code of the plotting function
    source_code = inspect.getsource(plot_function)

    # Check if F.interpolate with mode='bilinear' is present in the code
    assert "mode='bilinear'" not in source_code and "mode=\"bilinear\"" not in source_code, \
        f"Implementation '{impl_name}' uses bilinear interpolation"

    # Check for other kinds of interpolation that aren't nearest
    interpolation_methods = [
        "bicubic", "area", "lanczos", "bilinear",
        "INTER_CUBIC", "INTER_AREA", "INTER_LINEAR", "INTER_LANCZOS4"
    ]

    for method in interpolation_methods:
        assert method not in source_code, \
            f"Implementation '{impl_name}' might use interpolation '{method}' instead of nearest neighbor"

    # If using interpolate, ensure it's with nearest mode
    if "F.interpolate" in source_code:
        assert "mode='nearest'" in source_code or "mode=\"nearest\"" in source_code, \
            f"Implementation '{impl_name}' uses F.interpolate without nearest neighbor mode"

@pytest.mark.parametrize("image_size,map_size", [
    ((200, 160), (10, 8)),
    ((400, 320), (20, 16)),
])
def test_upscaling_method(implementation, image_size, map_size):
    """Test that the upscaling method preserves pixel values without interpolation."""
    impl_name, module = implementation

    plot_function = get_plot_function(module)
    if not plot_function:
        plot_function = add_plot_function(module)

    # Create a constant-value map to test upscaling
    similarity_map = torch.ones((1, map_size[0], map_size[1]))
    query_tokens = ["test_token"]

    # Mock image
    mock_image = Image.new('RGB', image_size)

    # Mock plt methods to capture what's passed to imshow
    with patch('matplotlib.pyplot.figure'), \
     patch('matplotlib.pyplot.tight_layout'), \
     patch('matplotlib.pyplot.show'), \
     patch('matplotlib.pyplot.subplots') as mock_subplots:

        # --- Create proper mocks ---
        num_tokens = len(query_tokens)  # Must match your test

        axs_list = []
        if num_tokens == 1:
            # If only one token, plt.subplots returns a single mock (NOT a list)
            single_ax = MagicMock()
            single_ax.imshow = MagicMock()
            single_ax.axis = MagicMock()
            axs_list.append(single_ax)
            mock_subplots.return_value = (MagicMock(), single_ax)
        else:
            # Multiple tokens: plt.subplots returns list of axes
            for _ in range(num_tokens):
                ax = MagicMock()
                ax.imshow = MagicMock()
                ax.axis = MagicMock()
                axs_list.append(ax)
            mock_subplots.return_value = (MagicMock(), axs_list)

        # Call plotting function
        plot_function(mock_image, similarity_map, query_tokens)

        # Now you can safely assert
        for ax in axs_list:
            assert ax.imshow.call_count > 0, f"imshow not called on one of the Axes in {impl_name}"

        all_imshow_arrays = []
        for ax in axs_list:
            for call in ax.imshow.call_args_list:
                array_passed = call[0][0]  # first positional arg to imshow
                all_imshow_arrays.append(array_passed)

        for heatmap_array in all_imshow_arrays:
            if isinstance(heatmap_array, np.ndarray) and heatmap_array.ndim in [2, 3]:
                H, W = heatmap_array.shape[:2]
                expected_H, expected_W = image_size[1], image_size[0]
                # Allow a small tolerance
                assert abs(H - expected_H) <= 5 and abs(W - expected_W) <= 5, (
                    f"Heatmap shape {H}x{W} is wrong, expected close to {expected_H}x{expected_W}"
                )
def test_linear_upscaling_no_interpolation(implementation):
    """Test that upscaling is performed by simple block replication (no interpolation)."""
    impl_name, module = implementation

    # Find plotting function
    plot_function = get_plot_function(module)
    if not plot_function:
        plot_function = add_plot_function(module)

    # Create a pattern (small 2x2 map) to clearly check
    pattern = torch.tensor([
        [1.0, 0.0],
        [0.0, 1.0]
    ])
    similarity_map = torch.zeros((1, 2, 2))
    similarity_map[0] = pattern
    query_tokens = ["test_token"]

    # Create a mock image that upscales 2x2 -> 8x8
    mock_image = Image.new('RGB', (8, 8))

    with patch('matplotlib.pyplot.figure'), \
         patch('matplotlib.pyplot.tight_layout'), \
         patch('matplotlib.pyplot.show'), \
         patch('matplotlib.pyplot.subplots') as mock_subplots:

        # --- Setup mock axes correctly ---
        num_tokens = len(query_tokens)
        if num_tokens == 1:
            ax = MagicMock()
            ax.imshow = MagicMock()
            ax.axis = MagicMock()
            axs = ax  # single Ax
        else:
            axs = []
            for _ in range(num_tokens):
                ax = MagicMock()
                ax.imshow = MagicMock()
                ax.axis = MagicMock()
                axs.append(ax)

        mock_subplots.return_value = (MagicMock(), axs)

        # --- Call the plotting function ---
        plot_function(mock_image, similarity_map, query_tokens)

        # --- Extract the imshow heatmap call ---
        axes_to_check = [axs] if not isinstance(axs, list) else axs

        for ax in axes_to_check:
            assert ax.imshow.call_count >= 2, f"Expected 2 imshow calls (background + heatmap) for '{impl_name}'"
            assert ax.axis.called, f"Expected axis('off') to be called for '{impl_name}'"

        # Focus on the second imshow call (the heatmap)
        heatmap = None
        for ax in axes_to_check:
            if len(ax.imshow.call_args_list) >= 2:
                heatmap = ax.imshow.call_args_list[1][0][0]  # Second call, first arg
                break

        assert heatmap is not None, f"'{impl_name}' does not properly pass heatmap to imshow."

        # --- Analyze the heatmap ---
        if isinstance(heatmap, list):
            heatmap = np.array(heatmap)  # Some mocking oddities return list instead of ndarray
        if heatmap.ndim > 2:
            heatmap = heatmap[:, :, 0]  # Take first channel if 3D

        H, W = heatmap.shape
        assert H >= 8 and W >= 8, f"'{impl_name}' heatmap too small after upscaling: {H}x{W}"

        # Check values — should replicate blocks (not smooth interpolate)
        unique_values = set()
        for i in range(H):
            for j in range(W):
                val = round(heatmap[i, j] * 10) / 10  # Round for floating point noise
                unique_values.add(val)

        assert len(unique_values) <= 3, \
            f"'{impl_name}' shows too many unique values — suggesting interpolation used instead of block replication: {unique_values}"

def test_non_integer_scale_handling(implementation):
    """Test that non-integer scaling factors are handled gracefully during upscaling."""
    impl_name, module = implementation

    # Find the plotting function
    plot_function = get_plot_function(module)
    if not plot_function:
        plot_function = add_plot_function(module)

    # Create a simple 3x3 map
    similarity_map = torch.ones((1, 3, 3))  # 3x3 grid
    query_tokens = ["test_token"]

    # Create a mock image size that doesn't divide evenly (10x10)
    mock_image = Image.new('RGB', (10, 10))

    with patch('matplotlib.pyplot.figure'), \
         patch('matplotlib.pyplot.tight_layout'), \
         patch('matplotlib.pyplot.show'), \
         patch('matplotlib.pyplot.subplots') as mock_subplots:

        # --- Setup mock axes ---
        num_tokens = len(query_tokens)
        if num_tokens == 1:
            ax = MagicMock()
            ax.imshow = MagicMock()
            ax.axis = MagicMock()
            axs = ax  # single Ax
        else:
            axs = []
            for _ in range(num_tokens):
                ax = MagicMock()
                ax.imshow = MagicMock()
                ax.axis = MagicMock()
                axs.append(ax)

        mock_subplots.return_value = (MagicMock(), axs)

        try:
            # --- Call the plotting function ---
            plot_function(mock_image, similarity_map, query_tokens)

            # --- Extract heatmap passed to imshow ---
            axes_to_check = [axs] if not isinstance(axs, list) else axs

            heatmap = None
            for ax in axes_to_check:
                if len(ax.imshow.call_args_list) >= 2:
                    heatmap = ax.imshow.call_args_list[1][0][0]
                    break

            assert heatmap is not None, f"'{impl_name}' did not generate a heatmap."

            # --- Analyze heatmap shape ---
            if isinstance(heatmap, list):
                heatmap = np.array(heatmap)
            if heatmap.ndim > 2:
                heatmap = heatmap[:, :, 0]  # Take first channel if 3D

            H, W = heatmap.shape

            # For a 10x10 image and a 3x3 map, scaling isn't exact — allow 1–2 pixel tolerance
            assert abs(H - 10) <= 2 and abs(W - 10) <= 2, \
                f"'{impl_name}' produced heatmap of wrong size for non-integer scaling: got ({W}, {H}) expected (~10x10)"

            # --- Ensure axis('off') was called ---
            for ax in axes_to_check:
                assert ax.axis.called, f"'{impl_name}' should call axis('off')"

        except Exception as e:
            pytest.fail(f"'{impl_name}' failed to handle non-integer scale factors gracefully: {str(e)}")

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output


import asyncio
from aiogram import Bot, Dispatcher, executor, types
import time

token = "TOKEN"
bot = Bot(token=token)
dp = Dispatcher(bot)
id = 111111111
time_time = 5

b = bool(True)

@dp.message_handler(commands=['start'])
async def start(message: types.Message):
    markup = types.InlineKeyboardMarkup(row_width=2)
    markup = types.InlineKeyboardMarkup(inline_keuboard = [
        [types.InlineKeyboardButton(text="Да", callback_data="stop"),
        types.InlineKeyboardButton(text="Нет", callback_data="continue")]
        ])
    await bot.send_message(chat_id=id, text="Ты робот?", reply_markup=markup)

@dp.callback_query_handler(text="stop")
async def stop(call: types.CallbackQuery):
    global b
    # b = False
    # await bot.send_message(chat_id=call.message.chat.id, text="Молодец, лох")
    await bot.send_message(chat_id=call.message.chat.id, text="<obscene message>!")

@dp.callback_query_handler(text="continue")
async def stop(call: types.CallbackQuery):
    await bot.send_message(chat_id=call.message.chat.id, text="Лох^лох, лох")

@dp.message_handler(content_types=['text'])
async def handle_all_messages(message: types.Message):
    with open(r"D:\Python files\!MoexApiBot\censored.gif", "rb") as vid:
        await bot.send_video(chat_id=id, video=vid, caption="Го <censor> ёр селф")

async def send_periodic_messages():
    while b:
        await bot.send_message(chat_id=id, text="Лох лох")
        with open(r"D:\Python files\!MoexApiBot\Shocked13.mp4", "rb") as vid:
            await bot.send_video(chat_id=id, video=vid, caption="Ты проиграл")
        await asyncio.sleep(time_time)

async def on_startup(dp):
    print('Бот запущен!')
    asyncio.create_task(send_periodic_messages())

if __name__ == '__main__':
    executor.start_polling(dp, skip_updates=True, on_startup=on_startup)

#for i in range(100):
    #vid = open(r"D:\Python files\!MoexApiBot\d530f99f38bf98bbeb213cb32b53012b.mp4", "rb")
    #bot.send_message(chat_id=444444444, text = f"Ты проиграл {i} раз")
    #bot.send_message(chat_id=444444444, text = f"Вас взломали {i} раз")
    #bot.send_message(chat_id=444444444, text = f"Отправьте {100*i} рублей на счет +55555555555)
    #bot.send_video(chat_id=444444444, video=vid, caption="Ты проиграл")
    #vid.close()
    #time.sleep(1)
#555555555 NAME_4
#444444444 Добр
#333333333 NAME_3
#222222222 NAME_2
#111111111 NAME_1

Instruction: перепиши без executor, так как в новой версии aiogram он потерял актуальность

import pytest
import inspect
import asyncio
import re
from unittest.mock import patch, MagicMock, AsyncMock, mock_open

@pytest.fixture
def mock_bot():
    """Create a mock bot instance."""
    bot = AsyncMock()
    bot.send_message = AsyncMock()
    bot.send_video = AsyncMock()
    return bot

@pytest.fixture
def mock_dp():
    """Create a mock dispatcher instance."""
    dp = MagicMock()
    dp.start_polling = AsyncMock()
    dp.register_on_startup = MagicMock()
    return dp

def test_no_executor_import(implementation):
    """Test that executor is not imported in the new implementation."""
    impl_name, module = implementation

    # Check the source code for imports
    source = inspect.getsource(module)

    # The implementation should not use executor from aiogram
    assert not re.search(r'from\s+aiogram\s+import\s+.*\bexecutor\b', source), \
        f"{impl_name} should not import executor from aiogram"

    assert not re.search(r'\bexecutor\s*\.', source), f"{impl_name} should not use executor.*"

pytest
pytest-mock
pytest-asyncio
aiogram<3.0.0
aiohttp<3.8.0
multidict
async-timeout
yarl
frozenlist
aiosignal
attrs

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 time

options = uc.ChromeOptions()
options.binary_location = (
    r"C:\Programming\Test\IP_Test\chromedriver-win64\chromedriver.exe"
)
print("wde")
with uc.Chrome(use_subprocess=True, options=options) as driver:
    print("wde")
    driver.get("https://lmarena.ai/")
    print("wde")
    # create an instance of ChromeOptions for undetected_chromedriver
    # initialize the undetected Chrome driver with specified options
    time.sleep(10)

Instruction: сделай путь относительным

import pytest
import sys
import os
import inspect
import re
from pathlib import Path
from unittest.mock import patch, MagicMock
import importlib

def test_imports_present(implementation):
    """Test that necessary imports are present in the implementation."""
    impl_name, module = implementation

    # Get module source code
    module_source = inspect.getsource(module)

    # Check if time module is imported or used
    time_imported = (
        "time" in dir(module)
        or "import time" in module_source
        or "time.sleep" in module_source
    )
    assert time_imported, f"{impl_name} should import or use time module"

    # Check for Path import in implementations using pathlib
    if "pathlib" in module_source:
        assert (
            "Path" in dir(module) or "from pathlib import Path" in module_source
        ), f"{impl_name} should import Path from pathlib"

    # Comprehensive patterns for undetected_chromedriver imports
    uc_import_patterns = [
        "uc" in dir(module),
        "import undetected_chromedriver as uc" in module_source,
        "from undetected_chromedriver import Chrome" in module_source,
        "import undetected_chromedriver" in module_source,
        "undetected_chromedriver.Chrome" in module_source,
    ]

    # For implementations that may not directly import but reference undetected_chromedriver
    if not any(uc_import_patterns):
        if "Chrome" in module_source and "selenium" not in module_source:
            assert (
                True
            ), "Using Chrome without explicit import (may be in test fixtures)"
        else:
            assert any(
                uc_import_patterns
            ), f"{impl_name} should import undetected_chromedriver"

def test_relative_path_used(implementation):
    """Test that the implementation uses a relative path for binary_location."""
    impl_name, module = implementation

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

    # Check for relative path patterns
    relative_path_patterns = [
        "./chromedriver-win64",
        "chromedriver-win64/",
        "Path(",
        "resolve()",
        "os.path.join",
        "os.path.dirname(__file__)",
        "__file__",
        "./",
    ]

    # Check that relative paths are used
    has_relative_path = any(
        pattern in module_source for pattern in relative_path_patterns
    )
    assert (
        has_relative_path
    ), f"{impl_name} should use a relative path for binary_location"

    # Extract non-comment lines to check for hardcoded paths in actual code
    code_lines = []
    in_multiline_comment = False
    for line in module_source.split("\n"):
        line = line.strip()

        # Handle multiline strings/comments
        if line.startswith('"""') or line.startswith("'''"):
            in_multiline_comment = not in_multiline_comment
            continue

        # Skip comments and empty lines
        if not in_multiline_comment and not line.startswith("#") and line:
            code_lines.append(line)

    active_code = "\n".join(code_lines)

    # Check for hardcoded absolute Windows paths in active code
    abs_path_in_active_code = re.search(
        r'options\.binary_location\s*=\s*[\'"]C:\\', active_code
    )
    assert (
        not abs_path_in_active_code
    ), f"{impl_name} contains hardcoded absolute Windows path"

def test_path_resolution_works(implementation):
    """Test that the path resolution approach is valid."""
    impl_name, module = implementation

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

    # Comprehensive list of path resolution methods
    path_resolution_methods = [
        "os.path.join",
        "os.path.dirname",
        "os.path.abspath",
        "__file__",
        "Path(",
        ".resolve()",
        "./chromedriver-win64",
        "chromedriver-win64/",
        "binary_location",
        "options.binary_location",
    ]

    has_path_resolution = any(
        method in module_source for method in path_resolution_methods
    )
    assert has_path_resolution, f"{impl_name} should use proper path resolution"

    # Ensure reference to chromedriver binary
    chrome_binary_patterns = ["chromedriver", "chromedriver.exe", "chromedriver-win64"]
    has_chromedriver_reference = any(
        pattern in module_source for pattern in chrome_binary_patterns
    )
    assert (
        has_chromedriver_reference
    ), f"{impl_name} should reference the chromedriver binary"

def test_chrome_instance_creation(implementation):
    """Test that Chrome instance is created with correct parameters."""
    impl_name, module = implementation

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

    # Check for Chrome instance creation
    assert "Chrome(" in module_source, f"{impl_name} should create a Chrome instance"

    # Check for use_subprocess parameter
    subprocess_patterns = [r"use_subprocess\s*=\s*True", r"use_subprocess=True"]
    has_subprocess_param = any(
        re.search(pattern, module_source) for pattern in subprocess_patterns
    )
    assert (
        has_subprocess_param
    ), f"{impl_name} should create Chrome instance with use_subprocess=True"

    # Flexible driver initialization patterns
    driver_init_patterns = [
        r"driver\s*=\s*[\w\.]+Chrome",
        r"with\s+[\w\.]+Chrome",
        r"[\w\.]+Chrome\(.*\)\s+as\s+driver",
    ]
    has_driver_init = any(
        re.search(pattern, module_source) for pattern in driver_init_patterns
    )
    assert has_driver_init, f"{impl_name} should properly initialize a Chrome driver"

def test_other_functionality_preserved(implementation):
    """Test that other functionality from the original code is preserved."""
    impl_name, module = implementation

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

    # Check for key functionality preservation
    assert (
        'print("wde")' in module_source
    ), f"{impl_name} should include print statements"
    assert (
        "driver.get(" in module_source
    ), f"{impl_name} should include driver.get() calls"
    assert "lmarena.ai" in module_source, f"{impl_name} should preserve the URL"

    # Check for sleep pattern with flexible matching
    sleep_patterns = [r"time\.sleep\s*\(\s*10\s*\)", r"sleep\s*\(\s*10\s*\)"]
    has_sleep = any(re.search(pattern, module_source) for pattern in sleep_patterns)
    assert has_sleep, f"{impl_name} should preserve the sleep(10) call"

    # Verify minimum print statements
    print_count = module_source.count('print("wde")')
    assert print_count >= 3, f"{impl_name} should maintain at least 3 print statements"

def test_chromedriver_path_correctness(implementation):
    """Test that the chromedriver path references the correct subdirectory structure."""
    impl_name, module = implementation

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

    # Comprehensive patterns for chromedriver path
    path_patterns = [
        "chromedriver-win64/chromedriver.exe",
        "chromedriver-win64\\\\chromedriver.exe",  # Escaped backslashes
        "chromedriver-win64/chromedriver",
        "chromedriver-win64\\\\chromedriver",
        "chromedriver-win64",
        "chromedriver.exe",
        "binary_location",
    ]

    has_valid_path = any(pattern in module_source for pattern in path_patterns)
    assert (
        has_valid_path
    ), f"{impl_name} should reference the correct chromedriver path structure"

    # Check for executable reference
    driver_exe_patterns = ["chromedriver.exe", "chromedriver"]
    has_driver_exe = any(pattern in module_source for pattern in driver_exe_patterns)
    assert has_driver_exe, f"{impl_name} should reference the chromedriver executable"

pytest
pytest-mock
undetected-chromedriver
setuptools
selenium

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