# SPDX-FileCopyrightText: 2024 Arthit Suriyawongkul
# SPDX-FileType: SOURCE
# SPDX-License-Identifier: MIT

"""
O-NET

# inspect eval inspect_evals/onet
"""

import hashlib
import re
from typing import Any, List, Literal, Tuple

from inspect_ai import Task, task
from inspect_ai.dataset import Sample, hf_dataset
from inspect_ai.scorer import choice
from inspect_ai.solver import multiple_choice, system_message

SYSTEM_MESSAGE = """The question can be in either English or Thai.""".strip()


@task
def onet_m6(
    dataset_name: Literal[
        "default", "english", "math", "science", "social", "thai"
    ] = "default",
) -> Task:
    """Inspect Task implementation of O-NET"""
    dataset = hf_dataset(
        path="openthaigpt/thai-onet-m6-exam",
        name=dataset_name,
        sample_fields=record_to_sample,
        split="test",
        shuffle=True,
    )

    # Pick only answerable multiple-choice questions with only one correct answer
    dataset = dataset.filter(
        lambda sample: sample.metadata is not None
        and sample.metadata["is_answerable"]
        and sample.metadata["is_multiple_choice"]
        and sample.metadata["is_single_choice_solution"]
    )

    return Task(
        dataset=dataset,
        solver=[
            system_message(SYSTEM_MESSAGE),
            multiple_choice(cot=True),
        ],
        scorer=choice(),
    )


# Utility function to normalize digits
_thai_arabic = {
    "๐": "0",
    "๑": "1",
    "๒": "2",
    "๓": "3",
    "๔": "4",
    "๕": "5",
    "๖": "6",
    "๗": "7",
    "๘": "8",
    "๙": "9",
}
_thai_arabic_translate_table = str.maketrans(_thai_arabic)


def thai_digit_to_arabic_digit(text: str) -> str:
    """
    Converts Thai digits in the given text to Arabic digits

    :param str text: Text containing Thai digits
    :return: Text with Thai digits converted to Arabic digits
    :rtype: str
    """
    if not text or not isinstance(text, str):
        return ""

    return text.translate(_thai_arabic_translate_table)


def split_numbered_list(text: str) -> Tuple[str, List[str]]:
    """
    Splits choices.

    Splits a string containing a numbered list into a tuple containing a
    substring before the numbered list and a list of strings,
    each containing one numbered item.

    A substring before the numbered list (which can be additional instructions)
    will be returned as the first item of the tuple.
    This can be blank if there is none.

    The numbered list will be returned as the second item of the tuple.

    :param str text: Text containing a numbered list.
    :return: A tuple containing the substring before the numbered list and a list of strings,
             each containing one numbered item.
    :rtype: Tuple[str, List[str]]
    """
    items = re.split(r"(?=\d+\.\s)", text)

    lead_text = items[0].strip()
    numbered_list = [re.sub(r"^\d+\.\s*", "", i).strip() for i in items[1:]]

    return lead_text, numbered_list


def extract_choice_letter(text: str) -> str:
    """
    Extracts the number at the beginning of the given text and converts it to a corresponding letter.

    :param str text: Text containing a number at the beginning.
    :return: The corresponding letter for the extracted number.
    :rtype: str
    """
    match = re.match(r"^\d+", text)
    if match:
        number = int(match.group(0))
        # Convert number to corresponding letter (1 -> A, 2 -> B, ..., 26 -> Z)
        # If the number is over 26, we run out of Latin characters for it,
        # so just don't try.
        # In practice, we will have at max only 5 choices (1-5 -> A-E).
        if 1 <= number <= 26:
            return chr(ord("A") + number - 1)
    return ""


def record_to_sample(record: dict[str, Any]) -> Sample:
    """
    Converts a dataset record to an Inspect Sample

    Some metadata are also being converted.
    """
    is_answerable = record["isAnswerable"]
    is_multiple_choice = record["isMultipleChoice"]
    is_single_choice_solution = record["isSingleChoiceSolution"]

    question = thai_digit_to_arabic_digit(record["instruction"]).strip()
    choices = thai_digit_to_arabic_digit(record["input"]).strip()
    additional_instructions, choices_list = split_numbered_list(choices)
    answer = (
        extract_choice_letter(thai_digit_to_arabic_digit(record["result"]).strip())
        if is_answerable
        else "unanswerable"
    )

    year = record["year"]
    no = (
        record['"no"'].replace('"', "").strip()
    )  # The column name is 'no' (with single quotes)
    hash = hashlib.md5(question.encode()).hexdigest()[:4]
    # The dataset contains duplicate sample ids (duplicate id: 2021_50),
    # so a hash is needed for a unique id.
    # We're not using auto_id because we want to see year/no in the log,
    # so we can look it up in the dataset more conveniently.
    id = f"{year}_{no}_{hash}"

    return Sample(
        input=f"{question}\n\n{additional_instructions}".strip(),
        choices=choices_list,
        target=answer,
        id=id,
        metadata={
            "year": year,
            "is_answerable": is_answerable,
            "is_multiple_choice": is_multiple_choice,
            "is_single_choice_solution": is_single_choice_solution,
        },
    )
