import logging
from functools import partial
from typing import Optional, cast
from collections.abc import Iterable, Callable

import torch
from pytorch_lightning import Trainer
from torch import Tensor, sigmoid, softmax
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForTokenClassification,
    AutoConfig,
    AutoTokenizer,
    DataCollatorWithPadding,
    BatchEncoding,
    PreTrainedTokenizerBase,
)
from transformers.file_utils import PaddingStrategy

from kazu.data.data import Document, Section
from kazu.data.pytorch import HFDataset
from kazu.steps import Step, document_batch_step
from kazu.steps.ner.hf_lightning_wrappers import PLAutoModelForTokenClassification
from kazu.steps.ner.tokenized_word_processor import TokenizedWordProcessor, TokenizedWord
from kazu.steps.ner.entity_post_processing import NonContiguousEntitySplitter
from kazu.utils.utils import documents_to_document_section_batch_encodings_map

logger = logging.getLogger(__name__)


class TransformersModelForTokenClassificationNerStep(Step):
    """
    A wrapper for
    :class:`transformers.AutoModelForTokenClassification`\\ . This implementation uses a sliding
    window concept to process large documents that don't fit into the maximum sequence length
    allowed by a model. Resulting token labels are then post processed by
    :class:`~kazu.steps.ner.tokenized_word_processor.TokenizedWordProcessor`.
    """

    def __init__(
        self,
        path: str,
        batch_size: int,
        stride: int,
        max_sequence_length: int,
        trainer: Trainer,
        labels: list[str],
        detect_subspans: bool = False,
        threshold: Optional[float] = None,
        entity_splitter: Optional[NonContiguousEntitySplitter] = None,
        strip_re: Optional[dict[str, str]] = None,
    ):
        """

        :param path: path to HF model, config and tokenizer. Passed to HF .from_pretrained()
        :param batch_size: batch size for dataloader
        :param stride: passed to HF tokenizers (for splitting long docs)
        :param max_sequence_length: passed to HF tokenizers (for splitting long docs)
        :param trainer:
        :param detect_subspans: attempt to detect nested entities (threshold must be configured)
        :param threshold: the confidence threshold used to detect nested entities
        :param entity_splitter: to detect non-contiguous entities if provided
        :param strip_re: passed to :class:`~kazu.steps.ner.tokenized_word_processor.TokenizedWordProcessor`
        """
        self.entity_splitter = entity_splitter
        if max_sequence_length % 2 != 0:
            raise RuntimeError(
                "max_sequence_length must %2 ==0 in order for correct document windowing"
            )
        self.max_sequence_length = max_sequence_length
        if stride % 2 != 0:
            raise RuntimeError("stride must %2 ==0 in order for correct document windowing")
        self.stride = stride
        self.batch_size = batch_size
        self.config = AutoConfig.from_pretrained(path)
        self.tokeniser: PreTrainedTokenizerBase = AutoTokenizer.from_pretrained(
            path, config=self.config
        )
        self.model = AutoModelForTokenClassification.from_pretrained(path, config=self.config)
        self.model = PLAutoModelForTokenClassification(self.model).eval()
        self.trainer = trainer
        self.activation_fn = cast(
            Callable[[Tensor], Tensor], sigmoid if detect_subspans else partial(softmax, dim=-1)
        )
        self.tokenized_word_processor = TokenizedWordProcessor(
            detect_subspans=detect_subspans,
            confidence_threshold=threshold,
            id2label=self.id2labels_from_label_list(labels),
            strip_re=strip_re,
        )

    @document_batch_step
    def __call__(self, docs: list[Document]) -> None:
        loader, id_section_map = self.get_dataloader(docs)
        # need this so mypy knows to expect the dataset to have encodings
        dataset = cast(HFDataset, loader.dataset)
        # run the transformer and get results
        activations = self.get_activations(loader)
        for section_index, section in id_section_map.items():
            words = self.section_frames_to_tokenised_words(
                section_index=section_index,
                batch_encoding=dataset.encodings,
                predictions=activations,
            )
            entities = self.tokenized_word_processor(
                words, text=section.text, namespace=self.namespace()
            )
            section.entities.extend(entities)
            if self.entity_splitter:
                for ent in entities:
                    section.entities.extend(self.entity_splitter(ent, section.text))

    def get_activations(self, loader: DataLoader) -> Tensor:
        """get a namedtuple_values_indices consisting of confidence and labels
        for a given dataloader (i.e. run bert)

        :param loader:
        :return:
        """
        results = torch.cat(
            [
                x.logits  # type: ignore[union-attr]
                for x in self.trainer.predict(  # type: ignore[union-attr]
                    model=self.model, dataloaders=loader, return_predictions=True
                )
            ]
        )
        return self.activation_fn(results)

    def get_dataloader(self, docs: list[Document]) -> tuple[DataLoader, dict[int, Section]]:
        """
        get a dataloader from a List of :class:`kazu.data.data.Document`. Collation is handled via
        :class:`transformers.DataCollatorWithPadding`\\ .

        :param docs:
        :return: The returned dict's keys map to overflow_to_sample_mapping in the
                underlying batch encoding, allowing the processing of docs longer than can fit within the maximum
                sequence length of a transformer
        """
        batch_encoding, id_section_map = documents_to_document_section_batch_encodings_map(
            docs, self.tokeniser, stride=self.stride, max_length=self.max_sequence_length
        )
        dataset = HFDataset(batch_encoding)
        collate_func = DataCollatorWithPadding(
            tokenizer=self.tokeniser, padding=PaddingStrategy.MAX_LENGTH
        )
        loader = DataLoader(dataset=dataset, batch_size=self.batch_size, collate_fn=collate_func)
        return loader, id_section_map

    def frame_to_tok_word(
        self,
        batch_encoding: BatchEncoding,
        number_of_frames: int,
        frame_index: int,
        section_frame_index: int,
        predictions: Tensor,
    ) -> list[TokenizedWord]:
        """Depending on the number of frames generated by a string of text, and
        whether it is the first or last frame, we need to return different
        subsets of the frame offsets and frame word_ids.

        :param batch_encoding:
        :param number_of_frames: number of frames created by the tokenizer for the string
        :param frame_index: the index of the query frame, relative to the total number of frames
        :param section_frame_index: the index of the section frame, relative to the whole BatchEncoding
        :param predictions:
        :return: Tuple of 2 lists: frame offsets and frame word ids
        """
        half_stride = int(self.stride / 2)
        # 1:-1 skip cls and sep tokens
        if number_of_frames == 1:
            start_index = 1
            end_index = -1
        elif number_of_frames > 1 and frame_index == 0:
            start_index = 1
            end_index = -(half_stride + 1)
        elif number_of_frames > 1 and frame_index == number_of_frames - 1:
            start_index = half_stride + 1
            end_index = -1
        else:
            start_index = half_stride + 1
            end_index = -(half_stride + 1)

        assert batch_encoding.encodings is not None
        frame_offsets = batch_encoding.encodings[section_frame_index].offsets[start_index:end_index]
        frame_word_ids: list[Optional[int]] = batch_encoding.encodings[
            section_frame_index
        ].word_ids[start_index:end_index]
        frame_token_ids = batch_encoding.encodings[section_frame_index].ids[start_index:end_index]
        frame_tokens = batch_encoding.encodings[section_frame_index].tokens[start_index:end_index]
        predictions = predictions[section_frame_index][start_index:end_index]
        prev_word_id: Optional[int] = None
        all_words = []

        word_id_index_start, offset_start, offset_end = 0, 0, 0

        for i, word_id in enumerate(frame_word_ids):
            if word_id != prev_word_id:
                # start a new word, and add the previous
                if prev_word_id is not None:
                    all_words.append(
                        TokenizedWord(
                            token_offsets=frame_offsets[word_id_index_start:i],
                            token_confidences=predictions[word_id_index_start:i],
                            token_ids=frame_token_ids[word_id_index_start:i],
                            tokens=frame_tokens[word_id_index_start:i],
                            word_char_start=offset_start,
                            word_char_end=offset_end - 1,
                            word_id=prev_word_id,
                        )
                    )
                word_id_index_start = i
                offset_start, offset_end = frame_offsets[i]
            if i == len(frame_word_ids) - 1 and word_id is not None:
                # if checking the last word in a frame, if word_id is not None, add it
                all_words.append(
                    TokenizedWord(
                        token_offsets=frame_offsets[word_id_index_start : i + 1],
                        token_confidences=predictions[word_id_index_start : i + 1],
                        token_ids=frame_token_ids[word_id_index_start : i + 1],
                        tokens=frame_tokens[word_id_index_start : i + 1],
                        word_char_start=offset_start,
                        word_char_end=offset_end,
                        word_id=word_id,
                    )
                )

            _, offset_end = frame_offsets[i]
            prev_word_id = word_id

        frame_word_ids_excluding_nones = set(frame_word_ids)
        frame_word_ids_excluding_nones.discard(None)
        assert len(set(frame_word_ids_excluding_nones)) == len(all_words)

        logger.debug(
            f"inputs this frame: {self.tokeniser.decode(frame_token_ids[:start_index])}<-IGNORED "
            f"START->{self.tokeniser.decode(frame_token_ids[start_index:end_index])}"
            f"<-END IGNORED->{self.tokeniser.decode(frame_token_ids[end_index:])}"
        )
        return all_words

    def section_frames_to_tokenised_words(
        self,
        section_index: int,
        batch_encoding: BatchEncoding,
        predictions: Tensor,
    ) -> list[TokenizedWord]:
        """

        :param section_index:
        :param batch_encoding:
        :param predictions:
        :return:
        """
        words = []

        section_frame_indices = self.get_list_of_batch_encoding_frames_for_section(
            batch_encoding, section_index
        )
        for frame_index, section_frame_index in enumerate(section_frame_indices):
            word_sub_list = self.frame_to_tok_word(
                batch_encoding=batch_encoding,
                number_of_frames=len(section_frame_indices),
                frame_index=frame_index,
                section_frame_index=section_frame_index,
                predictions=predictions,
            )
            words.extend(word_sub_list)
        return words

    @staticmethod
    def get_list_of_batch_encoding_frames_for_section(
        batch_encoding: BatchEncoding, section_index: int
    ) -> list[int]:
        """For a given dataloader with a HFDataset, return a list of frame
        indexes associated with a given section index.

        :param batch_encoding:
        :param section_index:
        :return:
        """
        section_frame_indices = [
            i
            for i, x in enumerate(batch_encoding.data["overflow_to_sample_mapping"])
            if x == section_index
        ]
        return section_frame_indices

    @staticmethod
    def id2labels_from_label_list(labels: Iterable[str]) -> dict[int, str]:
        return {idx: label for idx, label in enumerate(labels)}
