from functools import lru_cache
from typing import List, NamedTuple, Optional, Dict, Tuple, Union

import torch
from tokenizers import AddedToken

from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast

from src.utils.collections import flatten


class TokenizationOutput(NamedTuple):
    input_ids: torch.Tensor
    attention_mask: torch.Tensor
    token_type_ids: Optional[torch.Tensor]
    definitions_offsets: Dict[str, Tuple[int, int]]
    rel_infos_offsets: Optional[Dict[str, Tuple[int, int]]] = None
    relative_positions: Optional[torch.Tensor] = None
    padding_attention_mask: Optional[torch.Tensor] = None


# utilities to extract the gold definition indices in the input sequence
def extract_gold_indices(definition: str, definitions_sequence: str, offsets_mapping: List[Tuple[int, int]]) -> Tuple[int, int]:
    gold_start_offset = definitions_sequence.index(definition)
    gold_end_offset = gold_start_offset + len(definition)

    offset_start2bpe_idx = {}
    offset_end2bpe_idx = {}
    for i, (off_start, off_end) in enumerate(offsets_mapping):

        if off_start == off_end:
            continue  # specials

        offset_start2bpe_idx[off_start] = i
        offset_end2bpe_idx[off_end] = i

    start_bpe_idx = offset_start2bpe_idx[gold_start_offset]
    end_bpe_idx = offset_end2bpe_idx[gold_end_offset]

    assert start_bpe_idx < end_bpe_idx

    return start_bpe_idx, end_bpe_idx


class ConsecTokenizer:
    def __init__(
        self,
        transformer_model: Union[str, PreTrainedTokenizer, PreTrainedTokenizerFast],
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
    ):
        if type(transformer_model) == str:
            self.tokenizer = AutoTokenizer.from_pretrained(transformer_model)
        else:
            self.tokenizer = transformer_model
        self.target_marker = target_marker
        self.context_definitions_token = context_definitions_token
        self.context_markers = [
            (
                context_markers["pattern"][0].replace("#I#", f"{i}"),
                context_markers["pattern"][1].replace("#I#", f"{i}"),
            )
            for i in range(context_markers["number"])
        ]
        assert (
            len(set(self.context_markers)) == context_markers["number"]
        ), f"Error in given pattern: number of unique created patterns != specified number"

        additional_special_tokens = [
            *[AddedToken(t, single_word=True, lstrip=True) for t in self.target_marker],
            *[AddedToken(t, single_word=True, lstrip=True) for p in self.context_markers for t in p],
            AddedToken(context_definitions_token, single_word=True, lstrip=True),
        ]
        self.tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})
        self.n_added_special_tokens = len(additional_special_tokens)
        self.add_prefix_space = add_prefix_space

    def mark_token(self, token: str, marker: Tuple[str, str]) -> str:
        bom, eom = marker
        if self.add_prefix_space:
            return f"{bom} {token} {eom}"
        else:
            return f"{bom}{token}{eom}"

    def tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[Tuple[str, str]],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:
        raise NotImplementedError

    @property
    def pad_token_id(self) -> int:
        return self.tokenizer.pad_token_id

    @property
    def model_max_length(self) -> int:
        return self.tokenizer.model_max_length


class DeBERTaTokenizer(ConsecTokenizer):
    def __init__(
        self,
        transformer_model: str,
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
        optimize_relative_positions: bool = True,
        enforce_symmetry: bool = True,
    ):
        super().__init__(
            transformer_model,
            target_marker,
            context_definitions_token,
            context_markers,
            add_prefix_space,
        )
        self.optimize_relative_positions = optimize_relative_positions
        self.enforce_symmetry = enforce_symmetry

    def tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[Tuple[str, str]],
        context_definitions2positions: List[Tuple[str, int]],
        **kwargs,
    ) -> TokenizationOutput:
        if self.optimize_relative_positions:
            return self.power_tokenize(
                sentence,
                instance_idx,
                instance_possible_definitions,
                context_definitions2positions,
            )
        else:
            raise NotImplementedError

    def deberta_tokenize(self, text: str) -> List[int]:
        return self.tokenizer(text, return_attention_mask=False, return_token_type_ids=False, add_special_tokens=True,)[
            "input_ids"
        ][1:-1]

    def plain_tokenize(
        self,
        sentence: Union[str, List[int]],
        instance_possible_definitions: Union[List[str], List[Tuple[str, List[int]]]],
        context_definitions: Union[List[str], List[Tuple[str, List[int]]]],
        use_specials: bool = True,
    ) -> TokenizationOutput:
        if type(sentence) != list:
            sentence = f" {sentence}"
            sentence_input_ids = self.deberta_tokenize(sentence)
        else:
            sentence_input_ids = sentence

        final_input_ids = [self.tokenizer.cls_token_id] + sentence_input_ids + [self.tokenizer.sep_token_id]
        token_type_ids = [0] * len(final_input_ids)

        definitions_offsets = dict()
        rel_infos_offsets = dict()
        for definition in instance_possible_definitions:
            if type(definition) == tuple:
                (definition, rel_infos), definition_ids, rel_infos_ids = definition
            else:
                # definition_ids = self.deberta_tokenize(f" {definition}")
                raise NotImplementedError

            definitions_offsets[definition] = (
                len(final_input_ids),
                len(final_input_ids) + len(definition_ids) + len(rel_infos_ids),
            )
            rel_infos_offsets[definition] = (
                len(final_input_ids) + len(definition_ids),
                len(final_input_ids) + len(definition_ids) + len(rel_infos_ids),
            )

            final_input_ids += definition_ids + rel_infos_ids
            token_type_ids += [1] * len(definition_ids + rel_infos_ids)

        # "context" definitions token
        if self.context_definitions_token is not None and use_specials:
            final_input_ids += self.deberta_tokenize(self.context_definitions_token)
            token_type_ids.append(1)
        else:
            if use_specials:
                assert len(context_definitions) == 0

        for definition in context_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.deberta_tokenize(f" {definition}")

            final_input_ids += definition_ids
            token_type_ids += [1] * len(definition_ids)

        # last [SEP] token
        final_input_ids.append(self.tokenizer.sep_token_id)
        token_type_ids.append(1)

        final_input_ids = torch.tensor(final_input_ids, dtype=torch.long)
        attention_mask = torch.ones_like(final_input_ids)
        token_type_ids = torch.tensor(token_type_ids, dtype=torch.long)

        return TokenizationOutput(
            final_input_ids,
            attention_mask,
            token_type_ids,
            definitions_offsets,
            rel_infos_offsets,
        )

    @lru_cache(maxsize=10_000)
    def split_tokenize(self, word: str):
        return self.deberta_tokenize(f" {word.strip()}")

    @lru_cache(maxsize=10_000)
    def _tokenize_sentence_list(self, sentence: str) -> Tuple[List[int], List[Tuple[int, int]]]:
        sentence = sentence.split("\t")
        sentence_input_ids = [self.tokenizer.cls_token_id]
        token_offsets = []

        for token in sentence:
            token_ids = self.split_tokenize(f" {token.strip()}")
            token_offsets.append((len(sentence_input_ids), len(sentence_input_ids) + len(token_ids)))
            sentence_input_ids += token_ids

        sentence_input_ids.append(self.tokenizer.sep_token_id)

        return sentence_input_ids, token_offsets

    @staticmethod
    def mirror_arange(size: int, zero_pos: int):
        return torch.cat(
            [
                torch.flip(torch.arange(0, zero_pos + 1), dims=(0,)),
                -torch.arange(1, size - zero_pos),
            ]
        )

    def power_tokenize(
        self,
        sentence: List[str],  # tokens can be multiword
        instance_idx: int,
        instance_possible_definitions: List[Tuple[str, str]],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:

        sentence_input_ids, token_offsets = self._tokenize_sentence_list("\t".join(sentence))

        tmp_instance_possible_definitions_ids = [
            (self.deberta_tokenize(f" {ipd}"), self.deberta_tokenize(f" {ri}")) for ipd, ri in instance_possible_definitions
        ]
        instance_possible_definitions_ids = [item[0] + item[1] for item in tmp_instance_possible_definitions_ids]
        context_definitions_ids = [self.deberta_tokenize(f" {cd}") for cd, _ in context_definitions2positions]

        total_input_ids = len(sentence_input_ids) + len(flatten(instance_possible_definitions_ids)) + len(flatten(context_definitions_ids))

        relative_positions = torch.zeros((total_input_ids + 1, total_input_ids + 1), dtype=torch.long)

        # filling the sentences relative positions
        for id_idx in range(len(sentence_input_ids)):
            relative_positions[id_idx, :-1] = torch.cat(
                [self.mirror_arange(len(sentence_input_ids), id_idx)]
                + [
                    -torch.arange(
                        len(sentence_input_ids) - id_idx,
                        len(sentence_input_ids) + len(def_ids) - id_idx,
                    )
                    for def_ids in flatten(tmp_instance_possible_definitions_ids) + context_definitions_ids
                ]
            )

        # flipping the matrix
        relative_positions.T[: len(sentence_input_ids)] = -relative_positions[: len(sentence_input_ids)]

        curr_offset = len(sentence_input_ids)

        # computing the relative positions for the possible definitions
        definitions2positions = []
        for ipdid, riid in tmp_instance_possible_definitions_ids:
            definitions2positions.append((instance_idx, ipdid, True))
            definitions2positions.append((instance_idx, riid, False))
        for cdid, (_, rel_idx) in zip(context_definitions_ids, context_definitions2positions):
            definitions2positions.append((rel_idx, cdid, True))
        for def_num in range(len(definitions2positions)):

            rel_token_idx, def_ids, closing_tgt_word = definitions2positions[def_num]
            possible_def_token_pos = list(range(len(def_ids)))
            instance_start_pos, instance_end_pos = token_offsets[rel_token_idx]

            for token_pos in possible_def_token_pos:

                if closing_tgt_word:
                    for off_idx, inst_token_pos in enumerate(range(instance_start_pos, instance_end_pos)):
                        relative_positions[curr_offset + token_pos, inst_token_pos] = token_pos + (instance_end_pos - instance_start_pos) - off_idx
                        if self.enforce_symmetry or rel_token_idx != instance_idx:
                            relative_positions[inst_token_pos, curr_offset + token_pos] = -relative_positions[curr_offset + token_pos, inst_token_pos]

                relative_positions[curr_offset + token_pos, len(sentence_input_ids) : -1] = self.mirror_arange(
                    total_input_ids - len(sentence_input_ids),
                    curr_offset + token_pos - len(sentence_input_ids),
                )

            curr_offset += len(possible_def_token_pos)

        last_token_positions = torch.min(relative_positions, dim=-1)[0] - 1
        relative_positions[-1] = -last_token_positions
        relative_positions.T[-1] = last_token_positions
        relative_positions[-1, -1] = 0

        relative_positions[relative_positions == -0] = 0

        tokenization_output = self.plain_tokenize(
            sentence_input_ids[1:-1],
            [(x, y[0], y[1]) for x, y in zip(instance_possible_definitions, tmp_instance_possible_definitions_ids)],
            [
                (x, y)
                for x, y in zip(
                    [x[0] for x in context_definitions2positions],
                    context_definitions_ids,
                )
            ],
            use_specials=False,
        )

        assert tokenization_output.input_ids.shape[0] == relative_positions.shape[0]

        attention_mask = tokenization_output.attention_mask
        attention_mask = attention_mask.unsqueeze(-1).repeat(1, attention_mask.size(0))
        padding_attention_mask = tokenization_output.attention_mask

        # fix attention-mask and relative-positions
        cls_sep_positions = torch.tensor(
            [
                token_pos
                for token_pos, token_id in enumerate(tokenization_output.input_ids)
                if token_id in [self.tokenizer.cls_token_id, self.tokenizer.sep_token_id]
            ]
        )

        for (_, (rel_start_idx, rel_end_idx)), (_, (def_start_idx, _)) in zip(
            tokenization_output.rel_infos_offsets.items(),
            tokenization_output.definitions_offsets.items(),
        ):

            attention_mask[..., rel_start_idx:rel_end_idx] = 0  # all -x-> rel_infos
            attention_mask[def_start_idx:rel_start_idx, rel_start_idx:rel_end_idx] = 1  # def(all) --> rel(all)
            attention_mask[cls_sep_positions, rel_start_idx:rel_end_idx] = 1  # cls, sep --> rel_infos

            attention_mask[rel_start_idx:rel_end_idx, len(sentence_input_ids) :] = 0  # rel_infos -x-> all (対象文以外)
            attention_mask[rel_start_idx:rel_end_idx, def_start_idx:rel_start_idx] = 1  # rel(all) --> def(all)
            attention_mask[rel_start_idx:rel_end_idx, rel_start_idx:rel_end_idx] = 1  # rel_infos --> rel_infos
            attention_mask[rel_start_idx:rel_end_idx, cls_sep_positions] = 1  # rel_infos --> cls, sep

            # relative-position
            relative_positions[def_start_idx, rel_start_idx] = -1
            relative_positions[rel_start_idx, def_start_idx] = -1

        return TokenizationOutput(
            tokenization_output.input_ids,
            attention_mask,
            tokenization_output.token_type_ids,
            tokenization_output.definitions_offsets,
            relative_positions=relative_positions,
            padding_attention_mask=padding_attention_mask,
        )

    @property
    def model_max_length(self) -> int:
        return 24_528  # from paper: deberta-large


class MBartTokenizer(ConsecTokenizer):
    def __init__(
        self,
        transformer_model: str,
        target_marker: Tuple[str, str],
        context_definitions_token: str,
        context_markers: Dict,
        add_prefix_space: bool,
        source_language: str = "en_XX",
        target_language: str = "en_EN",
    ):
        tokenizer = AutoTokenizer.from_pretrained(transformer_model, src_lang=source_language, tgt_lang=target_language)
        self.source_language = source_language
        super().__init__(
            tokenizer,
            target_marker,
            context_definitions_token,
            context_markers,
            add_prefix_space,
        )

    def mbart_tokenize(self, text: str) -> List[int]:
        tokenization_out = self.tokenizer(
            text,
            return_attention_mask=False,
            return_token_type_ids=False,
            add_special_tokens=True,
        )["input_ids"]
        return tokenization_out[:-2]

    def tokenize(
        self,
        sentence: Union[str, List[int]],
        instance_idx: int,
        instance_possible_definitions: List[str],
        context_definitions2positions: List[Tuple[str, int]],
    ) -> TokenizationOutput:
        use_specials = False
        context_definitions = [x[0] for x in context_definitions2positions]

        sentence = f" {' '.join(sentence)}"
        sentence_input_ids = self.mbart_tokenize(sentence)

        final_input_ids = sentence_input_ids

        definitions_offsets = dict()
        for definition in instance_possible_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.mbart_tokenize(f" {definition}")

            definitions_offsets[definition] = (
                len(final_input_ids),
                len(final_input_ids) + len(definition_ids),
            )

            final_input_ids += definition_ids

        # "context" definitions token
        if self.context_definitions_token is not None and use_specials:
            final_input_ids += self.mbart_tokenize(self.context_definitions_token)
        else:
            if use_specials:
                assert len(context_definitions) == 0

        for definition in context_definitions:
            if type(definition) == tuple:
                definition, definition_ids = definition
            else:
                definition_ids = self.mbart_tokenize(f" {definition}")

            final_input_ids += definition_ids

        # last [SEP] token
        final_input_ids.append(self.tokenizer.sep_token_id)

        # lang id
        final_input_ids += self.mbart_tokenize(self.source_language)

        final_input_ids = torch.tensor(final_input_ids, dtype=torch.long)
        attention_mask = torch.ones_like(final_input_ids)

        return TokenizationOutput(final_input_ids, attention_mask, None, definitions_offsets)
