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

from typing import Any, Dict, List, Tuple

from evalscope.api.benchmark import BenchmarkMeta, DefaultDataAdapter
from evalscope.api.dataset import Sample
from evalscope.api.dataset.dataset import DatasetDict, MemoryDataset
from evalscope.api.metric.scorer import AggScore, SampleScore, Score
from evalscope.api.registry import register_benchmark
from evalscope.constants import Tags
from evalscope.report.report import Report, Subset
from evalscope.utils.logger import get_logger

logger = get_logger()

SUBSET_LIST = [
    'en', 'zh', 'ar', 'bn', 'de', 'es', 'fr', 'id', 'it', 'ja', 'ko', 'ms', 'pt', 'ru', 'sw', 'te', 'th', 'vi'
]
LEVEL_LIST = ['low', 'medium', 'high', 'top']


@register_benchmark(
    BenchmarkMeta(
        name='poly_math',
        pretty_name='PolyMath',
        tags=[Tags.MATH, Tags.REASONING, Tags.MULTI_LINGUAL],
        description="""
## Overview

PolyMath is a multilingual mathematical reasoning benchmark covering 18 languages and 4 difficulty levels with 9,000 high-quality problem samples. It ensures difficulty comprehensiveness, language diversity, and high-quality translation for discriminative multilingual evaluation.

## Task Description

- **Task Type**: Multilingual Mathematical Reasoning
- **Input**: Math problem in one of 18 languages
- **Output**: Numerical answer in \\boxed{} format
- **Domains**: Mathematics across multiple difficulty levels and languages

## Key Features

- 18 supported languages: en, zh, ar, bn, de, es, fr, id, it, ja, ko, ms, pt, ru, sw, te, th, vi
- 4 difficulty levels: low, medium, high, top
- 9,000 high-quality problems total
- Language-specific instructions for each problem
- High-quality human translations ensuring accuracy

## Evaluation Notes

- Default evaluation uses the **test** split
- Primary metric: **Accuracy** with numeric comparison
- Additional metric: **DW-ACC** (Difficulty-Weighted Accuracy)
  - Weights: low=1, medium=2, high=4, top=8
  - Provides balanced scoring across difficulty levels
- Results reported per language and overall
""",  # noqa: E501
        dataset_id='evalscope/PolyMath',
        subset_list=SUBSET_LIST,
        metric_list=[{
            'acc': {
                'numeric': True
            }
        }],
        eval_split='test',
        prompt_template='{question}',
    )
)
class PolyMathAdapter(DefaultDataAdapter):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def load(self) -> Tuple[DatasetDict, None]:
        """Load all difficulty levels, rename subsets with their level suffix, and merge them."""
        # Need to load all levels to get the full dataset
        dataset_list: List[Dict[str, MemoryDataset]] = []
        original_split = getattr(self, 'eval_split', None)
        try:
            for level in LEVEL_LIST:
                self.eval_split = level
                cur_level_dataset_dict, _ = super().load()
                # Build a renamed mapping without mutating during iteration
                renamed: Dict[str, MemoryDataset] = {
                    f'{subset}-{level}': ds
                    for subset, ds in cur_level_dataset_dict.items()
                }
                dataset_list.append(renamed)
        finally:
            # Restore original split to avoid side effects
            if original_split is not None:
                self.eval_split = original_split
        # Merge all levels into one dataset
        return DatasetDict.from_dataset_dicts(dataset_list), None

    def record_to_sample(self, record: Dict[str, Any]) -> Sample:
        """Convert a single record into a Sample with language-specific instruction."""
        from .utils.instruction import query_dic

        # e.g., 'high-en-1'
        question_id = record['id']
        level, language, index = question_id.split('-')
        # Get the instruction for the specific language
        instruction = query_dic[language]

        return Sample(
            input=record['question'] + '\n' + instruction,
            target=record['answer'],
            metadata={
                'level': level,
                'language': language,
                'index': index,
            },
        )

    def extract_answer(self, prediction: str, task_state):
        from evalscope.metrics.math.parser import extract_answer

        return extract_answer(prediction)

    def _on_generate_report_end(self, report: Report, output_dir, **kwargs):
        """
        Finalize the report generation process. Calculate the difficulty-weighted accuracy (DW-ACC)
        per language and the overall DW-ACC, and append as a new category to each metric.
        """
        from evalscope.report import Category, Metric

        WEIGHT_DENOMINATOR = 15  # 1 + 2 + 4 + 8 for ['low','medium','high','top']

        for metric in report.metrics:
            # Collect all subsets by name for easy lookup (e.g., "en-low")
            subset_dict: Dict[str, Subset] = {}
            for category in metric.categories:
                for subset in category.subsets:
                    subset_dict[subset.name] = subset

            # Compute per-language DW-ACC
            dw_subsets: List[Subset] = []
            for language in SUBSET_LIST:
                weighted_sum = 0.0
                total_num = 0
                for i, level in enumerate(LEVEL_LIST):
                    s = subset_dict.get(f'{language}-{level}')
                    if s is not None:
                        weighted_sum += (2**i) * s.score
                        total_num += s.num
                # missing subsets contribute 0 by design
                if total_num == 0:
                    continue
                dw_acc = weighted_sum / WEIGHT_DENOMINATOR
                dw_subsets.append(Subset(name=language, score=dw_acc, num=total_num))

            # Overall DW-ACC: unweighted average across languages
            if dw_subsets:
                overall_score = sum(s.score for s in dw_subsets) / len(dw_subsets)
                overall_num = sum(s.num for s in dw_subsets)
                dw_subsets.append(Subset(name='OVERALL', score=overall_score, num=overall_num))

        # Append DW-ACC metric
        if dw_subsets:
            report.metrics.append(Metric(name='DW-ACC', categories=[Category(name='-', subsets=dw_subsets)]))
