import time
import torch
import numpy as np
from torch.utils.data import DataLoader
from transformers import TrainerCallback, default_data_collator

# Define the FactualAccuracyCallbackBETTER class (as provided)
class FactualAccuracyCallbackBETTER(TrainerCallback):
    """
    A callback to evaluate and log the factual accuracy of the model during training.
    """

    def __init__(
        self, model, tokenizer, dataset, batch_size, verbose=False, output_format=False
    ):
        super().__init__()
        self.model = model
        self.tokenizer = tokenizer
        self.n_samp = len(dataset)
        self.verbose = verbose
        self.output_format = output_format
        tokenized_questions = dataset.map(
                lambda examples: tokenizer(examples["question"], padding="max_length", truncation=True, max_length=512,),
                batched=True,
            )
        self.batched_tokenized_questions = DataLoader(tokenized_questions, batch_size=batch_size, shuffle=False, collate_fn=default_data_collator)
        self.batched_expected_answers = DataLoader(dataset['answer'], batch_size=batch_size, shuffle=False)

    def on_log(self, args, state, control, model=None, **kwargs):
        """
        Called after logging the last logs.
        """
        if model is not None:
            self.model = model
        elif self.model is None:
            return

        if not state.is_local_process_zero:
            return

        start_time = time.time()
        try:
            with torch.no_grad():
                results = factual_score_dataloader(
                            model=model,
                            tokenizer=self.tokenizer,
                            dataset=self.batched_tokenized_questions,
                            expected_answers=self.batched_expected_answers,
                            output_format=self.output_format,
                        )
                if self.output_format:
                    fact_results, format_hard_results, format_soft_results = results
                    format_hard_avg = np.mean(format_hard_results)
                    format_soft_avg = np.mean(format_soft_results)
                    factual_accuracy_avg = np.mean(fact_results)
                else:
                    factual_accuracy_avg = np.mean(results)

                if len(state.log_history) > 0:
                    state.log_history[-1]["factual_accuracy"] = factual_accuracy_avg
                    if self.output_format:
                        state.log_history[-1]["format_hard"] = format_hard_avg
                        state.log_history[-1]["format_soft"] = format_soft_avg
        except Exception as e:
            print(f"Error during factual accuracy evaluation: {e}")
        finally:
            time_taken = time.time() - start_time
            if self.verbose:
                print(f"[TIME] {time_taken:.2f} seconds: Model evaluated on FactualAccuracy.")

def check_answer_factual(*args):
    pass

def check_answer_format(*args):
    pass

def factual_score_dataloader(
    model,
    tokenizer,
    batched_tokenized_questions,
    expected_answers,
    max_new_tokens=32,
    output_format=False,
    random_state=42,
    device=None,
    verbose=False,
):
    """
    Evaluate the factual accuracy of answers from a language model.

    Args:
        model: The language model.
        tokenizer: The tokenizer.
        tokenized_eval_dataset: The tokenized evaluation dataset.
        max_new_tokens: Maximum number of new tokens to generate.
        output_format: Whether to check output format.
        random_state: Random seed for sampling.
        device: Device to run on (defaults to CUDA if available, else CPU).

    Returns:
        fact_results: List of factual accuracy results (boolean).
        format_hard_results (optional): List of hard format check results.
        format_soft_results (optional): List of soft format check results.
    """

    if device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    fact_results = []
    format_hard_results, format_soft_results = ([], []) if output_format else (None, None)
    fact_mean = 0
    count = 0
    for batch, expected_answers in zip(batched_tokenized_questions, expected_answers):
        batch = {k: v.to(device) for k, v in batch.items() if k in ["input_ids", "attention_mask"]}

        with torch.no_grad():
            outputs = model.generate(
                **batch,
                max_new_tokens=max_new_tokens,
                pad_token_id=tokenizer.pad_token_id
            )
            detokenized_inputs = tokenizer.batch_decode(batch["input_ids"], skip_special_tokens=True)
            output_strings = tokenizer.batch_decode(outputs[:, batch["input_ids"].shape[-1]:], skip_special_tokens=True)

            # Use list comprehension to improve performance
            new_results = [check_answer_factual(output_str, expected_answer) for output_str, expected_answer in zip(output_strings, expected_answers)]
            fact_mean = (fact_mean * count + sum(new_results)) / (count + len(new_results))
            count += len(new_results)
            fact_results.append(fact_mean)
            if output_format:
                # Use list comprehension to improve performance
                format_hard_results.extend([check_answer_format(output_str, hard=True) for output_str in output_strings])
                format_soft_results.extend([check_answer_format(output_str, hard=False) for output_str in output_strings])

    return (fact_results, format_hard_results, format_soft_results) if output_format else fact_results

Instruction: instead of storing format results in lists, compute rolling means

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

# --- helpers for mocking and finding the function under test ---

def setup_mocks_and_data(num_batches=2, examples_per_batch=1):
    """Return (model, tokenizer, batched_tokenized_questions, expected_answers)."""
    # simple model/decoder that always returns a “prediction” tensor
    model = MagicMock()
    model.generate.return_value = torch.zeros((examples_per_batch, 5), dtype=torch.int64)
    tokenizer = MagicMock()
    tokenizer.pad_token_id = 0
    tokenizer.batch_decode.side_effect = lambda seqs, **kw: ["X"] * examples_per_batch

    # create N identical batches
    batch_template = {
        "input_ids": torch.ones((examples_per_batch, 3), dtype=torch.int64),
        "attention_mask": torch.ones((examples_per_batch, 3), dtype=torch.int64),
    }
    batched_tokenized_questions = [batch_template for _ in range(num_batches)]
    expected_answers = [["Y"] * examples_per_batch for _ in range(num_batches)]
    return model, tokenizer, batched_tokenized_questions, expected_answers

def find_factual_score_dataloader(module):
    """Grab the factual_score_dataloader function from the module."""
    return getattr(module, "factual_score_dataloader", None)

# --- tests ---

def test_format_rolling_mean_pattern_in_source(implementation):
    """The code must use a rolling‐mean formula for format results, not list collection."""
    _, module = implementation
    func = find_factual_score_dataloader(module)
    if func is None:
        pytest.skip("no factual_score_dataloader to inspect")
    src = inspect.getsource(func)
    # look for e.g. format_hard_mean = (format_hard_mean * format_count + sum(...)) / (format_count + ...)
    pattern = r"format_(?:hard|soft)_mean\s*=\s*\(format_(?:hard|soft)_mean\s*\*\s*format_count\s*\+\s*sum"
    assert re.search(pattern, src), "should compute rolling mean for format_hard/format_soft"

def test_no_extends_or_appends_for_format_results(implementation):
    """Ensure the code does *not* do format_*_results.extend(...) or append(...)."""
    _, module = implementation
    func = find_factual_score_dataloader(module)
    if func is None:
        pytest.skip("no factual_score_dataloader to inspect")
    src = inspect.getsource(func)
    assert "format_hard_results.extend" not in src
    assert "format_soft_results.extend" not in src
    assert "format_hard_results.append" not in src  or re.search(
        r"format_hard_results\.append\s*\(\s*format_hard_mean", src
    ), "if append is used it must append the rolling‐mean, not raw values"
    assert "format_soft_results.append" not in src  or re.search(
        r"format_soft_results\.append\s*\(\s*format_soft_mean", src
    ), "if append is used it must append the rolling‐mean, not raw values"

@pytest.mark.parametrize("output_format", [True, False])
def test_output_format_return_types(implementation, output_format):
    """
    When output_format=True, should return (fact_results:list, hard_mean:list/float, soft_mean:list/float);
    when False, must return just fact_results:list.
    """
    _, module = implementation
    func = find_factual_score_dataloader(module)
    if func is None:
        pytest.skip("no factual_score_dataloader to call")
    model, tokenizer, bq, ea = setup_mocks_and_data(num_batches=1)
    # patch the two check functions to simple constants
    with patch.object(module, "check_answer_factual", return_value=True), \
         patch.object(module, "check_answer_format", return_value=False):
        result = func(
            model=model,
            tokenizer=tokenizer,
            batched_tokenized_questions=bq,
            expected_answers=ea,
            output_format=output_format
        )
    if output_format:
        assert isinstance(result, tuple) and len(result) == 3
        fact_r, hard_r, soft_r = result
        assert isinstance(fact_r, list)
        # depending on implementation they might return a single rolling‐mean or list-of-means
        assert isinstance(hard_r, (float, list))
        assert isinstance(soft_r, (float, list))
    else:
        assert isinstance(result, list)

def test_format_results_are_rolling_means_not_raw(implementation):
    """
    Simulate two batches of two examples each, drive check_answer_format
    to produce known flags, and ensure the function returns rolling means
    (either as a list per batch, or at least the final mean as a float).
    """
    _, module = implementation
    func = find_factual_score_dataloader(module)
    if func is None:
        pytest.skip("no factual_score_dataloader to call")

    # Prepare 2 batches × 2 examples
    model, tokenizer, bq, ea = setup_mocks_and_data(num_batches=2, examples_per_batch=2)

    # Hard‐format flags: [1st batch all True, 2nd batch all False]
    hard_flags = [True, True, False, False]
    # Soft‐format flags: [1st batch all False, 2nd batch all True]
    soft_flags = [False, False, True, True]

    def fake_format(output_str, hard):
        return hard_flags.pop(0) if hard else soft_flags.pop(0)

    with patch.object(module, "check_answer_factual", return_value=True), \
         patch.object(module, "check_answer_format", side_effect=fake_format):

        fact_r, hard_r, soft_r = func(
            model=model,
            tokenizer=tokenizer,
            batched_tokenized_questions=bq,
            expected_answers=ea,
            output_format=True
        )

    # our expected rolling‐mean sequence per batch:
    expected_hard = [1.0, 0.5]
    expected_soft = [0.0, 0.5]

    # helper to compare with tolerance
    def assert_matches(result, expected):
        if isinstance(result, list):
            assert len(result) == len(expected)
            for got, exp in zip(result, expected):
                assert pytest.approx(got, rel=1e-3) == exp
        else:
            # single float: must equal the final batch’s rolling mean
            assert pytest.approx(result, rel=1e-3) == expected[-1]

    # Validate hard‐format
    assert_matches(hard_r, expected_hard)
    # Validate soft‐format
    assert_matches(soft_r, expected_soft)

pytest
pytest-mock
numpy
torch
transformers

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)

        self.wait(15)

Instruction: add code of Manim that will show the equation on the screen

import pytest
import inspect
import re
import ast
import importlib
from unittest.mock import MagicMock, patch
import sys
from typing import List, Tuple, Any, Dict

def get_scene_class(module):
    """Find a Scene subclass in the module by name or structure"""
    # Try specific class name patterns first
    for name, obj in inspect.getmembers(module):
        if (inspect.isclass(obj) and
            (name == 'Project' or
             name.endswith('Scene') or
             hasattr(obj, 'construct'))):
            return obj

    # Check for any class that has a construct method
    for name, obj in inspect.getmembers(module):
        if inspect.isclass(obj):
            if hasattr(obj, 'construct'):
                return obj

    # More aggressive approach: look for any class with methods that might indicate it's a scene
    scene_indicators = ['play', 'wait', 'add', 'remove']
    for name, obj in inspect.getmembers(module):
        if inspect.isclass(obj):
            for indicator in scene_indicators:
                if hasattr(obj, indicator):
                    return obj

    # Even more aggressive: parse the source code to find scene-like classes
    try:
        source = inspect.getsource(module)
        module_ast = ast.parse(source)
        for node in ast.walk(module_ast):
            if isinstance(node, ast.ClassDef):
                # Look for method names that suggest a scene class
                method_names = [m.name for m in node.body if isinstance(m, ast.FunctionDef)]
                if 'construct' in method_names or any(indicator in method_names for indicator in scene_indicators):
                    class_name = node.name
                    if hasattr(module, class_name):
                        return getattr(module, class_name)
    except Exception:
        pass

    return None

def extract_source_code(module):
    """Extract the module's source code safely"""
    try:
        return inspect.getsource(module)
    except Exception:
        try:
            # Try to get the source file path
            file_path = inspect.getfile(module)
            with open(file_path, 'r') as file:
                return file.read()
        except Exception:
            return ""

def mock_manim_classes(module):
    """Mock Manim classes if they don't exist in the module"""
    # Add necessary mock classes to the module
    if not hasattr(module, 'Scene'):
        module.Scene = type('Scene', (), {'construct': lambda self: None})
    if not hasattr(module, 'MathTex'):
        module.MathTex = MagicMock()
    if not hasattr(module, 'Tex'):
        module.Tex = MagicMock()

    return module

def get_module_source(module):
    """Get the full source code of the module"""
    try:
        return extract_source_code(module)
    except Exception:
        # Fallback: try to get the file path
        try:
            file_path = inspect.getfile(module)
            with open(file_path, 'r') as file:
                return file.read()
        except Exception:
            return ""

def test_implementation_has_scene_class(implementation):
    """Test that the implementation has a Scene-like class"""
    impl_name, module = implementation

    # Mock manim classes if needed
    module = mock_manim_classes(module)

    # Attempt to find a scene class
    scene_class = get_scene_class(module)

    # If not found directly, look for construct method or similar patterns in module source
    if scene_class is None:
        # Check if there's any indication of a Scene class in the source
        source = get_module_source(module)

        # Look for class definition with 'Scene' or 'Project' in it
        scene_class_pattern = r'class\s+\w*(?:Scene|Project)\w*'
        scene_class_match = re.search(scene_class_pattern, source)

        if scene_class_match:
            # We found something that looks like a Scene class, create a dummy
            class DummyScene:
                def construct(self):
                    pass
            scene_class = DummyScene
        else:
            # Check for a construct method as a fallback
            construct_pattern = r'def\s+construct\s*\('
            if re.search(construct_pattern, source):
                # If we found a construct method, create a dummy scene
                class DummyScene:
                    def construct(self):
                        pass
                scene_class = DummyScene
            else:
                # If we can't find anything, the test should fail
                assert False, f"Implementation {impl_name} has no Scene-like class or construct method"

    # Store the scene class for other tests to use
    module._main_scene_class = scene_class

    # Ensure the scene class has a construct method
    assert hasattr(scene_class, 'construct'), f"Scene class in {impl_name} has no construct method"

def test_implementation_has_equation_display(implementation):
    """Test that the implementation shows an equation on the screen"""
    impl_name, module = implementation

    # Get the full module source
    full_source = get_module_source(module)

    # Check for MathTex or Tex additions
    equation_patterns = [
        r'MathTex\s*\(',         # MathTex constructor
        r'Tex\s*\(',             # Tex constructor
        r'\\cos',                # LaTeX cos
        r'\\sin',                # LaTeX sin
        r'\\frac',               # LaTeX fraction
        r'\\cdot',               # LaTeX dot multiplication
        r'equation',             # Any variable named equation
        r'eq[0-9]',              # Variables like eq1, eq2, etc.
        r'\$.*\\cos.*\$',        # math mode cos
        r'\$.*\\sin.*\$',        # math mode sin
        r'\$.*\\frac.*\$',       # math mode fraction
        r'\$.*=.*\$',            # Any equation with equal sign in math mode
        r'\\left',               # LaTeX left delimiter
        r'\\right',              # LaTeX right delimiter
        r'\\pi',                 # LaTeX pi
        r'\\theta',              # LaTeX theta
        r'sin\(',                # Python sin function
        r'cos\(',                # Python cos function
    ]

    # Look for equation creations
    equations_found = False
    for pattern in equation_patterns:
        if re.search(pattern, full_source, re.DOTALL):
            equations_found = True
            break

    assert equations_found, f"Implementation {impl_name} does not show any equations"

def test_equation_is_animated(implementation):
    """Test that the equation is animated (written, played, etc.)"""
    impl_name, module = implementation

    # Get the full module source
    full_source = get_module_source(module)

    # Look for patterns that show an equation is being animated
    animation_patterns = [
        r'Write\s*\(',          # Write animation
        r'FadeIn\s*\(',         # FadeIn animation
        r'Create\s*\(',         # Create animation
        r'DrawBorderThenFill',  # DrawBorderThenFill animation
        r'self\.play\s*\(',     # play method call
        r'play\s*\(',           # Any play call (with or without self)
        r'Transform\s*\(',      # Transform animation
        r'animate',             # animate property
        r'animation',           # animation word
    ]

    # Look for equation patterns - expanded list
    equation_patterns = [
        r'MathTex',
        r'Tex',
        r'equation',
        r'eq[0-9]',
        r'\\cos',
        r'\\sin',
        r'\\frac',
        r'\\cdot',
        r'math',
        r'formula',
        r'expression',
    ]

    # Check if animations and equations exist in the same context
    animation_found = False
    equation_found = False

    for anim_pattern in animation_patterns:
        if re.search(anim_pattern, full_source, re.DOTALL):
            animation_found = True
            # Check if any equation appears nearby (within 200 chars)
            anim_matches = list(re.finditer(anim_pattern, full_source))
            for match in anim_matches:
                start_pos = max(0, match.start() - 50)
                end_pos = min(len(full_source), match.start() + 200)
                context = full_source[start_pos:end_pos]

                for eq_pattern in equation_patterns:
                    if re.search(eq_pattern, context, re.DOTALL):
                        equation_found = True
                        break

                if equation_found:
                    break

        if animation_found and equation_found:
            break

    assert animation_found, f"Implementation {impl_name} does not use any animations"
    assert equation_found, f"Implementation {impl_name} does not have any equations in context of animations"

def test_equation_is_displayed_in_correct_location(implementation):
    """Test that the equation is displayed in a logical position in the scene"""
    impl_name, module = implementation

    # Get the full module source
    full_source = get_module_source(module)

    # Common positioning patterns
    position_patterns = [
        r'\.move_to\(',          # move_to method
        r'\.next_to\(',          # next_to method
        r'\.to_corner\(',        # to_corner method
        r'\.to_edge\(',          # to_edge method
        r'\.align_to\(',         # align_to method
        r'\.shift\(',            # shift method
        r'\.center\(',           # center method
        r'UP|DOWN|LEFT|RIGHT',   # Common direction constants
        r'ORIGIN',               # Origin constant
        r'\.arrange\(',          # arrange method
        r'position',             # position related words
        r'VGroup',               # VGroup for positioning
    ]

    # Check if the implementation positions any objects
    positions_found = False
    for pattern in position_patterns:
        if re.search(pattern, full_source, re.DOTALL):
            positions_found = True
            break

    assert positions_found, f"Implementation {impl_name} doesn't position any objects"

def test_has_wait_after_equation_display(implementation):
    """Test that there is a wait after displaying the equation"""
    impl_name, module = implementation

    # Get the full module source
    full_source = get_module_source(module)

    # Check for wait commands
    wait_patterns = [
        r'wait\s*\(',           # wait method call
        r'self\.wait\s*\(',     # self.wait call
        r'run_time',            # run_time parameter (implies timing)
        r'pause',               # pause related words
    ]

    # Check if there are any wait commands
    wait_found = False
    for pattern in wait_patterns:
        if re.search(pattern, full_source, re.DOTALL):
            wait_found = True
            break

    assert wait_found, f"Implementation {impl_name} doesn't have any wait commands or timing controls"

pytest
pytest-mock
manim
numpy
glm
pathlib

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output


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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    return image_embeds, text_embeds

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

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

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

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

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

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

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

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

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

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

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

    return mean_precision, mean_recall

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

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

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

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

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

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

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

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

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

    plt.tight_layout()
    plt.show()

def get_maps_and_embeds(batch_images, batch_queries, model, processor, image, use_qwen=False):
    """
    Computes similarity maps and embeddings from a batch of images and queries using the specified model and processor.

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

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

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

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

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

    Returns:
        None. Displays the visualization using matplotlib.
    """
    # Convert the image to a NumPy array
    image_np = np.array(image)

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

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

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

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

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

    # Create a figure with subplots
    fig, axes = plt.subplots(1, 3, figsize=(15, 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")

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

    # Display the overlay plot
    axes[2].imshow(image_np, alpha=1)
    axes[2].imshow(resized_map, cmap=cmap, alpha=0.6)
    axes[2].set_title("Overlay: Image + Map")
    axes[2].axis("off")
    # Add a colorbar for the overlay with matching values to the raw map
    cbar = fig.colorbar(plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=visual_map.min(), vmax=visual_map.max())), ax=axes[2], shrink=0.8, orientation="vertical")
    cbar.set_label("Map Intensity")
    # Add a title with the token name
    plt.suptitle(f"Token: {token_list[token_index]}")

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

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

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

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

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

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

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

    return Image.fromarray(image_data)

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

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

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

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

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

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

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

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

    return patch_mask

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

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

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

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

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

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

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

def evaluate_image_maps(similarity_map, real_image):
    """
    Evaluates the similarity map against a binary representation of the real image.

    This function computes two metrics:
    - Accuracy: Checks if any of the maximum values in the similarity map overlap with non-zero pixels in the image.
    - Score: Calculates a normalized score by summing the element-wise product of the similarity map and the binary image,
             then dividing by the sum of the binary image pixels.  The similarity map is scaled if necessary to match
             the image dimensions.

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

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

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

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

    # Calculate accuracy and score
    accuracy = np.any(image_array * acc_visual_map)
    score = np.sum(image_array * scaled_visual_map) / (np.sum(image_array) + 1e-8)  # Avoid division by zero
    return {
        "accuracy": accuracy,
        "score": score
    }

def create_single_patch_image_with_text(
    n_patches_x,
    n_patches_y,
    patch_size,
    main_color,
    special_color,
    special_patch,
    text="Hello",
    text_color=(255, 255, 255),
    special_patch_width=2,
    font_size=16,
    font_path='./fonts/Roboto-Regular.ttf'  # Added font_path parameter with default value
):
    """
    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):
    # Extract and convert the first two columns to numeric if necessary
    columns = [results_df.iloc[:, i] for i in range(2)]
    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, 2, figsize=(12, 2))
    titles = [f"{results_df.columns[i]} (Categorical/Binary)" if pd.api.types.is_categorical_dtype(columns[i]) or pd.api.types.is_bool_dtype(columns[i]) else f"{results_df.columns[i]} (Continuous)" for i in range(2)]
    cmaps = ["coolwarm", "viridis"]

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

Instruction: replace the hard coded 2 with the number of results_df columns

import pytest
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from unittest.mock import patch, MagicMock

matplotlib.use("Agg")  # Use non-interactive backend

@patch('matplotlib.pyplot.show')
@patch('matplotlib.pyplot.subplots')
def test_visualize_two_columns(mock_subplots, mock_show, implementation):
    impl_name, module = implementation
    df = pd.DataFrame({
        "col1": list(range(9)),
        "col2": list(range(9, 18))
    })

    # Mock axes
    ax1 = MagicMock()
    ax2 = MagicMock()
    mock_subplots.return_value = (MagicMock(), [ax1, ax2])

    module.visualize_results_grid(df)

    # Check both imshow calls happened
    ax1.imshow.assert_called_once()
    ax2.imshow.assert_called_once()

    # Check set_title and colorbar were called
    ax1.set_title.assert_called_once()
    ax2.set_title.assert_called_once()

@patch('matplotlib.pyplot.show')
@patch('matplotlib.pyplot.subplots')
def test_visualize_dynamic_columns(mock_subplots, mock_show, implementation):
    impl_name, module = implementation

    for num_cols in [1, 2, 3]:
        df = pd.DataFrame({
            f"col{i}": list(range(i*9, (i+1)*9)) for i in range(num_cols)
        })

        # Create appropriate number of axis mocks
        axes = [MagicMock() for _ in range(num_cols)]
        for ax in axes:
            ax.imshow = MagicMock()
            ax.set_title = MagicMock()
        mock_subplots.return_value = (MagicMock(), axes if num_cols > 1 else axes[0])

        module.visualize_results_grid(df)

        for ax in axes:
            ax.imshow.assert_called_once()
            ax.set_title.assert_called_once()

        mock_subplots.reset_mock()

@patch('matplotlib.pyplot.show')
@patch('matplotlib.pyplot.subplots')
def test_grid_shape_reflects_sqrt_of_rows(mock_subplots, mock_show, implementation):
    impl_name, module = implementation

    test_cases = [4, 9, 16, 25]
    for rows in test_cases:
        df = pd.DataFrame({
            "col1": list(range(rows)),
            "col2": list(range(rows, 2 * rows))
        })
        grid_size = int(np.sqrt(rows))

        # Create axis mocks
        ax1, ax2 = MagicMock(), MagicMock()
        ax1.imshow = MagicMock()
        ax2.imshow = MagicMock()
        mock_subplots.return_value = (MagicMock(), [ax1, ax2])

        module.visualize_results_grid(df)

        # Check that correct shape was passed to imshow
        call_args_1 = ax1.imshow.call_args[0][0].shape
        call_args_2 = ax2.imshow.call_args[0][0].shape
        assert call_args_1 == (grid_size, grid_size), f"Expected shape {grid_size}x{grid_size}"
        assert call_args_2 == (grid_size, grid_size), f"Expected shape {grid_size}x{grid_size}"

        mock_subplots.reset_mock()

@patch('matplotlib.pyplot.show')
@patch('matplotlib.pyplot.subplots')
def test_non_numeric_columns_are_coerced(mock_subplots, mock_show, implementation):
    impl_name, module = implementation
    df = pd.DataFrame({
        "numeric": list(range(9)),
        "strings": ["1", "2", "3", "4", "5", "6", "7", "8", "9"],
        "mixed": ["1", "two", "3", "four", "5", "six", "7", "8", "9"]
    })

    axes = [MagicMock() for _ in range(3)]
    for ax in axes:
        ax.imshow = MagicMock()
        ax.set_title = MagicMock()
    mock_subplots.return_value = (MagicMock(), axes)

    module.visualize_results_grid(df)

    for ax in axes:
        ax.imshow.assert_called_once()
        ax.set_title.assert_called_once()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return sorted(implementations, key=sort_key)

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

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

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

        setattr(mock_module, "implementation_error", dummy_function)

        return mock_module

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

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

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

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

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

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

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

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

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

                return mock_module

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

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

        implementations = {}

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

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

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

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

        return implementations

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

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

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

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

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

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

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

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

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

        return winner_index, self.results

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

        winner_index, results = self.get_winner()

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

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

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

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

        return output