from time import time
from typing import List

from inspect_ai import Task, task
from inspect_ai.dataset import MemoryDataset, csv_dataset
from inspect_ai.scorer import accuracy, match, stderr
from inspect_ai.solver import generate

# https://github.com/openai/simple-evals/blob/main/mgsm_eval.py
ALL_LANGUAGES = ["bn", "de", "en", "es", "fr", "ja", "ru", "sw", "te", "th", "zh"]

LANG_TO_FPATH = {
    "bn": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_bn.tsv",
    "de": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_de.tsv",
    "en": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_en.tsv",
    "es": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_es.tsv",
    "fr": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_fr.tsv",
    "ja": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_ja.tsv",
    "ru": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_ru.tsv",
    "sw": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_sw.tsv",
    "te": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_te.tsv",
    "th": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_th.tsv",
    "zh": "https://openaipublic.blob.core.windows.net/simple-evals/mgsm_zh.tsv",
}
LANG_TO_INSTRUCTIONS = {
    "en": """Solve this math problem. Give the reasoning steps before giving the final answer on the last line by itself in the format of "Answer:". Do not add anything other than the integer answer after "Answer:".

{prompt}""",
    "bn": """এই গণিতের সমস্যাটি সমাধান করুন। চূড়ান্ত উত্তর দেওয়ার আগে যুক্তিসম্পন্ন পদক্ষেপ প্রদান করুন। চূড়ান্ত উত্তরটি একক সংখ্যা হিসাবে "উত্তর:" এর পরে শেষ লাইনে দিন। "উত্তর:" এর পরে অন্য কিছু যুক্ত করবেন না।.

{prompt}""",
    "de": """Löse dieses Mathematikproblem. Gib die Schritte zur Begründung an, bevor du die endgültige Antwort in der letzten Zeile alleine im Format "Antwort:" gibst. Füge nichts anderes als die ganzzahlige Antwort nach "Antwort:" hinzu.

{prompt}""",
    "es": """Resuelve este problema matemático. Proporciona los pasos de razonamiento antes de dar la respuesta final en la última línea por sí misma en el formato de "Respuesta:". No añadas nada más que la respuesta entera después de "Respuesta:".

{prompt}""",
    "fr": """Résolvez ce problème de mathématiques. Donnez les étapes de raisonnement avant de fournir la réponse finale sur la dernière ligne elle-même dans le format de "Réponse:". N'ajoutez rien d'autre que la réponse entière après "Réponse:".

{prompt}""",
    "ja": """の数学の問題を解いてください。最終的な答えを出す前に、解答の推論過程を記述してください。そして最後の行には "答え:" の形式で答えを記述し、その後には整数の答え以外何も追加しないでください。

{prompt}""",
    "ru": """Решите эту математическую задачу. Объясните шаги рассуждения перед тем, как дать окончательный ответ в последней строке сам по себе в формате "Ответ:". Не добавляйте ничего, кроме целочисленного ответа после "Ответ:".

{prompt}""",
    "sw": """Suluhisha tatizo hili la hesabu. Toa hatua za mantiki kabla ya kutoa jibu la mwisho kwenye mstari wa mwisho peke yake katika muundo wa "Jibu:". Usiongeze chochote kingine isipokuwa jibu la integer baada ya "Jibu:".

{prompt}""",
    "te": """ఈ గణిత సమస్యను పరిష్కరించండి. చివరి సమాధానాన్ని ఇవ్వదానికి ముందు తర్కాత్మక అదుగులను ఇవ్వండి. చివరి పంక్తిలో మాత్రమే 'సమాధానం:' అనే ఆకారంలో చివరి సమాధానాద్ని ఇవ్వండి సమాధానం: తర్వాత పూర్ణాంక సమాధానానికి తప్పించి ఎదేనా చేర్చవద్దు.

{prompt}""",
    "th": """แก้ปัญหาคณิตศาสตร์นี้ ให้ให้ขั้นตอนการใช้เหตุผลก่อนที่จะให้คำตอบสุดท้ายในบรรทัดสุดท้ายโดยอยู่ในรูปแบบ "คำตอบ:" ไม่ควรเพิ่มอะไรนอกจากคำตอบที่เป็นจำนวนเต็มหลังจาก "คำตอบ:"

{prompt}""",
    "zh": """解决这个数学问题。在最后一行给出答案前，请提供推理步骤。最后一行应该以 "答案: " 的形式独立给出答案。在 "答案：" 后不要添加除整数答案之外的任何内容。

{prompt}""",
}


def load_mgsm_dataset(
    languages: str | List[str] = "all",
    limit_samples_per_lang: None | int = None,
    use_cot: bool = True,
    shuffle: bool = False,
) -> MemoryDataset:
    if languages == "all":
        languages = ALL_LANGUAGES

    samples = []
    seed = int(time())  # use same shuffling seed for all languages

    for lang in languages:
        csv_filename = LANG_TO_FPATH[lang]
        lang_dataset = csv_dataset(
            csv_filename,
            fieldnames=["input", "target"],
            dialect="excel-tab",
            limit=limit_samples_per_lang,
            shuffle=shuffle,
            seed=seed,
            auto_id=True,
            delimiter="\t",
        )

        lang_samples = lang_dataset.samples  # type: ignore
        cot_template = LANG_TO_INSTRUCTIONS[lang]

        for sample in lang_samples:
            sample.metadata = {"language": lang, "question_id": sample.id}
            sample.id = f"{lang}_{sample.id}"
            if use_cot:
                cot_prompt = cot_template.format(prompt=sample.input)
                sample.input = cot_prompt

        samples.extend(lang_samples)

    mgsm_dataset = MemoryDataset(samples=samples, name="mgsm")
    return mgsm_dataset


@task
def mgsm(
    languages: str | List[str] = "all",
    limit_samples_per_lang: None | int = None,
    use_cot: bool = True,
    shuffle: bool = False,
) -> Task:
    dataset = load_mgsm_dataset(
        languages=languages,
        limit_samples_per_lang=limit_samples_per_lang,
        use_cot=use_cot,
        shuffle=shuffle,
    )

    task = Task(
        dataset=dataset,
        solver=[generate()],
        scorer=match(numeric=True),
        metrics=[accuracy(), stderr(cluster="question_id")],
    )
    return task
