from dataclasses import dataclass, field
from datasets import load_dataset, Dataset
from functools import cached_property
from tqdm.auto import tqdm
from typing import Any, Optional, Protocol, Iterable, Callable
import logging
import pandas as pd
from functools import partial
from datasets.utils.logging import disable_progress_bar

from .utils import *

from evaluate import load
from collections import defaultdict
import sys

# if sys.version_info >= (3, 9):
#     from functools import cache
# else:
#     from functools import lru_cache as cache

disable_progress_bar()


def fake_pipeline(prompts: Iterable[str]) -> list[str]:
    return [prompt for prompt in tqdm(prompts)]


@dataclass
class Task:
    dataset_name: str | tuple[str, str] = ("gsm8k", "main")
    split: str = "test"
    # metrics: list[str] = field(default_factory=list)
    metric_name: str | tuple[str, str] = ("sustech/tlem", "gsm8k")
    input_column: str = "question"
    label_column: str = "answer"
    prompt: Optional[Callable | str] = None
    few_shot: int = 0
    few_shot_from: Optional[str] = None
    # results: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        names = (
            [self.dataset_name]
            if isinstance(self.dataset_name, str)
            else list(self.dataset_name)
        )
        names[0] = names[0].split("/")[-1]

        self.name = "-".join(names) + f"-{self.split}"
        if isinstance(self.prompt, str):
            self.prompt = lambda example: {
                self.input_column: self.prompt.format(
                    input_column=example[self.input_column]
                )
            }

    @cached_property
    def samples(self):
        return self.dataset[self.input_column]

    @cached_property
    def dataset(self):
        ds = load_dataset(
            *self.dataset_name
            if isinstance(self.dataset_name, tuple)
            else self.dataset_name,
            # split=self.split,
        )
        test_ds = ds[self.split]
        if self.prompt is not None:
            test_ds = test_ds.map(self.prompt)

        if self.few_shot:
            if self.few_shot_from is None:
                for name in ["train", "validation", "val", "dev"]:
                    if name in ds:
                        self.few_shot_from = name
                        break

            shots = ds[self.few_shot_from].select(range(self.few_shot))
            if self.prompt is not None:
                shots = shots.map(self.prompt)

            shots = shots.map(
                lambda example: {
                    self.input_column: example[self.input_column]
                    + example[self.label_column],
                }
            )[self.input_column]
            few_shot_prompts = "\n\n".join(shots)

            test_ds = test_ds.map(
                lambda example: {
                    self.input_column: few_shot_prompts
                    + "\n\n"
                    + example[self.input_column],
                }
            )

        return test_ds

    @cached_property
    def metric(self):
        metric = (
            load(self.metric_name)
            if isinstance(self.metric_name, str)
            else load(*self.metric_name)
        )
        return metric

    # @cache
    def run(
        self,
        pipeline,
    ):
        if (outputs := pipeline(self.samples)) is None:
            logging.warning("pipeline returns None")
            return
        self.outputs = outputs
        try:
            result = self.metric._compute(
                responses=outputs, references=self.dataset[self.label_column]
            )
        except Exception as e:
            result = self.metric.compute(
                responses=outputs, references=self.dataset[self.label_column]
            )
        # if log:
        #     name = name or pipeline.__name__
        #     self.results[name] = result

        return result


def multichoice(responses: Any, references: list[str]):
    if isinstance(responses[0], str):
        responses = [extract_choice(response) for response in responses]
    else:
        responses = decode_choice(responses)

    return responses, references


def multichoice_zh(responses: Any, references: list[str]):
    if isinstance(responses[0], str):
        responses = [extract_choice_zh(response) for response in responses]
    else:
        responses = decode_choice(responses)

    return responses, references


class Metrics:
    cmmlu = multichoice_zh
    mmlu = multichoice

    def gsm8k(responses: list[str], answers: list[str | int]):
        # scores = []
        # for response, answer in zip(responses, answers):
        #     pred = extract_numeric(response)
        #     gold = extract_numeric(answer) if isinstance(answer, str) else str(answer)
        #     scores.append(1.0 * (pred == gold))
        responses = [extract_numeric(response) for response in responses]
        answers = [
            extract_numeric(answer) if isinstance(answer, str) else str(answer)
            for answer in answers
        ]

        return responses, answers

    def MATH(responses: list[str], answers: list[str]):
        scores = []

        for response, answer in zip(responses, answers):
            indices = [pos for pos, char in enumerate(response) if char == "$"]
            if len(indices) <= 2:
                scores.append(0)
                continue
            else:
                result = response[indices[-2] + 1 : indices[-1]]
                gold = get_answer(answer)
                scores.append(1.0 * is_equiv(result, gold))

        return scores

    def math23k(responses: list[str], answers: list[str]):
        scores = []
        for response, answer in zip(responses, answers):
            pred = extract_numeric(response, pattern=NUMERIC_IN_ZH)
            gold = extract_numeric(answer, pattern=NUMERIC_IN_ZH)
            scores.append(1.0 * (pred == gold))
        return scores


class CMMLU:
    input_column = "prompt"
    label_column = "Answer"

    def prompt_cmmlu(example, chat=False):
        prefix = "以下是一道多项选择题，请从A、B、C和D中选择最合适的答案作为这个问题的答案。\n\n" if chat else "问题："
        prompt = prefix + example["Question"]
        for choice in list("ABCD"):
            prompt += f"\n{choice}. {example[choice]}"

            prompt += "\n答案："
        return {"prompt": prompt}

    subcategories = {
        "agronomy": ["other"],
        "anatomy": ["biology"],
        "ancient_chinese": ["linguistics", "china specific"],
        "arts": ["arts"],
        "astronomy": ["physics"],
        "business_ethics": ["business"],
        "chinese_civil_service_exam": ["politics", "china specific"],
        "chinese_driving_rule": ["other", "china specific"],
        "chinese_food_culture": ["culture", "china specific"],
        "chinese_foreign_policy": ["politics", "china specific"],
        "chinese_history": ["history", "china specific"],
        "chinese_literature": ["literature", "china specific"],
        "chinese_teacher_qualification": ["education", "china specific"],
        "college_actuarial_science": ["math"],
        "college_education": ["education"],
        "college_engineering_hydrology": ["engineering"],
        "college_law": ["law"],
        "college_mathematics": ["math"],
        "college_medical_statistics": ["statistics"],
        "clinical_knowledge": ["other"],
        "college_medicine": ["other"],
        "computer_science": ["computer science"],
        "computer_security": ["other"],
        "conceptual_physics": ["physics"],
        "construction_project_management": ["other", "china specific"],
        "economics": ["economics"],
        "education": ["education"],
        "elementary_chinese": ["linguistics", "china specific"],
        "elementary_commonsense": ["other", "china specific"],
        "elementary_information_and_technology": ["other"],
        "electrical_engineering": ["engineering"],
        "elementary_mathematics": ["math"],
        "ethnology": ["culture", "china specific"],
        "food_science": ["other"],
        "genetics": ["biology"],
        "global_facts": ["global"],
        "high_school_biology": ["biology"],
        "high_school_chemistry": ["chemistry"],
        "high_school_geography": ["geography"],
        "high_school_mathematics": ["math"],
        "high_school_physics": ["physics"],
        "high_school_politics": ["politics", "china specific"],
        "human_sexuality": ["other"],
        "international_law": ["law"],
        "journalism": ["sociology"],
        "jurisprudence": ["law"],
        "legal_and_moral_basis": ["other"],
        "logical": ["philosophy"],
        "machine_learning": ["computer science"],
        "management": ["business"],
        "marketing": ["business"],
        "marxist_theory": ["philosophy"],
        "modern_chinese": ["linguistics", "china specific"],
        "nutrition": ["other"],
        "philosophy": ["philosophy"],
        "professional_accounting": ["business"],
        "professional_law": ["law"],
        "professional_medicine": ["other"],
        "professional_psychology": ["psychology"],
        "public_relations": ["politics"],
        "security_study": ["politics"],
        "sociology": ["culture"],
        "sports_science": ["other"],
        "traditional_chinese_medicine": ["other", "china specific"],
        "virology": ["biology"],
        "world_history": ["history"],
        "world_religions": ["global"],
    }

    categories = {
        "STEM": [
            "physics",
            "chemistry",
            "biology",
            "computer science",
            "math",
            "engineering",
            "statistics",
        ],
        "Humanities": ["history", "philosophy", "law", "arts", "literature", "global"],
        "Social Science": [
            "linguistics",
            "business",
            "politics",
            "culture",
            "economics",
            "geography",
            "psychology",
            "education",
            "sociology",
        ],
        "Other": ["other"],
        "China specific": ["china specific"],
        "Test": ["computer science"],
    }

    @classmethod
    def suite(cls, chat=False):
        finer_categories = (
            pd.Series(cls.subcategories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        cls.categories["all"] = list(finer_categories.keys())
        for k, v in cls.categories.items():
            for subject in v:
                suite[k].extend(
                    [
                        Task(
                            ("haonan-li/cmmlu", subcategories),
                            metric_name=("sustech/tlem", "cmmlu"),
                            input_column=cls.input_column,
                            label_column=cls.label_column,
                            prompt=partial(cls.prompt_cmmlu, chat=chat),
                            few_shot=0 if chat else 5,
                            few_shot_from="dev",
                        )
                        for subcategories in finer_categories[subject]
                    ]
                )
        return suite


class MMLU:
    input_column = "prompt"
    label_column = "target"

    @classmethod
    def prompt_mmlu(cls, example, chat=False):
        prefix = (
            "The following is a multiple-choice question. Please choose the most suitable one among A, B, C and D as the answer to this question.\n\n"
            if chat
            else "Question: "
        )
        prompt = prefix + example["input"]
        for choice in list("ABCD"):
            prompt += f"\n{choice}. {example[choice]}"

        prompt += "\nAnswer："
        return {"prompt": prompt}

    subcategories = {
        "abstract_algebra": ["math"],
        "anatomy": ["health"],
        "astronomy": ["physics"],
        "business_ethics": ["business"],
        "clinical_knowledge": ["health"],
        "college_biology": ["biology"],
        "college_chemistry": ["chemistry"],
        "college_computer_science": ["computer science"],
        "college_mathematics": ["math"],
        "college_medicine": ["health"],
        "college_physics": ["physics"],
        "computer_security": ["computer science"],
        "conceptual_physics": ["physics"],
        "econometrics": ["economics"],
        "electrical_engineering": ["engineering"],
        "elementary_mathematics": ["math"],
        "formal_logic": ["philosophy"],
        "global_facts": ["other"],
        "high_school_biology": ["biology"],
        "high_school_chemistry": ["chemistry"],
        "high_school_computer_science": ["computer science"],
        "high_school_european_history": ["history"],
        "high_school_geography": ["geography"],
        "high_school_government_and_politics": ["politics"],
        "high_school_macroeconomics": ["economics"],
        "high_school_mathematics": ["math"],
        "high_school_microeconomics": ["economics"],
        "high_school_physics": ["physics"],
        "high_school_psychology": ["psychology"],
        "high_school_statistics": ["math"],
        "high_school_us_history": ["history"],
        "high_school_world_history": ["history"],
        "human_aging": ["health"],
        "human_sexuality": ["culture"],
        "international_law": ["law"],
        "jurisprudence": ["law"],
        "logical_fallacies": ["philosophy"],
        "machine_learning": ["computer science"],
        "management": ["business"],
        "marketing": ["business"],
        "medical_genetics": ["health"],
        "miscellaneous": ["other"],
        "moral_disputes": ["philosophy"],
        "moral_scenarios": ["philosophy"],
        "nutrition": ["health"],
        "philosophy": ["philosophy"],
        "prehistory": ["history"],
        "professional_accounting": ["other"],
        "professional_law": ["law"],
        "professional_medicine": ["health"],
        "professional_psychology": ["psychology"],
        "public_relations": ["politics"],
        "security_studies": ["politics"],
        "sociology": ["culture"],
        "us_foreign_policy": ["politics"],
        "virology": ["health"],
        "world_religions": ["philosophy"],
    }

    categories = {
        "STEM": [
            "physics",
            "chemistry",
            "biology",
            "computer science",
            "math",
            "engineering",
        ],
        "humanities": ["history", "philosophy", "law"],
        "social sciences": [
            "politics",
            "culture",
            "economics",
            "geography",
            "psychology",
        ],
        "other": ["other", "business", "health"],
        "Test": ["culture"],
    }

    @classmethod
    def suite(cls, chat=False):
        finer_categories = (
            pd.Series(cls.subcategories)  # noqa # type: ignore
            .explode()
            .reset_index()
            .set_index(0)
            .groupby(0)
            .agg(list)["index"]
            .to_dict()
        )
        suite = defaultdict(list)
        cls.categories["all"] = list(finer_categories.keys())
        for k, v in cls.categories.items():
            for subject in v:
                suite[k].extend(
                    [
                        Task(
                            ("lukaemon/mmlu", subcategories),
                            metric_name=("sustech/tlem", "mmlu"),
                            input_column=cls.input_column,
                            label_column=cls.label_column,
                            prompt=partial(cls.prompt_mmlu, chat=chat),
                            few_shot=0 if chat else 5,
                            few_shot_from="validation",
                        )
                        for subcategories in finer_categories[subject]
                    ]
                )
        return suite