import pandas as pd
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)
    # Example augmentation: horizontal flip
    augmented_images = [ToTensor()(image).flip(-1) for image in images]
    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
):
    """
    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,
    )
    # (query_length, n_patches_x, n_patches_y)
    original_maps = original_batched_maps[0].permute(0, 2, 1).contiguous()
    return original_maps, original_image_embeddings, original_query_embeddings

def visualize_token_map(image, original_maps, token_list, token_index=2, cmap="Greens", figsize=(15, 2), show_text=True):
    """
    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, 2))

    # 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")

    if(show_text):
        # 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.astype(np.float32), 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 evaluate_image_maps(similarity_map, real_image):
    """
    Evaluates the quality of similarity maps by comparing them to a real image.

    This function assesses the alignment between a similarity map and a corresponding
    real image. It calculates several metrics:

    - Accuracy: Checks if any of the maximum values in the similarity map overlap with
      non-zero pixels in the real image (converted to grayscale).
    - Score: Computes a normalized score by summing the element-wise product of the
      similarity map and the normalized grayscale image, divided by the sum of the
      grayscale image pixel values.  This measures the weighted overlap, giving more
      importance to brighter regions in the real image.
    - Rank: Determines the rank of the average value within the special patch in the sorted
      list of all values in the similarity map. This indicates how strongly the map
      highlights the special patch compared to other regions.

    Args:
        similarity_map (np.ndarray): The similarity map to evaluate.
        real_image (PIL.Image.Image): The corresponding real image.

    Returns:
        dict: A dictionary containing the calculated metrics: accuracy, score, and rank.
    """
    # Convert the real image to a binary array (1 - normalized grayscale)
    image_array = 1 - np.array(real_image.convert("L"), dtype=np.float32) / 255.0

    # Create a mask for the maximum values in the similarity map
    acc_visual_map = np.where(similarity_map == similarity_map.max(), similarity_map, 0)
    visual_map = np.copy(similarity_map)

    # Check if scaling is necessary
    if image_array.shape != visual_map.shape:
        scale_factor = image_array.shape[0] // visual_map.shape[0]
        scaled_visual_map = np.kron(
            np.abs(visual_map), np.ones((scale_factor, scale_factor))
        )
        rank_map = np.kron(np.abs(visual_map), np.ones((scale_factor, scale_factor)))
        acc_visual_map = np.kron(
            np.abs(acc_visual_map), np.ones((scale_factor, scale_factor))
        )
    else:
        scaled_visual_map = visual_map

    # Calculate accuracy and score
    accuracy = np.any(image_array * acc_visual_map)
    score = np.sum(image_array * scaled_visual_map) / (
        np.sum(image_array) + 1e-8
    )  # Avoid division by zero
    bin_image = (image_array != 0).astype(int)
    rank = np.sum(bin_image * rank_map) / np.sum(bin_image)  # Avoid division by zero
    rank = np.where(
        np.isclose(sorted(list(np.abs(similarity_map.ravel())))[::-1], rank)
    )[0][0]

    return {
        "accuracy": accuracy,
        "score": score,
        "rank": rank,
    }

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,
    # Added font_path parameter with default value
    font_path="./fonts/Roboto-Regular.ttf",
):
    """
    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(font_path, font_size)
    except IOError:
        print(f"Error loading font from {font_path}. Using default font.")
        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

def visualize_results_grid(results_df):
    columns = [results_df.iloc[:, i] for i in range(len(results_df.columns))]
    columns = [
        (
            pd.to_numeric(col, errors="coerce")
            if not pd.api.types.is_numeric_dtype(col)
            else col
        )
        for col in columns
    ]

    # Deduce the grid shape from the number of results rows
    grid_size = int(np.sqrt(len(results_df)))
    # Reshape columns into matrices
    matrices = [col.to_numpy().reshape(grid_size, grid_size) for col in columns]

    # Visualization setup
    fig, axes = plt.subplots(1, len(results_df.columns), figsize=(12, 2))
    titles = [
        (
            f"{results_df.columns[i]} (Categorical/Binary)"
            if i == 0
            else f"{results_df.columns[i]} (Continuous)"
        )
        for i in range(len(results_df.columns))
    ]
    # Added colormap for the fourth plot
    cmaps = ["coolwarm"] * len(results_df.columns)
    # Plot each matrix
    for i, (matrix, ax, title, cmap) in enumerate(zip(matrices, axes, titles, cmaps)):
        im = ax.imshow(matrix, cmap=cmap, interpolation="none")
        ax.set_title(title)
        ax.set_xticks(range(grid_size))
        ax.set_yticks(range(grid_size))
        fig.colorbar(im, ax=ax)

    # Display the plot
    plt.tight_layout()
    plt.show()

def run_expe_word_square(
    word_to_write,
    token,
    n_patches_x,
    n_patches_y,
    patch_size,
    model,
    processor,
    device,
    use_qwen,
    main_color=[255, 255, 255],
    special_color=(0, 0, 0),
):

    all_images_text = [
        create_single_patch_image_with_text(
            n_patches_x=n_patches_x,
            n_patches_y=n_patches_y,
            patch_size=patch_size,
            main_color=main_color,
            special_color=main_color,
            special_patch=(row, col),
            text=word_to_write,
            text_color=(0,0,0),  # text_color,
            font_size=9,
        )
        for row in range(0, n_patches_y, 2)
        for col in range(0, n_patches_x, 2)
    ]

    all_maps = []
    for image in all_images_text:
        batch_images = processor.process_images([image]).to(device)
        batch_queries = processor.process_queries([token]).to(device)
        original_maps, original_image_embeddings, original_query_embeddings = (
            get_maps_and_embeds(
                batch_images, batch_queries, model, processor, image, use_qwen=use_qwen
            )
        )
        original_maps = original_maps.to(dtype=torch.float32).cpu().numpy()
        all_maps.append(original_maps)

    input_ids = batch_queries["input_ids"][0]  # shape: (num_subtokens,)
    token_list = [processor.tokenizer.decode([token_id]) for token_id in input_ids]
    # print(token_list)
    indexes = [i for i, x in enumerate(token_list) if "<" not in x and ">" not in x][2:]
    # print(indexes)
    # print(np.array(token_list)[[indexes]])

    results_df = pd.DataFrame(columns=["accuracy", "score", "rank"])
    for i, (this_map, image) in enumerate(zip(all_maps, all_images_text)):
        visual_map = this_map[token_index]
        metrics = evaluate_image_maps(visual_map, image)
        results_df.loc[i] = metrics.values()
    return results_df

Instruction: --------------------------------------------------------------------------- AttributeError                            Traceback (most recent call last) Cell In[26], line 24      20 visual_map = this_map[token_index]      22 print(visual_map.shape, patch_mask.shape) ---> 24 metrics = evaluate_map_quality(visual_map, patch_mask)      25 results_df.loc[i] = metrics.values()      26 # Display results  Cell In[25], line 16, in evaluate_map_quality(similarity_map, patch_mask)      14 patch_mask_flat = patch_mask.flatten()      15 # (A) Correlation ---> 16 correlation = np.corrcoef(sim_map_flat.astype(np.float32), patch_mask_flat)[0, 1]      17 # (B) Peak Signal Location      18 max_location = np.unravel_index(np.argmax(similarity_map), similarity_map.shape)  AttributeError: 'Tensor' object has no attribute 'astype'

import pytest
import numpy as np
import torch
from unittest.mock import Mock, patch

def mock_module_dependencies(module):
    """Mock any missing dependencies in the module"""
    # Mock colpali_engine.interpretability imports if they don't exist
    if not hasattr(module, "get_similarity_maps_from_embeddings") and hasattr(
        module, "get_maps_and_embeds"
    ):
        # Create a mock for get_similarity_maps_from_embeddings
        mock_get_maps = Mock()
        mock_get_maps.return_value = [torch.rand(1, 5, 5)]  # Return random tensor
        module.get_similarity_maps_from_embeddings = mock_get_maps
        print("Mocked get_similarity_maps_from_embeddings function")

def test_evaluate_map_quality_with_tensor_input(implementation):
    """
    Test that evaluate_map_quality correctly handles tensor inputs.
    """
    impl_name, module = implementation

    # First, mock any missing dependencies
    mock_module_dependencies(module)

    # Get the original function
    original_func = module.evaluate_map_quality

    # Define a patched version that handles tensor inputs
    def patched_evaluate_map_quality(similarity_map, patch_mask):
        """Patched version to handle tensor inputs"""
        # Convert tensor to numpy if needed
        if isinstance(similarity_map, torch.Tensor):
            similarity_map = similarity_map.detach().cpu().numpy()
        # Call the original function with numpy arrays
        return original_func(similarity_map, patch_mask)

    # Temporarily replace the function
    module.evaluate_map_quality = patched_evaluate_map_quality

    try:
        # Test with tensor input
        similarity_map = torch.tensor([[0.1, 0.2], [0.3, 0.4]])
        patch_mask = np.array([[0, 0], [0, 1]])

        # Run the function
        result = module.evaluate_map_quality(similarity_map, patch_mask)

        # Check result structure
        assert isinstance(result, dict)
        assert "correlation" in result
        assert "peak_accuracy" in result
        assert "overlap_score" in result

        # Test with input similar to what caused the original error
        token_index = 2
        this_map = torch.rand(10, 5, 5)
        visual_map = this_map[token_index]

        patch_mask = np.zeros((5, 5))
        patch_mask[2, 3] = 1

        # This should now work with our patch
        result = module.evaluate_map_quality(visual_map, patch_mask)
        assert isinstance(result, dict)

        print(f"Tensor input test passed for {impl_name}")
    finally:
        # Restore the original function
        module.evaluate_map_quality = original_func

def test_evaluate_map_quality_with_numpy_input(implementation):
    """
    Test that evaluate_map_quality works correctly with numpy arrays.
    """
    impl_name, module = implementation

    # First, mock any missing dependencies
    mock_module_dependencies(module)

    # Get the original function
    original_func = module.evaluate_map_quality

    # Define a patched version that handles tensor inputs
    def patched_evaluate_map_quality(similarity_map, patch_mask):
        """Patched version to handle tensor inputs"""
        # Convert tensor to numpy if needed
        if isinstance(similarity_map, torch.Tensor):
            similarity_map = similarity_map.detach().cpu().numpy()
        # Call the original function with numpy arrays
        return original_func(similarity_map, patch_mask)

    # Temporarily replace the function
    module.evaluate_map_quality = patched_evaluate_map_quality

    try:
        # Test with numpy array input
        similarity_map = np.array([[0.1, 0.2], [0.3, 0.9]])
        patch_mask = np.array([[0, 0], [0, 1]])

        # Run the function
        result = module.evaluate_map_quality(similarity_map, patch_mask)

        # Check result structure
        assert isinstance(result, dict)
        assert "correlation" in result
        assert "peak_accuracy" in result
        assert "overlap_score" in result

        # Check values make sense
        assert -1 <= result["correlation"] <= 1
        assert result["peak_accuracy"] in [0, 1]
        assert result["overlap_score"] > 0

        print(f"NumPy input test passed for {impl_name}")
    finally:
        # Restore the original function
        module.evaluate_map_quality = original_func

numpy
torch
pytest
pytest-mock
matplotlib
pandas
pillow
scikit-learn
seaborn

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output


import pygame
import math

# Initialize Pygame
pygame.init()

# Set up display
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Basic Screen')

# Main loop
running = True
while running:
import pygame
import math

# Initialize Pygame
pygame.init()

# Set up display
screen_width = 800
screen_height = 800
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Flappy Bird Clone')

# Bird properties
bird_x = 100
bird_y = screen_height // 2
bird_width = 40
bird_height = 30
bird_flap_height = 10
gravity = 1
bird_velocity = 0

# Main loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                bird_velocity = -bird_flap_height

    # Apply gravity
    bird_velocity += gravity
    bird_y += bird_velocity

    # Fill the screen with a color (e.g., black)
    screen.fill((0, 0, 0))

    # Draw the bird
    pygame.draw.rect(screen, (255, 255, 0), (bird_x, bird_y, bird_width, bird_height))

    # Update the display
    pygame.display.flip()

# Quit Pygame
pygame.quit()
# Set up display
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Basic Screen')
screen_height = 800
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Flappy Bird Clone')

# Bird properties
bird_x = 100
bird_y = screen_height // 2
bird_width = 40
bird_height = 30
bird_flap_height = 10
gravity = 1
bird_velocity = 0

# Main loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                bird_velocity = -bird_flap_height

    # Apply gravity
    bird_velocity += gravity
    bird_y += bird_velocity

    # Fill the screen with a color (e.g., black)
    screen.fill((0, 0, 0))

    # Draw the bird
    pygame.draw.rect(screen, (255, 255, 0), (bird_x, bird_y, bird_width, bird_height))

    # Update the display
    pygame.display.flip()

# Quit Pygame

        if event.type == pygame.QUIT:
            running = False

    # Fill the screen with a color (e.g., black)
    screen.fill((0, 0, 0))

    # Update the display
    pygame.display.flip()

# Quit Pygame
pygame.quit()

Instruction: fix the quitting and the start of the main loop

import unittest.mock
import sys
import ast
import pytest
from test_utils import TestUtils

def test_pygame_init_present(implementation):
    """Test that the code initializes pygame"""
    impl_name, module = implementation

    # Get source code without executing the module
    source_code = TestUtils.get_source_code(module)
    assert "pygame.init()" in source_code, f"{impl_name}: Pygame initialization missing"

def test_single_game_loop(implementation):
    """Test that there is only one game loop in the code"""
    impl_name, module = implementation

    # Skip if module has errors
    if hasattr(module, "__error__"):
        pytest.skip(f"Module has errors: {module.__error__}")

    # Extract the source code and parse the AST
    source_code = TestUtils.get_source_code(module)
    tree = ast.parse(source_code)

    # Count the number of while loops with 'running' condition
    while_loops = [node for node in ast.walk(tree) if isinstance(node, ast.While)]
    while_running_loops = [
        loop
        for loop in while_loops
        if isinstance(loop.test, ast.Name) and loop.test.id == "running"
    ]

    assert (
        len(while_running_loops) == 1
    ), f"{impl_name}: There should be exactly one main game loop"

def test_proper_game_loop_execution(implementation):
    """Test that the game loop runs properly and quits correctly when requested"""
    impl_name, module = implementation

    # Skip if module has errors
    if hasattr(module, "__error__"):
        pytest.skip(f"Module has errors: {module.__error__}")

    # Run the module in a subprocess
    result = TestUtils.run_module_in_subprocess(module)

    # Check results
    if not result["success"]:
        pytest.fail(f"{impl_name}: {result['error']}")

    assert result["quit_called"], f"{impl_name}: pygame.quit() was not called"

def test_duplicated_code_removed(implementation):
    """Test that duplicate code has been removed"""
    impl_name, module = implementation

    # Skip if module has errors
    if hasattr(module, "__error__"):
        pytest.skip(f"Module has errors: {module.__error__}")

    # Extract the source code
    source_code = TestUtils.get_source_code(module)

    # Count occurrences of certain key lines to check for duplicates
    pygame_init_count = source_code.count("pygame.init()")
    pygame_quit_count = source_code.count("pygame.quit()")

    assert (
        pygame_init_count == 1
    ), f"{impl_name}: pygame.init() should appear exactly once"
    assert (
        pygame_quit_count == 1
    ), f"{impl_name}: pygame.quit() should appear exactly once"

def test_only_one_flappy_bird_game(implementation):
    """Test that there's only one version of the Flappy Bird game in the code"""
    impl_name, module = implementation

    # Skip if module has errors
    if hasattr(module, "__error__"):
        pytest.skip(f"Module has errors: {module.__error__}")

    # Extract the source code
    source_code = TestUtils.get_source_code(module)

    # Count occurrences of the caption setting
    flappy_bird_caption_count = source_code.count(
        "pygame.display.set_caption('Flappy Bird Clone')"
    )

    assert (
        flappy_bird_caption_count == 1
    ), f"{impl_name}: 'Flappy Bird Clone' caption should appear exactly once"

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

# 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

# Setup for Pygame tests - initialize once per session
@pytest.fixture(scope="session", autouse=True)
def setup_pygame():
    """Initialize pygame once at the start of the session."""
    # Initialize pygame in headless mode if no display is available
    if os.environ.get("CI") or not os.environ.get("DISPLAY"):
        os.environ["SDL_VIDEODRIVER"] = "dummy"

    # Initialize pygame
    pygame.init()

    # Clean up at the end of the session
    yield
    pygame.quit()

# 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."""
    # Make sure pygame is properly cleaned up
    pygame.quit()

    # Save test results
    test_results.save_results()

import os
import sys
import glob
import re
import importlib.util
import traceback
import types
import subprocess
import tempfile
import json
from typing import Dict, List, Optional, Any, Tuple
import pygame
import threading
import time
import inspect

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_without_execution(
        file_path: str, module_name: Optional[str] = None
    ) -> Any:
        """
        Load a module from a file path WITHOUT executing its code.
        This prevents pygame windows from opening during module loading.
        """
        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:
            # Read the source code
            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 a new module object
            module = types.ModuleType(unique_module_name)
            module.__file__ = file_path
            module.__source_code__ = source_code  # Store source code for inspection
            module.__display_name__ = module_name

            # Add the module to sys.modules
            sys.modules[unique_module_name] = module

            return 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 without executing them."""
        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_without_execution(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

    @staticmethod
    def get_source_code(module):
        """Get the source code of a module."""
        # First try to get it from our stored attribute
        if hasattr(module, "__source_code__"):
            return module.__source_code__

        # If that fails, try to use inspect
        try:
            return inspect.getsource(module)
        except Exception as e:
            raise ValueError(f"Could not get source code: {e}")

    @staticmethod
    def run_module_in_subprocess(module, timeout=5.0):
        """Run a module in a subprocess with a timeout and check if it closes properly."""
        # Get source code
        try:
            source_code = TestUtils.get_source_code(module)
        except Exception as e:
            return {
                "success": False,
                "error": f"Could not get source code: {e}",
                "quit_called": False,
            }

        # Indent source code for inclusion in the wrapper script
        indented_source = "\n".join("    " + line for line in source_code.splitlines())

        # Create a wrapper script that will run the module and check if pygame.quit() is called
        wrapper_code = """
import sys
import pygame
import time

# Track if pygame.quit is called
original_quit = pygame.quit
quit_called = False

def mock_quit():
    global quit_called
    quit_called = True
    original_quit()

pygame.quit = mock_quit

# Set up automated event injection
def post_quit_event():
    try:
        pygame.event.post(pygame.event.Event(pygame.QUIT))
    except Exception as e:
        print(f"Error posting event: {{e}}")

# Use a timer to post a quit event after 1 second
import threading
timer = threading.Timer(1.0, post_quit_event)
timer.daemon = True
timer.start()

# Execute the module code
try:
{}
except SystemExit:
    pass
except Exception as e:
    print(f"ERROR: {{e}}")
    sys.exit(1)

# Report results
print(f"QUIT_CALLED: {{quit_called}}")
sys.exit(0)
""".format(
            indented_source
        )

        # Create temporary file with the wrapped code
        with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as temp_file:
            temp_file_path = temp_file.name
            temp_file.write(wrapper_code.encode("utf-8"))

        try:
            # Set environment variable to use dummy video driver (headless mode)
            env = os.environ.copy()
            env["SDL_VIDEODRIVER"] = "dummy"

            # Run the wrapper script in a subprocess
            process = subprocess.Popen(
                [sys.executable, temp_file_path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=env,
            )

            # Wait for the process with timeout
            try:
                stdout, stderr = process.communicate(timeout=timeout)
                stdout = stdout.decode("utf-8")
                stderr = stderr.decode("utf-8")

                # Check if there was an error
                if process.returncode != 0:
                    return {
                        "success": False,
                        "error": f"Process exited with code {process.returncode}: {stderr}",
                        "quit_called": False,
                    }

                # Check if pygame.quit() was called
                quit_called = "QUIT_CALLED: True" in stdout

                return {"success": True, "error": None, "quit_called": quit_called}

            except subprocess.TimeoutExpired:
                # Kill the process if it times out
                process.kill()
                return {
                    "success": False,
                    "error": f"Process timed out after {timeout} seconds",
                    "quit_called": False,
                }
        finally:
            # Clean up the temporary file
            try:
                os.unlink(temp_file_path)
            except Exception:
                pass

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output


from pyspark.sql import SparkSession
from pyspark.sql.functions import col

# Crear una sesión Spark
spark = SparkSession.builder.appName("EscrituraParquetADL2").getOrCreate()

# Ejemplo de datos (reemplazar con tus datos reales)
data = [
    {"id": 1, "nombre": "Juan", "edad": 30},
    {"id": 2, "nombre": "Ana", "edad": 25},
    {"id": 3, "nombre": "Pedro", "edad": 40}
]

# Crear un DataFrame a partir de los datos

# Configurar la conexión a ADL2 usando la identidad de Microsoft ID
# No es necesario proporcionar credenciales explícitamente en un notebook de Synapse
# Spark utilizará la identidad administrada del notebook para autenticarse.

# Especificar la ruta al contenedor y la carpeta en ADL2
container_name = "<your_container_name>"  # Reemplazar con el nombre de tu contenedor
folder_path = "<your_folder_path>"  # Reemplazar con la ruta a la carpeta dentro del contenedor
adl2_path = f"abfss://{container_name}@{<your_storage_account_name>}.dfs.core.windows.net/{folder_path}"

# Escribir el DataFrame en formato parquet en ADL2
df.write.parquet(adl2_path, mode="overwrite")

# Opcional: leer el archivo parquet para verificar
df_leido = spark.read.parquet(adl2_path)
df_leido.show()

# Detener la sesión Spark
spark.stop()

Instruction: No se como vienen los datos porque provienen de una api. Primero tengo que analizarlos

import pytest
import inspect
import re
from unittest.mock import patch, MagicMock, ANY
import json

def test_includes_data_analysis(implementation):
    """Test that the implementation includes data analysis functionality."""
    impl_name, module = implementation
    source_code = inspect.getsource(module)

    # Define various ways to interact with the data variable
    interaction_patterns = [
        r'\bdata\s*\[',                   # data[...] access
        r'\bdata\s*\.',                   # data.method or data.attribute (not common unless it's a custom object)
        r'for\s+\w+\s+in\s+data',         # iterating over data
        r'len\s*\(\s*data\s*\)',          # checking length
        r'isinstance\s*\(\s*data',        # type checking
        r'pd\.DataFrame\s*\(\s*data'      # creating a DataFrame
    ]

    # At least one form of analysis should be present
    assert any(re.search(pattern, source_code) for pattern in interaction_patterns), \
        f"{impl_name} should include at least one interaction with the data variable."

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output


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

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

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

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

        cnt = 0
        for replies in post_replies:

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

Instruction: Quiero crear un diccionario jerarquico que sea una lista de diccionario [{post_generator}:[{replies}]]

import pytest
import inspect
import re
import sys
from unittest.mock import MagicMock, patch

# Mock the beem modules since they're not available
sys.modules['beem'] = MagicMock()
sys.modules['beem.discussions'] = MagicMock()
sys.modules['beem.comment'] = MagicMock()
sys.modules['random'] = MagicMock()
sys.modules['string'] = MagicMock()

# Create mock classes with more comprehensive behavior
class MockComment:
    def __init__(self, **kwargs):
        self.data = kwargs

    def get_replies(self):
        # Simulate different replies
        return [
            {'author': 'user1', 'body': 'count me in', 'permlink': 'reply1'},
            {'author': 'user2', 'body': 'count me too', 'permlink': 'reply2'},
            {'author': 'user3', 'body': 'count me please', 'permlink': 'reply3'},
            {'author': 'user4', 'body': 'random text', 'permlink': 'reply4'},
            {'author': 'subidu', 'body': 'I am the author', 'permlink': 'reply5'},
        ]

class MockQuery:
    def __init__(self, **kwargs):
        self.params = kwargs

class MockDiscussions:
    def get_discussions(self, sort, query, limit=None):
        # Return a list of mock posts with varying properties
        return [
            {
                'author': 'post_author1',
                'permlink': 'post1',
                'children': 10,
                'title': 'Test Post 1'
            },
            {
                'author': 'post_author2',
                'permlink': 'post2',
                'children': 3,
                'title': 'Test Post 2'
            },
            {
                'author': 'post_author3',
                'permlink': 'post3',
                'children': 20,
                'title': 'Test Post 3'
            }
        ]

# Update mock modules with enhanced mock classes
sys.modules['beem.discussions'].Discussions = MockDiscussions
sys.modules['beem.discussions'].Query = MockQuery
sys.modules['beem.comment'].Comment = MockComment

# Patch random and string modules
sys.modules['random'].choices = lambda chars, k: ['1'] * k
sys.modules['string'].digits = '0123456789'

def test_hierarchical_structure_implementation(implementation):
    """Test that the implementation creates a hierarchical data structure."""
    impl_name, module = implementation

    source_code = inspect.getsource(module)

    # Check for evidence of hierarchical structure - expanded and improved patterns
    hierarchy_patterns = [
        # Dictionary with post_author/post_permlink as key
        r"diccionario\[\s*f[\"']?{.*?post_author.*?post_permlink",
        r"post_replies_dict\[\s*f[\"']?{.*?post_author.*?post_permlink",

        # Nested data structure with post and replies fields
        r"['\"]post['\"]\s*:.*?['\"]replies['\"]\s*:",
        r"post_data\s*=\s*{.*?post.*?replies.*?}",

        # Lists of dictionaries or nested structures
        r"append\(\s*{.*?['\"]post['\"]\s*:.*?['\"]replies['\"]\s*:",
        r"diccionario\.append\(\s*{.*?post.*?replies",

        # Dictionary assignment with list of replies
        r"diccionario\[.*?\]\s*=\s*.*?replies",

        # Other hierarchical patterns
        r"hierarchical_data\s*=",
        r"post_data\[['\"](replies|post)['\"]",
        r"post_data\[['\"]replies['\"]\]\.append"
    ]

    has_hierarchical_structure = any(re.search(pattern, source_code, re.DOTALL)
                                    for pattern in hierarchy_patterns)

    assert has_hierarchical_structure, \
        f"{impl_name} should implement a hierarchical structure to store posts and replies"

pytest
pytest-mock
pycryptodome
cryptography
pyscrypt
argon2-cffi
websockets
backports.zoneinfo;python_version<"3.9"
beem

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output