from abc import ABC, abstractmethod
from torch.utils.data import Dataset as tDataset
from typing import Union, Dict, List


class PromptedDataset(tDataset, ABC):
    @abstractmethod
    def get_problem_system_instruction(self):
        pass


class DatasetWithReference(PromptedDataset, ABC):
    @abstractmethod
    def calculate_correctness(self, item, answer, *args, **kwargs) -> Dict:
        pass

    @abstractmethod
    def get_answer(self, item):
        pass

class DatasetWithExactSolution(PromptedDataset, ABC):
    @abstractmethod
    def calculate_exact_correctness(self, item, answer, *args, **kwargs) -> Dict:
        pass

class DatasetWithOODInfo(PromptedDataset, ABC):
    @abstractmethod
    def get_ood_identifier(self, item, *args, **kwargs) -> Dict:
        pass

class DatasetWithPerturbation(PromptedDataset, ABC):
    @abstractmethod
    def get_perturbation_strength(self, item, *args, **kwargs) -> Dict:
        pass

    @abstractmethod
    def get_perutrbation_group_ids(self, item) -> List:
        pass


OPEN_BOOK_QA_MSG = """Assistant is an intelligent question answering assistant that gives answers to questions relating a given text. \
Read the following text carefully and answer the users question concisely:\n"""

CLOSED_BOOK_QA_MSG = """Assistant is an intelligent question answering assistant that gives concise answers to questions asked by the user."""

STRUCTURED_TEXT_GEN_MSG = """You are an assistant that creates text according to provided requirements. \
 Read the users requirements carefully and compose the output text as specified. \
 Only provide the text that fits the given criteria and nothing else.\
 DO NOT PROVIDE quotation marks.\
 DO NOT PROVIDE assertions (e.g. 'Here is a sentence with' or 'Here is a paragraph that' and similar).\
 DO NOT PROVIDE confirmations or summaries (e.g. 'the number of characters in this text is ...')!
 The correctness of the answers will be checked automatically and such unnecessary things will fail the otherwise correct response!"""

R1_REASON_PROMPT = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it.
The assistant first thinks about the reasoning process in the mind and then provides the user
with the answer. The reasoning process and answer are enclosed within <think> </think> and
<answer> </answer> tags, respectively, i.e., <think> reasoning process here </think>
<answer> answer here </answer>.
"""
