"""
SciCode: A Research Coding Benchmark Curated by Scientists

Minyang Tian, Luyu Gao, Shizhuo Dylan Zhang, Xinan Chen, Cunwei Fan, Xuefei Guo, Roland Haas, Pan Ji, Kittithat Krongchon, Yao Li, Shengyan Liu, Di Luo, Yutao Ma, Hao Tong, Kha Trinh, Chenyu Tian, Zihan Wang, Bohao Wu, Yanyu Xiong, Shengzhu Yin, Minhui Zhu, Kilian Lieret, Yanxin Lu, Genglin Liu, Yufeng Du, Tianhua Tao, Ofir Press, Jamie Callan, Eliu Huerta, Hao Peng

https://arxiv.org/abs/2407.13168
"""

from pathlib import Path

from inspect_ai import Task, task
from inspect_ai.solver import (
    system_message,
)

from .dataset import get_dataset
from .prompt_templates import (
    INITIAL_PROMPT,
    INITIAL_PROMPT_PROVIDE_BACKGROUND,
    SUBPROBLEM_PROMPT,
    SUBPROBLEM_PROMPT_PROVIDE_BACKGROUND,
)
from .scorer import verify
from .solver import solve_scicode_problem


@task
def scicode(
    *,
    provide_scientific_background: bool = False,
    timeout: int = 300,
    include_dev_set: bool = False,
) -> Task:
    """
    Run SciCode.

    Parameters:
    -----------
    provide_scientific_background
       Include scientific background information written by scientists for the problem in the model's prompt.
    timeout
        Maximum time (in seconds) allowed for execution of generated code.
    include_dev_set

    Returns:
    --------
    Task

    """
    SYS_MESSAGE = (
        INITIAL_PROMPT_PROVIDE_BACKGROUND
        if provide_scientific_background
        else INITIAL_PROMPT
    )

    PROMPT_TEMPLATE = (
        SUBPROBLEM_PROMPT_PROVIDE_BACKGROUND
        if provide_scientific_background
        else SUBPROBLEM_PROMPT
    )

    return Task(
        dataset=get_dataset(include_dev_set),
        solver=[
            system_message(SYS_MESSAGE),
            solve_scicode_problem(PROMPT_TEMPLATE),
        ],
        scorer=verify(timeout),
        sandbox=("docker", (Path(__file__).parent / "Dockerfile").as_posix()),
    )
