import re

from datasets import load_dataset

from opencompass.registry import LOAD_DATASET, TEXT_POSTPROCESSORS

from .base import BaseDataset


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

    @staticmethod
    def load(path: str):
        return load_dataset('csv', data_files={'test': path})


@TEXT_POSTPROCESSORS.register_module('TheoremQA')
def TheoremQA_postprocess(text: str) -> str:
    text = text.strip()
    matches = re.findall(r'answer is ([^\s]+)', text)
    if len(matches) == 0:
        return text
    else:
        text = matches[0].strip().strip('.,?!\"\';:')
        return text


def TheoremQA_postprocess_v2(text: str) -> str:
    prediction = text.strip().strip('\n').split('\n')[-1]
    tmp = ''
    for entry in prediction.split(' ')[::-1]:
        if entry == 'is' or entry == 'be' or entry == 'are' or entry.endswith(
                ':'):
            break
        tmp = entry + ' ' + tmp
    prediction = tmp.strip().strip('.')
    return prediction