from abc import ABC, abstractmethod
from typing import Optional

from transformers import PreTrainedTokenizerBase

from trapper.common import Registrable
from trapper.data.data_processors.data_processor import IndexedInstance
from trapper.data.label_mapper import LabelMapper
from trapper.data.tokenizers import TokenizerWrapper


class DataAdapter(ABC, Registrable):
    """
    This callable class is responsible from converting the data instances that
    are already tokenized and indexed into a format suitable for feeding into a
    transformer model. Typically, it receives its inputs from a `DataProcessor`
    and adapts the input fields by renaming them to the names accepted by the
    models e.g. "input_ids", "token_type_ids" etc. Moreover, it also handles
    the insertion of the special tokens signaling the beginning or ending of a
    sequence such as `[CLS]`, `[SEP]` etc. To extend this class, you need to
    implement the `__call__` method as suitable for your task. See
    `DataAdapterForQuestionAnswering` for an example.

    Args:
        tokenizer_wrapper (): Required to access the ids of special tokens such
            as BOS, EOS etc
        label_mapper (): Only used in some tasks that require mapping between
            categorical labels and integer ids such as token classification.
    """

    def __init__(
        self,
        tokenizer_wrapper: TokenizerWrapper,
        label_mapper: Optional[LabelMapper] = None,
    ):
        tokenizer = tokenizer_wrapper.tokenizer
        self._bos_token_id: int = tokenizer.bos_token_id
        self._eos_token_id: int = tokenizer.eos_token_id
        self._tokenizer: PreTrainedTokenizerBase = tokenizer
        self._label_mapper = label_mapper

    @property
    def tokenizer(self) -> PreTrainedTokenizerBase:
        return self._tokenizer

    @abstractmethod
    def __call__(self, instance: IndexedInstance) -> IndexedInstance:
        """
        Takes a raw `IndexedInstance`, performs some processing on it, and returns
        an `IndexedInstance` again. Typically, you should return an instance having
        the fields specified in
        `trapper.models.auto_wrappers._TASK_TO_INPUT_FIELDS[task]` (except for
        "attention_mask" since it will be generated by the DataCollator for you)
        where `task` is your choice of modeling the downstream task you want to
        solve using `transformers`. Each task in the `_TASK_TO_INPUT_FIELDS`
        corresponds to an `AutoModel...` class implemented in `transformers`.

        You may look at
        `DataAdapterForQuestionAnswering.__call__` for an example.

        Args:
            instance ():
        """
        raise NotImplementedError
