"""MLE-bench: Evaluating Machine Learning Agents on Machine Learning Engineering

Jun Shern Chan, Neil Chowdhury, Oliver Jaffe, James Aung, Dane Sherburn, Evan Mays, Giulio Starace, Kevin Liu, Leon Maksin, Tejal Patwardhan, Lilian Weng, Aleksander Mądry
https://arxiv.org/abs/2410.07095
"""

import json
import logging
from importlib.util import find_spec
from pathlib import Path
from textwrap import dedent

from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, Sample
from inspect_ai.scorer import Score, Scorer, Target, accuracy, scorer, stderr
from inspect_ai.solver import Solver, TaskState, basic_agent, system_message
from inspect_ai.tool import bash
from inspect_ai.util import SandboxEnvironmentSpec, sandbox
from platformdirs import user_cache_dir

FILE_DIR = Path(__file__).parent
DEFAULT_DATA_DIR = Path(user_cache_dir("inspect_mlebench_eval")) / "data"
COMPOSE_FILES_DIR = Path(user_cache_dir("inspect_mlebench_eval")) / "compose_files"
DEFAULT_IMAGE_NAME = "mlebench-inspect-env"


logger = logging.getLogger(__name__)


with open(Path(__file__).parent / "instructions.txt", "r") as f:
    INSTRUCTIONS = f.read()


@task
def mle_bench(
    split: Path | str = "spaceship-titanic.txt",
    solver: Solver | None = None,
    max_messages: int = 30,
    scorer: Scorer | list[Scorer] | None = None,
    epochs: int = 1,
    docker_image_name: str | None = None,
    force_rebuild: bool = False,
    data_dir: Path | str | None = None,
) -> Task:
    """Returns a Task, representing an evaluation on MLE-bench.

    Args.
        split : str | Path
            The split to use. This should be the name of a .txt file within the splits directory, or a path to a custom split file. low.txt corresponts to MLE-bench-LITE, while all.txt is the full benchmark. The default is a split continaining only the spaceship-titanic competition.
        solver : Solver
            The solver to use when creating the task. If None, uses the default solver.
        max_messages : int
            The maximum number of messages to allow for each sample. Only applies when using the default solver.
        scorer: Scorer | list[Scorer] | None
            The scorer to use when evaluating mle_bench. If None, uses the default scorer which returns boolean values for the criteria of "valid_submission", "above_median" and "any_medal".
        epochs: int
            Number of times to repeat each sample.
        docker_image_name: str | None
            The name of the docker image to be used for each sandbox. If such an image does not already exist, it will be built from inspect_evals/mle_bench/Dockerfile.
        force_rebuild: bool
            Whether to rebuild when an image already exists with the name docker_image_name.
        data_dir : Path | str | None
            The directory that will be used for per competition data. If it does not exist, the data will be prepared here. If None, a directory in the user's cache will be used by default.

    """
    assert find_spec("mlebench"), (
        "To run MLE-bench, please install the optional MLE-bench dependency, by running `pip install inspect-evals[mle_bench]`"
    )
    from mlebench.data import (  # type: ignore
        download_and_prepare_dataset,
        is_dataset_prepared,
    )
    from mlebench.registry import registry  # type: ignore

    # If the split is a path, we assume it is a path to a custom split file and leave it as is.
    if Path(split).parent == Path("."):
        split = FILE_DIR / "splits" / split

    with open(split) as f:
        competition_ids = f.read().splitlines()

    new_registry = registry.set_data_dir(data_dir or DEFAULT_DATA_DIR)
    samples = []

    for competition_id in competition_ids:
        competition = new_registry.get_competition(competition_id)
        if not is_dataset_prepared(competition):
            print(f"Preparing competition {competition_id}")
            download_and_prepare_dataset(competition)

        samples.append(
            Sample(
                input=INSTRUCTIONS,
                id=competition_id,
                sandbox=SandboxEnvironmentSpec(
                    type="docker",
                    config=get_compose_file(
                        competition_id=competition_id,
                        data_dir=data_dir or DEFAULT_DATA_DIR,
                        docker_image_name=docker_image_name or DEFAULT_IMAGE_NAME,
                        force_rebuild=force_rebuild,
                    ),
                ),
            )
        )

    return Task(
        dataset=MemoryDataset(samples=samples),
        solver=solver or default_solver(max_messages),
        epochs=epochs,
        scorer=scorer or mle_bench_scorer(),
    )


@task
def mle_bench_full(
    solver: Solver | None = None,
    max_messages: int = 30,
    scorer: Scorer | list[Scorer] | None = None,
    epochs: int = 1,
    docker_image_name: str | None = None,
    force_rebuild: bool = False,
    data_dir: Path | str | None = None,
) -> Task:
    return mle_bench(
        split="all.txt",
        solver=solver,
        max_messages=max_messages,
        scorer=scorer,
        epochs=epochs,
        docker_image_name=docker_image_name,
        force_rebuild=force_rebuild,
        data_dir=data_dir,
    )


@task
def mle_bench_lite(
    solver: Solver | None = None,
    max_messages: int = 30,
    scorer: Scorer | list[Scorer] | None = None,
    epochs: int = 1,
    docker_image_name: str | None = None,
    force_rebuild: bool = False,
    data_dir: Path | str | None = None,
) -> Task:
    return mle_bench(
        split="low.txt",
        solver=solver,
        max_messages=max_messages,
        scorer=scorer,
        epochs=epochs,
        docker_image_name=docker_image_name,
        force_rebuild=force_rebuild,
        data_dir=data_dir,
    )


def default_solver(max_messages: int = 30) -> Solver:
    return basic_agent(
        init=[
            system_message(
                "Please solve the coding task below. Once you are done, use your submit tool."
            )
        ],
        tools=[bash(timeout=180, user="nonroot")],
        max_messages=max_messages,
    )


@scorer(
    metrics={
        "valid_submission": [accuracy(), stderr()],
        "above_median": [accuracy(), stderr()],
        "any_medal": [accuracy(), stderr()],
    }
)
def mle_bench_scorer() -> Scorer:
    async def scorer(state: TaskState, target: Target) -> Score:
        grading_script = dedent(f"""
            from pathlib import Path
            import json
            from mlebench.grade import grade_csv
            from mlebench.registry import registry
            competition = registry.set_data_dir(Path("/private/data")).get_competition("{state.sample_id}")
            report = grade_csv(Path("/home/submission/submission.csv"), competition)
            print(json.dumps(report.to_dict()))
        """)
        result = await sandbox().exec(["python", "-c", grading_script])

        try:
            report_data = json.loads(result.stdout)
        except json.JSONDecodeError:
            logger.info("Failed to parse grading report")
            report_data = {}

        return Score(
            value={
                "valid_submission": report_data.get("valid_submission", False),
                "above_median": report_data.get("above_median", False),
                "any_medal": report_data.get("any_medal", False),
            },
            explanation=result.stdout,
        )

    return scorer


def get_compose_file(
    competition_id: str,
    data_dir: Path | str,
    docker_image_name: str,
    force_rebuild: bool = False,
) -> str:
    import docker  # type: ignore

    image_compose_file = COMPOSE_FILES_DIR / f"{competition_id}.yaml"
    image_compose_file.parent.mkdir(parents=True, exist_ok=True)

    with image_compose_file.open(mode="w+") as f:
        f.write(
            dedent(f"""\
            services:
              default:
                image: {docker_image_name}
                platform: linux/amd64
                network_mode: none
                volumes:
                  - {Path(data_dir) / competition_id / "prepared" / "public"}:/home/data:ro
                  - {Path(data_dir) / competition_id / "prepared" / "private"}:/private/data/{competition_id}/prepared/private:ro
                environment:
                  - COMPETITION_ID={competition_id}
                x-local: true
                deploy:
                  resources:
                    limits:
                      cpus: '1'
        """)
        )

        try:
            docker.from_env().images.get(docker_image_name)
            image_exists = True
        except docker.errors.ImageNotFound:
            image_exists = False

        if force_rebuild or not image_exists:
            f.write(f"""    build:
      context: {FILE_DIR}
      network: host""")

    return str(image_compose_file)
