import json

from datasets import Dataset, DatasetDict

from opencompass.registry import LOAD_DATASET

from .base import BaseDataset


@LOAD_DATASET.register_module()
class dropDataset(BaseDataset):

    @staticmethod
    def get_answers(validated_answers):
        answers = []
        for answer_item in validated_answers:
            if answer_item['number']:
                answers.append(answer_item['number'])
            elif any(answer_item['date'][i] for i in ['day', 'month', 'year']):
                d = [answer_item['date'][i] for i in ['day', 'month', 'year']]
                answers.append(' '.join(d).strip())
            else:
                for span in answer_item['spans']:
                    answers.append(span)
        answers = list(set(answers))
        return answers

    @staticmethod
    def load(path, only_number=True):
        with open(path, 'r', encoding='utf-8') as f:
            lines = json.load(f)
        dataset_list = []
        for line in lines.values():
            for qa_pair in line['qa_pairs']:
                validated_answers = qa_pair['validated_answers']
                if only_number and not any(i['number']
                                           for i in validated_answers):
                    continue
                item = {
                    'prompt': line['passage'],
                    'question': qa_pair['question'],
                    'answers': dropDataset.get_answers(validated_answers),
                }
                dataset_list.append(item)

        dataset_list = Dataset.from_list(dataset_list)
        return DatasetDict({'validation': dataset_list})