# Copyright (c) Alibaba, Inc. and its affiliates.

from typing import Any, Dict

from evalscope.api.benchmark import BenchmarkMeta, MultiChoiceAdapter
from evalscope.api.dataset import Sample
from evalscope.api.registry import register_benchmark
from evalscope.constants import Tags
from evalscope.utils.logger import get_logger

logger = get_logger()

SUBJECT_MAPPING = {
    'computer_network': ['Computer Network', '计算机网络', 'STEM'],
    'operating_system': ['Operating System', '操作系统', 'STEM'],
    'computer_architecture': ['Computer Architecture', '计算机组成', 'STEM'],
    'college_programming': ['College Programming', '大学编程', 'STEM'],
    'college_physics': ['College Physics', '大学物理', 'STEM'],
    'college_chemistry': ['College Chemistry', '大学化学', 'STEM'],
    'advanced_mathematics': ['Advanced Mathematics', '高等数学', 'STEM'],
    'probability_and_statistics': ['Probability and Statistics', '概率统计', 'STEM'],
    'discrete_mathematics': ['Discrete Mathematics', '离散数学', 'STEM'],
    'electrical_engineer': ['Electrical Engineer', '注册电气工程师', 'STEM'],
    'metrology_engineer': ['Metrology Engineer', '注册计量师', 'STEM'],
    'high_school_mathematics': ['High School Mathematics', '高中数学', 'STEM'],
    'high_school_physics': ['High School Physics', '高中物理', 'STEM'],
    'high_school_chemistry': ['High School Chemistry', '高中化学', 'STEM'],
    'high_school_biology': ['High School Biology', '高中生物', 'STEM'],
    'middle_school_mathematics': ['Middle School Mathematics', '初中数学', 'STEM'],
    'middle_school_biology': ['Middle School Biology', '初中生物', 'STEM'],
    'middle_school_physics': ['Middle School Physics', '初中物理', 'STEM'],
    'middle_school_chemistry': ['Middle School Chemistry', '初中化学', 'STEM'],
    'veterinary_medicine': ['Veterinary Medicine', '兽医学', 'STEM'],
    'college_economics': ['College Economics', '大学经济学', 'Social Science'],
    'business_administration': ['Business Administration', '工商管理', 'Social Science'],
    'marxism': ['Marxism', '马克思主义基本原理', 'Social Science'],
    'mao_zedong_thought': ['Mao Zedong Thought', '毛泽东思想和中国特色社会主义理论体系概论', 'Social Science'],
    'education_science': ['Education Science', '教育学', 'Social Science'],
    'teacher_qualification': ['Teacher Qualification', '教师资格', 'Social Science'],
    'high_school_politics': ['High School Politics', '高中政治', 'Social Science'],
    'high_school_geography': ['High School Geography', '高中地理', 'Social Science'],
    'middle_school_politics': ['Middle School Politics', '初中政治', 'Social Science'],
    'middle_school_geography': ['Middle School Geography', '初中地理', 'Social Science'],
    'modern_chinese_history': ['Modern Chinese History', '近代史纲要', 'Humanities'],
    'ideological_and_moral_cultivation': ['Ideological and Moral Cultivation', '思想道德修养与法律基础', 'Humanities'],
    'logic': ['Logic', '逻辑学', 'Humanities'],
    'law': ['Law', '法学', 'Humanities'],
    'chinese_language_and_literature': ['Chinese Language and Literature', '中国语言文学', 'Humanities'],
    'art_studies': ['Art Studies', '艺术学', 'Humanities'],
    'professional_tour_guide': ['Professional Tour Guide', '导游资格', 'Humanities'],
    'legal_professional': ['Legal Professional', '法律职业资格', 'Humanities'],
    'high_school_chinese': ['High School Chinese', '高中语文', 'Humanities'],
    'high_school_history': ['High School History', '高中历史', 'Humanities'],
    'middle_school_history': ['Middle School History', '初中历史', 'Humanities'],
    'civil_servant': ['Civil Servant', '公务员', 'Other'],
    'sports_science': ['Sports Science', '体育学', 'Other'],
    'plant_protection': ['Plant Protection', '植物保护', 'Other'],
    'basic_medicine': ['Basic Medicine', '基础医学', 'Other'],
    'clinical_medicine': ['Clinical Medicine', '临床医学', 'Other'],
    'urban_and_rural_planner': ['Urban and Rural Planner', '注册城乡规划师', 'Other'],
    'accountant': ['Accountant', '注册会计师', 'Other'],
    'fire_engineer': ['Fire Engineer', '注册消防工程师', 'Other'],
    'environmental_impact_assessment_engineer': ['Environmental Impact Assessment Engineer', '环境影响评价工程师', 'Other'],
    'tax_accountant': ['Tax Accountant', '税务师', 'Other'],
    'physician': ['Physician', '医师资格', 'Other']
}

# Based on the prompt template for Chinese evaluation
USER_PROMPT_TEMPLATE = """以下是中国关于{subject}的单项选择题，请选出其中的正确答案。你的回答的最后一行应该是这样的格式："答案：[LETTER]"（不带引号），其中 [LETTER] 是 A、B、C、D 中的一个。

问题：{question}
选项：
{choices}
""".lstrip()  # noqa: E501

FEWSHOT_TEMPLATE = """以下是一些示例问题：

{fewshot}

""".lstrip()


@register_benchmark(
    BenchmarkMeta(
        name='ceval',
        pretty_name='C-Eval',
        tags=[Tags.KNOWLEDGE, Tags.MULTIPLE_CHOICE, Tags.CHINESE],
        description="""
## Overview

C-Eval is a comprehensive Chinese evaluation benchmark designed to assess the knowledge and reasoning abilities of language models in Chinese. It covers 52 subjects ranging from STEM to humanities and social sciences, with questions from middle school to professional examination levels.

## Task Description

- **Task Type**: Multiple-Choice Question Answering (Chinese)
- **Input**: Chinese question with four answer choices (A, B, C, D)
- **Output**: Single correct answer letter
- **Subjects**: 52 subjects organized into 4 categories (STEM, Social Science, Humanities, Other)

## Key Features

- 13,948 multiple-choice questions across 52 subjects
- Questions sourced from Chinese middle school, high school, college, and professional exams
- Covers diverse domains including mathematics, physics, law, medicine, and more
- Includes explanations for validation split questions
- Standard benchmark for Chinese language model evaluation

## Evaluation Notes

- Default configuration uses **5-shot** examples from the dev split
- Questions and prompts are in Chinese
- Answers should follow the format: "答案：[LETTER]"
- Results can be aggregated by subject or category
- Use `subset_list` parameter to evaluate specific subjects
""",
        dataset_id='evalscope/ceval',
        subset_list=list(SUBJECT_MAPPING.keys()),
        metric_list=['acc'],
        few_shot_num=5,
        train_split='dev',
        eval_split='val',
        prompt_template=USER_PROMPT_TEMPLATE,
        few_shot_prompt_template=FEWSHOT_TEMPLATE,
    )
)
class CEVALAdapter(MultiChoiceAdapter):

    def __init__(self, **kwargs):

        super().__init__(**kwargs)

        self.category_map = {k: v[-1] for k, v in SUBJECT_MAPPING.items()}

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        # Build choices list from A, B, C, D fields
        choices = [record['A'], record['B'], record['C'], record['D']]
        subset = self.current_subset_name

        return Sample(
            input=record['question'],
            choices=choices,
            target=record['answer'],
            metadata={
                'id': record.get('id', ''),
                'explanation': record.get('explanation', ''),
                'subject': subset
            },
        )

    def sample_to_fewshot(self, sample: Sample) -> str:
        q_str = f"""问题：{sample.input}"""
        choices = sample.choices if sample.choices is not None else []
        opt_str_list = []
        for i, choice in enumerate(choices):
            opt_str_list.append(f"""{chr(65 + i)}. {choice}""")
        opt_str = '\n'.join(opt_str_list)
        opt_str = f"""选项：\n{opt_str}"""
        exp_str = f"""解析：{sample.metadata.get('explanation', '')}"""
        ans_str = f"""答案：{sample.target}"""
        final_str = '\n'.join([q_str, opt_str, exp_str, ans_str])

        return final_str

    def format_fewshot_template(self, fewshot, sample):
        fewshot_str = FEWSHOT_TEMPLATE.format(fewshot=fewshot)
        prompt_str = self.format_prompt_template(sample)
        return fewshot_str + '\n' + prompt_str

    def format_prompt_template(self, sample):
        subject_name = SUBJECT_MAPPING.get(sample.metadata['subject'])[1]
        choices = sample.choices if sample.choices is not None else []
        choices_str = '\n'.join([f'{chr(65 + i)}. {choice}' for i, choice in enumerate(choices)])

        return USER_PROMPT_TEMPLATE.format(subject=subject_name, question=sample.input, choices=choices_str)

    def extract_answer(self, prediction, task_state) -> str:
        """
        Extract the answer from the prediction based on the task state.

        Args:
            prediction (str): The model's prediction string
            task_state (dict): The current task state containing metadata

        Returns:
            str: The extracted answer from the prediction
        """
        import re

        # Use regex to find the answer in the format "答案：LETTER"
        match = re.search(r'答案：([A-D])', prediction)
        if match:
            return match.group(1)
        else:
            logger.warning(f'No valid answer found in prediction: {prediction}')
            return ''
