import json
import random
from groq import Groq
import os
from dotenv import load_dotenv

load_dotenv()

api_key = os.getenv("GROQ_API_KEY")
client = Groq(api_key=api_key)

class SCP_Object:
    def __init__(self, name, description, triggers, initial_state):
        self.name = name
        self.description = description
        self.triggers = triggers
        self.state = initial_state

class D_Personnel:
    def __init__(self, name, initial_state):
        self.name = name
        self.state = initial_state
        self.memory = [] # List of strings, what he remembers from experiments

class Room:
    def __init__(self, name, description, initial_state):
        self.name = name
        self.description = description
        self.state = initial_state

class Game:
    def __init__(self):
        self.scp_objects = self.load_scp_objects()
        self.current_scp = random.choice(self.scp_objects)
        self.d_personnel = D_Personnel("D-" + str(random.randint(1000, 9999)), initial_state="calm")
        self.room = Room("Test Chamber", "A standard containment chamber.", initial_state="clean")
        self.player_report = ""
        self.experiment_log = []

    def load_scp_objects(self):
        # Example SCP objects, can be loaded from a file later
        return [
            SCP_Object(
                name="SCP-173",
                description="SCP-173 is a concrete statue that moves when not directly observed.",
                triggers={
                    "not_observed": "SCP-173 moves quickly towards the nearest person.",
                    "touch": "SCP-173 does nothing.",
                },
                initial_state="immobile"
            ),
            SCP_Object(
                name="SCP-096",
                description="SCP-096 is a humanoid creature that becomes extremely hostile when its face is viewed.",
                triggers={
                    "view_face": "SCP-096 will scream and chase the viewer.",
                    "touch": "SCP-096 does nothing.",
                },
                initial_state="docile"
            ),
            SCP_Object(
                name="SCP-999",
                description="SCP-999 is a large, amorphous, gelatinous mass of translucent orange slime with a consistency similar to that of peanut butter. SCP-999's behavior is best described as playful and dog-like.",
                triggers={
                    "touch": "SCP-999 will make a happy sound and try to hug the person.",
                    "attack": "SCP-999 will try to hug the person.",
                },
                initial_state="happy"
            )
        ]

    def llm_request(self, system_prompt, user_prompt):
        completion = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": user_prompt
                },
            ],
            temperature=0.02,
            max_tokens=1024,
            top_p=1,
            stream=False,
            response_format={"type": "json_object"},
            stop=None,
        )
        try:
            return json.loads(completion.choices[0].message.content)
        except json.JSONDecodeError:
            print(f"Error decoding JSON: {completion.choices[0].message.content}")
            return {"message": "Error decoding JSON", "intention": "error"}

    def check_instruction_clarity(self, instruction):
        system_prompt = "You are a helpful assistant that checks if the instruction to people is clear and understandable. Provide a json response with the following fields: 'reasoning' (a step-by-step analysis of the instruction), 'decision' (True if the instruction is clear, False otherwise), and 'why' (a brief explanation of your decision)."
        user_prompt = f"Analyze the clarity of this instruction: {instruction}"
        response = self.llm_request(system_prompt, user_prompt)
        return response

    def check_instruction_possibility(self, instruction):
        system_prompt = "You are a helpful assistant that checks if the instruction is physically possible. Provide a json response with the following fields: 'reasoning' (a step-by-step analysis of the instruction's feasibility), 'decision' (True if the instruction is possible, False otherwise), and 'why' (a brief explanation of your decision)."
        user_prompt = f"Analyze the possibility of this instruction: {instruction}. Current room state: {self.room.state}"
        response = self.llm_request(system_prompt, user_prompt)
        return response

    def check_d_willingness(self, instruction):
        system_prompt = "You are a D-class personnel in a test chamber with an unknown SCP object. Provide a json response with the following fields: 'reasoning' (a step-by-step analysis of your willingness), 'decision' (True if you agree to perform the instruction, False otherwise), and 'why' (a brief explanation of your decision)."
        user_prompt = f"Will you perform this instruction? Instruction: {instruction}. Your current state: {self.d_personnel.state}. Your current knowledge about SCP: {self.d_personnel.memory}"
        response = self.llm_request(system_prompt, user_prompt)
        return response
    def emulate_experiment(self, instruction):
        self.experiment_log.append(f"Instruction: {instruction}")

        # Initial checks
        clarity_check = self.check_instruction_clarity(instruction)
        if clarity_check["decision"] == "False":
            self.experiment_log.append(f"Instruction not clear: {clarity_check['why']}")
            return f"Instruction not clear. Please clarify. Here is why: {clarity_check['why']}"

        possibility_check = self.check_instruction_possibility(instruction)
        if possibility_check["decision"] == "False":
            self.experiment_log.append(f"Instruction impossible: {possibility_check['why']}")
            return f"Instruction impossible. Please provide a possible instruction. Here is why: {clarity_check['why']}"

        willingness_check = self.check_d_willingness(instruction)
        if willingness_check["decision"] == "False":
            self.experiment_log.append(f"D-personnel refused: {willingness_check['why']}")
            return f"D-personnel refused. Reason: {willingness_check['why']}"

        self.experiment_log.append("All checks passed. Starting emulation.")

        # Emulation loop
        current_actor = "d_personnel"
        count_of_iterations = 0
        action_history = [instruction]  # Start with the initial instruction

        while True and count_of_iterations < 5:
            count_of_iterations += 1
            if current_actor == "d_personnel":
                actions = self.generate_possible_actions(action_history, self.d_personnel)
                if not actions:
                    self.experiment_log.append("No possible actions for D-personnel.")
                    break

                chosen_action = self.choose_action(actions)
                self.experiment_log.append(f"D-personnel action: {chosen_action}")

                outcomes = self.generate_outcomes(chosen_action, self.d_personnel, self.current_scp, self.room)
                self.experiment_log.append(f"Outcomes: {outcomes}")

                self.apply_outcomes(outcomes)
                action_history.append({"d_personnel": {"action": chosen_action, "outcomes": outcomes}})

                current_actor = "scp"
            elif current_actor == "scp":
                scp_actions = self.generate_possible_actions(action_history, self.current_scp)
                if not scp_actions:
                    self.experiment_log.append("No possible actions for SCP.")
                    break

                chosen_scp_action = self.choose_action(scp_actions)
                self.experiment_log.append(f"SCP action: {chosen_scp_action}")

                scp_outcomes = self.generate_outcomes(chosen_scp_action, self.d_personnel, self.current_scp, self.room)
                self.experiment_log.append(f"SCP Outcomes: {scp_outcomes}")

                self.apply_outcomes(scp_outcomes)
                action_history.append({"scp": {"action": chosen_scp_action, "outcomes": scp_outcomes}})

                current_actor = "d_personnel"

    def generate_possible_actions(self, action_history, actor):
        if isinstance(actor, D_Personnel):
            system_prompt = """You are a helpful assistant that generates possible actions for D-class personnel.
            Answer in json format in format: {actions: ["action1", "action2", ... ]}.
            Generate 3-5 possible actions based on the instruction and action history."""

            user_prompt = f"""Generate possible actions for D-class personnel based on this history:
            Initial instruction: {action_history[0]}
            Action history: {action_history[1:]}
            Current state: {actor.state}
            Current knowledge about SCP: {actor.memory}"""
            response = self.llm_request(system_prompt, user_prompt)
            if "actions" in response:
                return response["actions"]
            else:
                return []
        elif isinstance(actor, SCP_Object):

    def choose_action(self, actions):
        if not actions:
            return None
        return random.choice(actions)

    def generate_outcomes(self, action, d_personnel, scp_object, room):
        system_prompt = "You are a helpful assistant that generates possible outcomes of an action. Answer in json format in format: {outcomes: [{\"description\": \"\", \"d_personnel_state\": \"\", \"scp_state\": \"\", \"room_state\": \"\"}, ... ]}. Generate 3-5 possible outcomes based on the action and current state."
        user_prompt = f"Generate possible outcomes for this action: {action}. D-personnel state: {d_personnel.state}. SCP state: {scp_object.state}. Room state: {room.state}. SCP description: {scp_object.description}. SCP triggers: {scp_object.triggers}"
        response = self.llm_request(system_prompt, user_prompt)
        if "outcomes" in response:
            return response["outcomes"]
        else:
            return []

    def apply_outcomes(self, outcomes):
        if not outcomes:
            return
        chosen_outcome = random.choice(outcomes)
        self.experiment_log.append(f"Chosen outcome: {chosen_outcome}")
        if "d_personnel_state" in chosen_outcome:
            self.d_personnel.state = chosen_outcome["d_personnel_state"]
        if "scp_state" in chosen_outcome:
            self.current_scp.state = chosen_outcome["scp_state"]
        if "room_state" in chosen_outcome:
            self.room.state = chosen_outcome["room_state"]

    def get_d_report(self):
        if self.d_personnel.state == "dead":
            return "D-personnel is dead. No report available."
        system_prompt = "You are a D-class personnel. You need to describe what happened during the experiment. Answer in json format in format: {report: \"\"}. Describe what you remember from the experiment."
        user_prompt = f"Describe what happened during the experiment. Your current state: {self.d_personnel.state}. Your current knowledge about SCP: {self.d_personnel.memory}. Experiment log: {self.experiment_log}"
        response = self.llm_request(system_prompt, user_prompt)
        if "report" in response:
            self.d_personnel.memory.append(response["report"])
            return response["report"]
        else:
            return "No report available."

    def get_d_death_report(self):
        if self.d_personnel.state != "dead":
            return "D-personnel is alive. No death report available."
        system_prompt = "You are a forensic expert. You need to describe the state of the dead D-class personnel. Answer in json format in format: {report: \"\"}. Describe the state of the body."
        user_prompt = f"Describe the state of the dead D-class personnel. Experiment log: {self.experiment_log}"
        response = self.llm_request(system_prompt, user_prompt)
        if "report" in response:
            return response["report"]
        else:
            return "No death report available."

    def start_experiment(self, instruction):
        self.experiment_log = []
        result = self.emulate_experiment(instruction)
        if result:
            return result
        if self.d_personnel.state == "dead":
            report = self.get_d_death_report()
        else:
            report = self.get_d_report()
        self.d_personnel = D_Personnel("D-" + str(random.randint(1000, 9999)), initial_state="calm")
        return report

    def submit_report(self, player_report):
        self.player_report = player_report
        system_prompt = "You are a helpful assistant that checks if the player report is factually correct. Answer in json format in format: {message: \"\", score: 0-100}. Compare the player report with the SCP description and triggers. Score should be 0 if the report is completely wrong and 100 if the report is completely correct."
        user_prompt = f"Compare the player report with the SCP description and triggers. Player report: {player_report}. SCP description: {self.current_scp.description}. SCP triggers: {self.current_scp.triggers}"
        response = self.llm_request(system_prompt, user_prompt)
        return response

    def play(self):
        print("Welcome to SCPIE!")
        instruction = input("Enter instruction for D-personnel: ")
        result = self.start_experiment(instruction)
        print("Experiment result:", result)
        print("\n\n\n")
        for log in self.experiment_log:
            if isinstance(log, dict):
                json.dumps(log, indent=4, ensure_ascii=False)
            else:
                print(log)
            print()
        # print(self.experiment_log)

if __name__ == "__main__":
    game = Game()
    game.play()
    # print("Yes")

Instruction: continue this function with SCP object prompting, just like D_personal

import ast
import inspect
import pytest

def extract_generate_possible_actions_branches(module_code, d_class_name, scp_class_name):
    """Extract the AST bodies of the D_Personnel and SCP_Object branches in generate_possible_actions."""
    try:
        parsed = ast.parse(module_code)
    except SyntaxError as e:
        print(f"SyntaxError while parsing module code: {e}")
        return None, None
    except Exception as e:
        print(f"Unexpected error while parsing module code: {e}")
        return None, None

    for node in ast.walk(parsed):
        if isinstance(node, ast.ClassDef) and node.name == "Game":
            for item in node.body:
                if isinstance(item, ast.FunctionDef) and item.name == "generate_possible_actions":
                    d_branch = None
                    scp_branch = None

                    for subnode in ast.walk(item):
                        if isinstance(subnode, ast.If):
                            test = subnode.test
                            if (
                                isinstance(test, ast.Call)
                                and isinstance(test.func, ast.Name)
                                and test.func.id == "isinstance"
                                and isinstance(test.args[1], ast.Name)
                            ):
                                class_name = test.args[1].id
                                if class_name == d_class_name:
                                    d_branch = subnode.body
                                elif class_name == scp_class_name:
                                    scp_branch = subnode.body

                    return d_branch, scp_branch

    return None, None

def ast_structure_summary(ast_nodes):
    """Extracts structure summary from AST nodes to compare similarity."""
    summary = []
    for node in ast_nodes:
        if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant):
            summary.append(("assign", node.targets[0].id, type(node.value.value).__name__))
        elif isinstance(node, ast.Assign):
            summary.append(("assign", node.targets[0].id, type(node.value).__name__))
        elif isinstance(node, ast.Expr) and isinstance(node.value, ast.Call):
            summary.append(("call", getattr(node.value.func, 'id', 'unknown')))
        elif isinstance(node, ast.Return):
            summary.append(("return",))
        elif isinstance(node, ast.If):
            summary.append(("if",))
        else:
            summary.append((type(node).__name__,))
    return summary

def test_generate_possible_actions_structure_similarity(implementation):
    impl_name, module = implementation
    module_code = inspect.getsource(module)

    # Extract the class names to match against isinstance checks
    d_class_name = "D_Personnel"
    scp_class_name = "SCP_Object"

    # Get the AST branches
    d_branch, scp_branch = extract_generate_possible_actions_branches(module_code, d_class_name, scp_class_name)

    assert d_branch is not None, "Could not extract D_Personnel branch"
    assert scp_branch is not None, "Could not extract SCP_Object branch"

    # Compare structure
    d_summary = ast_structure_summary(d_branch)
    scp_summary = ast_structure_summary(scp_branch)

    assert d_summary == scp_summary, f"Mismatch in structure:\nD: {d_summary}\nSCP: {scp_summary}"

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 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 (torch.Tensor): 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.
    """
    # Ensure similarity_map is in float32 and on the CPU
    similarity_map = similarity_map.to(dtype=torch.float32).cpu().numpy()

    # Flatten the map and mask for easier computation
    sim_map_flat = similarity_map.flatten()
    patch_mask_flat = patch_mask.flatten()

    # Ensure the shapes are compatible
    if sim_map_flat.shape != patch_mask_flat.shape:
        raise ValueError(
            f"Shape mismatch: similarity_map has {sim_map_flat.shape} elements, "
            f"but patch_mask has {patch_mask_flat.shape} elements."
        )

    # (A) Correlation
    correlation = np.corrcoef(
        sim_map_flat, patch_mask_flat.astype(np.float32))[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.

    Args:
        similarity_map (torch.Tensor): 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

    # Ensure similarity_map is float32 and on the CPU before using numpy operations
    similarity_map_cpu = similarity_map.to(dtype=torch.float32).cpu().numpy()

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

    # Check if scaling is necessary
    if image_array.shape != similarity_map_cpu.shape:
        scale_factor = image_array.shape[0] // similarity_map_cpu.shape[0]
        scaled_visual_map = np.kron(
            np.abs(similarity_map_cpu), np.ones((scale_factor, scale_factor))
        )
        rank_map = np.kron(
            np.abs(similarity_map_cpu), 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 = similarity_map_cpu
        rank_map = similarity_map_cpu  # Add this to avoid missing variable

    # 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

    # Calculate rank
    bin_image = (image_array != 0).astype(int)
    rank_value = np.sum(bin_image * rank_map) / np.sum(
        bin_image
    )  # Avoid division by zero
    sorted_values = sorted(np.abs(similarity_map_cpu.ravel()))[::-1]
    rank = np.where(np.isclose(sorted_values, rank_value))[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[indexes[0]]
        metrics = evaluate_image_maps(visual_map, image)
        results_df.loc[i] = metrics.values()
    return results_df

Instruction: add a check to avoid this operation if it is already a numpy format

import ast
import inspect
import pytest

def test_similarity_map_cpu_guarded(implementation):
    """
    Ensure that within `evaluate_image_maps`, the line with
    `similarity_map.to(dtype=torch.float32).cpu().numpy()` is preceded by
    an `if` statement that includes 'np' or 'numpy'.
    """
    impl_name, module = implementation
    module_code = inspect.getsource(module)

    lines = module_code.split('\n')

    # Strip comments and blank lines
    cleaned_lines = []
    for line in lines:
        stripped = line.strip()
        if not stripped or stripped.startswith('#'):
            continue
        # Remove inline comments
        line_no_comment = line.split('#')[0].strip()
        cleaned_lines.append(line_no_comment)

    # Flag to track whether we're inside the evaluate_image_maps function
    inside_target_function = False
    function_lines = []

    for line in cleaned_lines:
        if line.startswith("def evaluate_image_maps("):
            inside_target_function = True
            continue

        # Stop if we’re out of the function by checking indentation
        if inside_target_function:
            # We know we're out of the target function because the original code is succeeded by a new method
            if line.startswith("def ") or line.startswith("class "):
                inside_target_function = False
                break

            function_lines.append(line)

    if not function_lines:
        pytest.fail("Function evaluate_image_maps not found or is empty")

    target_expr = "similarity_map.to(dtype=torch.float32).cpu().numpy()"

    for idx, line in enumerate(function_lines):
        if target_expr in line:
            if idx == 0:
                pytest.fail("Expected 'if' statement before similarity_map conversion, got empty line.")
            prev_line = function_lines[idx - 1].strip()
            assert prev_line.startswith("if"), \
                f"Expected 'if' statement before similarity_map conversion, got: {prev_line}"
            assert "np" in prev_line or "numpy" in prev_line, \
                f"'if' statement before similarity_map conversion does not mention numpy: {prev_line}"
            return

    pytest.fail(f"Could not find line with: {target_expr}")

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

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 ast import Add
from asyncio import wait
from curses import COLOR_BLUE, COLOR_RED
from re import A
from shutil import move
from glm import degrees
from manim import *
from numpy import size, square

class Project(Scene):
    def construct(self):
        text = Tex("Double Angle")
        self.play( Write(text))

        self.wait(5)

        transform_text = Tex("What is Double Angle?")
        transform_text.to_corner(UP)
        box = SurroundingRectangle(transform_text)
        box.set_color(WHITE)
        box.set_stroke(width=1.5)
        self.play(
            Transform(text, transform_text)
        )
        self.wait(0.5)
        self.play(Create(box))

        explanation = Paragraph("A double angle is an angle measurement", "that has been multiplied by 2 or added to itself.", line_spacing=0.5, font_size=32)
        explanation.move_to(ORIGIN)

        self.play(
            Write(explanation)
        )

        self.wait(3)

        self.play(
            Transform(explanation, explanation.copy().shift(UP))
        )

        trig_cos2 = MathTex(
            r"\cos2x = \cos^2x - \sin^2x",

            substrings_to_isolate=["cos2x"]
            )
        trig_cos2.set_color_by_tex("cos2x", BLUE)
        trig_cos2.move_to(DOWN)
        transform_formula = Tex("Double Angle Formula")
        transform_formula.to_corner(UP)

        self.wait(1)

        self.play(
        Write(trig_cos2)
        )

        self.wait(2)

        self.play(
            FadeOut(trig_cos2, explanation)
        )

        self.wait(1)

        axes = Axes(
            x_range=[-2, 2, 2],
            y_range=[-2, 2, 2],
            x_length=4,
            y_length=4,
        )
        self.add(axes)

        # 単位円の作成
        circle = Circle(radius=2, color=BLUE)
        self.add(circle)

        # 原点 (Origin)
        dot = Dot(ORIGIN, color=RED)
        self.add(dot)

        # 角度を表す線分 (Line representing the angle)
        line = Line(ORIGIN, RIGHT * 2)
        self.add(line)

        # 角度のラベル (Angle label)
        # Create an Arc for the angle
        angle = Arc(
            radius=2,
            start_angle=0,  # Start at the positive x-axis
            angle=line.get_angle(),  # Use line's angle
            arc_center=ORIGIN,
            color=GREEN
        )
        angle_label = MathTex(r"\theta = 0^{\circ}").next_to(angle, RIGHT) # Changed Tex to MathTex and added \\
        self.add(angle, angle_label)

        intersection_dot = Dot(color=YELLOW)

        angle_tracker = ValueTracker(0)

        def update_line(mobject):
            mobject.become(Line(ORIGIN, RIGHT * 2).rotate(angle_tracker.get_value(), about_point=ORIGIN))

        def update_angle(mobject):
            mobject.become(Arc(
                radius=2,
                start_angle=0,
                angle=angle_tracker.get_value(),
                arc_center=ORIGIN,
                color=GREEN
            ))

        line.add_updater(update_line)
        angle.add_updater(update_angle)

        # Update the angle label
        def update_label(mobject):
            angle_in_degrees = np.degrees(angle_tracker.get_value())
            mobject.become(MathTex(rf"\\theta = {angle_in_degrees:.0f}^{{\circ}}")) # Added double brackets
            mobject.next_to(angle, RIGHT)

        angle_label.add_updater(update_label)

        def update_intersection_dot(mobject):
            angle = angle_tracker.get_value()
            x = 2 * np.cos(angle)  # x-coordinate on the circle
            y = 2 * np.sin(angle)  # y-coordinate on the circle
            mobject.move_to([x, y, 0])

        intersection_dot.add_updater(update_intersection_dot)

        self.add(intersection_dot)
        # Animate the angle
        self.play(
            angle_tracker.animate.set_value(PI / 6),
            run_time=2
        )
        self.wait(3)

        line.clear_updaters()
        intersection_dot.clear_updaters()
        angle.clear_updaters()
        angle_label.clear_updaters()

        # Change their color to indicate they are fixed
        fixed_line = line.copy().set_color(ORANGE)
        fixed_dot = intersection_dot.copy().set_color(ORANGE)
        fixed_angle = angle.copy().set_color(ORANGE)
        self.add(fixed_line, fixed_dot, fixed_angle)

        # Prepare a new line for the next animation
        new_line = Line(ORIGIN, RIGHT * 2, color=GREEN)
        new_intersection_dot = Dot(color=YELLOW)
        new_angle = Arc(
            radius=0.5,
            start_angle=PI / 6,  # Start from 30 degrees
            angle=0,
            arc_center=ORIGIN,
            color=GREEN
        )
        new_label = MathTex(rf"\theta = 30^\circ").next_to(new_angle, RIGHT).set_color(ORANGE)

        # Updaters for the new objects
        new_line.add_updater(lambda m: m.become(
            Line(ORIGIN, RIGHT * 2).rotate(angle_tracker.get_value(), about_point=ORIGIN)
        ))

        new_intersection_dot.add_updater(lambda m: m.move_to([
            2 * np.cos(angle_tracker.get_value()),
            2 * np.sin(angle_tracker.get_value()),
            0
        ]))

        new_angle.add_updater(lambda m: m.become(
            Arc(
                radius=0.5,
                start_angle=0,
                angle=angle_tracker.get_value(),
                arc_center=ORIGIN,
                color=GREEN
            )
        ))

        new_label.add_updater(lambda m: m.become(
            MathTex(rf"\theta = {np.degrees(angle_tracker.get_value()):.0f}^\circ").next_to(new_angle, LEFT)
        ))

        # Add the new objects
        self.add(new_line, new_intersection_dot, new_angle, new_label)

        # Animate from 30 degrees to 60 degrees
        self.play(
            angle_tracker.animate.set_value(PI / 3),  # 60 degrees
            run_time=2
        )
        self.wait(1)

        self.wait(10)

        self.play(
            FadeOut(circle, dot, line, angle, angle_label, axes, line, angle, intersection_dot, angle_label, new_line, new_angle, new_label, new_intersection_dot, fixed_line, fixed_angle, fixed_dot, angle_tracker)
        )

        self.play(
            FadeOut(transform_text, explanation),
            Transform(trig_cos2 , trig_cos2.copy().shift(UP + UP + UP)),
            Transform(text, transform_formula),
        )
        self.wait(2)

        cos_xx = MathTex(
            r"\cos2x = \cos(A+B)"
        )
        cos_xx.move_to(ORIGIN + UP)

        cos_ab = MathTex (
            r"\cos(A+B) =(\cos A \cdot \cos B) - (\sin A \cdot \sin B)"
        )
        cos_ab.move_to(ORIGIN)

        let_AB = Tex("Let A = B")
        let_AB.move_to(ORIGIN + DOWN)

        ab_simple = MathTex(
            r"\cos(A+A) = \cos^2A - \sin^2A"
        )
        ab_simple.move_to(ORIGIN + DOWN + DOWN)

        ab_finalize = MathTex(
            r"= 1-2\sin^2x"
        )
        ab_finalize.move_to(ORIGIN + DOWN + DOWN + DOWN + RIGHT)

        self.play(
            Write(cos_xx)
        )
        self.wait(0.5)
        self.play(
            Write(cos_ab),
        )
        self.wait(0.5)
        self.play(
            Write(let_AB)
        )
        self.wait(0.5)
        self.play(
            Write(ab_simple)
        )
        self.wait(0.5)
        self.play(
            Write(ab_finalize)
        )

        arrow = Arrow(2*UP, 2*DOWN)
        VGroup(arrow).set_x(0).arrange(buff=2)
        arrow.move_to(ORIGIN + RIGHT + RIGHT + RIGHT + RIGHT + RIGHT + RIGHT)
        self.play(Write(arrow))

        self.wait(15)

        self.play(
            FadeOut(text, transform_text, trig_cos2, cos_xx, cos_ab, let_AB, ab_simple, ab_finalize, arrow, box, transform_formula)
        )

        self.wait(1)
        #moving to the explanation of example

        #What is proof in Math?
        proof = Tex("What is proof?", font_size = 48)
        self.play(Write(proof))
        self.wait(3)

        self.play(
            Transform(proof, proof.copy().shift(UP).shift(UP))
        )

        proof_exp = Paragraph("In trigonometry, a proof is a way to show that ", "two trigonometric expressions are equivalent, regardless of the angle. ","This process is called validating or proving trigonometric identities.", font_size=28)
        self.play(Write(proof_exp))

        self.wait(8)
        self.play(
            FadeOut(proof, proof_exp)
        )

        #starting with Sin and Cos graph identity

        ax = Axes()
        sine = ax.plot(np.sin, color = RED)
        cosine = ax.plot(np.cos, color = BLUE)
        self.play(
            FadeIn(ax, sine, cosine)
            )

        red_square = Square(fill_opacity = 1, side_length=0.5, fill_color = RED_C).to_corner(UL)
        blue_square = Square(fill_opacity=1, side_length=0.5, fill_color=BLUE_C).to_corner(UL - DOWN)

        self.play(DrawBorderThenFill(red_square))
        self.play(DrawBorderThenFill(blue_square))
        text_sin = MathTex(r"\sin(x)")
        text_cos = MathTex(r"\cos(x)")
        text_sin.next_to(Square(fill_opacity=1, side_length=0.5, fill_color=RED_C).to_corner(UL))
        text_cos.next_to(Square(fill_opacity=1, side_length=0.5, fill_color=BLUE_C).to_corner(UL - DOWN))
        # Correct usage of next_to: Multiply RIGHT by a scala

        self.play(Write(text_sin))
        self.wait(0.5)

        self.play(Write(text_cos))
        self.wait(0.5)

        self.wait(8)
        self.play(FadeOut(sine, cosine, text_sin, text_cos, ax, red_square, blue_square))
        self.wait(2)

        prob_cos = Tex(r"Prove that $\cos\left(x - \frac{\pi}{2}\right)$ is the same as $\sin x$")
        self.play(Write(prob_cos))
        self.wait(2)

        self.play(
            Transform(prob_cos, prob_cos.copy().to_corner(UP))
        )
        self.wait(10)

        step1 = Tex(r"1. Make balance equation $\cos\left(x - \frac{\pi}{2}\right) = \sin x$")
        step2 = Tex("2. Identify which side is easier to change form, or simplify.")
        step3 = Tex("3. Formulate and make it equal to the other side.")

        steps = VGroup(step1, step2, step3).arrange(DOWN, aligned_edge=LEFT)
        steps.move_to(ORIGIN)
        steps.next_to(prob_cos, DOWN, buff=0.5)

        self.play(
            Write(steps)
        )

        self.wait(3)

        self.play(Circumscribe(step1, Rectangle, time_width=4))

        self.play(
            FadeOut(step2, step3)
        )

        step1_exp = MathTex(r"\cos\left(x-\frac{\pi}{2}\right) = \sin x")
        step1_exp.move_to(ORIGIN)

        self.play(
            Write(step1_exp)
        )

        self.wait(6)

        self.play(
            FadeOut(step1, step1_exp),
        )

        self.wait(1)

        self.play(
            FadeIn(steps),
        )

        self.wait(3)

        self.play(
            Circumscribe(step2, Rectangle, time_width=4)
        )

        self.play(
            FadeOut(step1, step3),
            Transform(step2, step2.copy().shift(UP))
        )

        self.wait(3)

        step2_exp = MathTex(r"\cos\left(x-\frac{\pi}{2}\right)", color=BLUE)
        step2_exp.move_to(ORIGIN)
        self.play(Write(step2_exp))
        self.wait(2)

        step2_exp2 = Tex("Left side is easier to change form", color=BLUE)
        step2_exp2.next_to(step2_exp, DOWN)

        self.play(Write(step2_exp2))
        self.wait(2)

        step2_exp3 = MathTex(r"\cos\left(x-\frac{\pi}{2}\right) = \cos(A-B)", color=WHITE)
        step2_exp3.move_to(ORIGIN)

        self.play(
            Transform(step2_exp, step2_exp3),
            FadeOut(step2_exp2)
        )
        self.wait(2)

        step2_exp4 = MathTex(r"\cos(A-B) = \cos A \cos B + \sin A \sin B", color=BLUE)
        step2_exp4.next_to(step2_exp3, DOWN)

        self.play(Write(step2_exp4))
        self.wait(2)

        step2_exp5 = MathTex(r"A = x, B = \frac{\pi}{2}", color=BLUE)
        step2_exp5.next_to(step2_exp4, DOWN)

        self.play(Write(step2_exp5))
        self.wait(2)

        step2_exp6 = MathTex(r"\cos x \cos \frac{\pi}{2} + \sin x \sin \frac{\pi}{2}", color=WHITE)
        step2_exp6.move_to(ORIGIN)

        self.play(
            FadeOut(step2_exp, step2_exp4, step2_exp5),
            Write(step2_exp6)
        )
        self.wait(2)

        step2_exp7 = MathTex(r"\cos \frac{\pi}{2} = 0, \sin \frac{\pi}{2} = 1", color=BLUE)
        step2_exp7.next_to(step2_exp6, DOWN)

        self.play(Write(step2_exp7))
        self.wait(2)

step2_exp8 = MathTex(r"\cos x (0) + \sin x (1) = \sin x", color=WHITE)
        step2_exp8.move_to(ORIGIN)

        self.play(
            FadeOut(step2_exp6, step2_exp7),
            Write(step2_exp8)
        )
        self.wait(2)

        self.play(FadeOut(step2_exp8, step2))

        self.wait(15)

Instruction: Move the proved sinx to center of the screen and fade out rest of equation

import pytest
import re
import inspect
from typing import List
import ast

def get_source_code(impl_name, module) -> str:
    """Get the source code of the implementation module"""
    try:
        return inspect.getsource(module)
    except Exception:
        return ""

import re
from typing import List

def test_moves_sinx_equation_to_center(implementation):
    """Test if sinx (step2_exp8) is moved to the center of the screen"""
    impl_name, module = implementation
    code = get_source_code(impl_name, module)

    # Look for .move_to(ORIGIN) or .animate.move_to(ORIGIN) applied to sinx object
    moved = re.search(r'step2_exp8(\.animate)?\.move_to\s*\(\s*ORIGIN\s*\)', code)
    assert moved, f"{impl_name} does not move sinx (step2_exp8) to center using move_to(ORIGIN)"

def test_fades_out_other_equations(implementation):
    """Test if other equations (e.g. step2_exp6, step2_exp7) are faded out"""
    impl_name, module = implementation
    code = get_source_code(impl_name, module)

    # Look for FadeOut involving other step2 expressions
    fadeout_other = re.search(r'FadeOut\s*\(\s*step2_exp6\s*,\s*step2_exp7\s*\)', code) or \
                    re.search(r'FadeOut\s*\(\s*step2_exp\d+', code)
    assert fadeout_other, f"{impl_name} does not fade out other equations like step2_exp6, step2_exp7"

pytest
pytest-mock
manim
numpy
pyglm
pydub

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