import itertools
import random
import re
from collections.abc import Sequence, Generator
from typing import Callable

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader, Subset
from transformers import PreTrainedTokenizerFast

from matchescu.matching.ml.ditto._ref_adapter import to_text
from matchescu.reference_store.comparison_space import BinaryComparisonSpace
from matchescu.reference_store.id_table import IdTable
from matchescu.typing import EntityReferenceIdentifier


def alphanumeric(token):
    return "".join([ch if ch.isalnum() else " " for ch in token])


class Augmenter(object):
    """Data augmentation operator.

    Support both span and attribute level augmentation operators.
    """

    def __init__(self):
        self._augmentations: dict[
            str, Callable[[Sequence, Sequence], tuple[Sequence, Sequence]]
        ] = {
            "del": self._del,
            "swap": self._swap,
            "drop_len": self._drop_len,
            "drop_sym": self._drop_symbols,
            "drop_same": self._drop_same,
            "drop_token": self._drop_token,
            "ins": self._insert,
            "append_col": self._append_col,
            "drop_col": self._drop_col,
        }

    def augment(self, tokens: Sequence[str], labels: Sequence[str], op="del"):
        augmentation = self._augmentations.get(op, lambda x, y: (x, y))
        return augmentation(labels, tokens)

    @staticmethod
    def _drop_col(labels, tokens):
        col_starts = [i for i in range(len(tokens)) if tokens[i] == "COL"]
        col_ends = [0] * len(col_starts)
        col_lens = [0] * len(col_starts)
        for i, pos in enumerate(col_starts):
            if i == len(col_starts) - 1:
                col_lens[i] = len(tokens) - pos
                col_ends[i] = len(tokens) - 1
            else:
                col_lens[i] = col_starts[i + 1] - pos
                col_ends[i] = col_starts[i + 1] - 1

            if tokens[col_ends[i]] == "[SEP]":
                col_ends[i] -= 1
                col_lens[i] -= 1
        candidates = [i for i, le in enumerate(col_lens) if le <= 8]
        if len(candidates) > 0:
            idx = random.choice(candidates)
            start, end = col_starts[idx], col_ends[idx]
            new_tokens = tokens[:start] + tokens[end + 1 :]
            new_labels = labels[:start] + labels[end + 1 :]
        else:
            new_tokens, new_labels = tokens, labels
        return new_labels, new_tokens

    @staticmethod
    def _append_col(labels, tokens):
        col_starts = [i for i in range(len(tokens)) if tokens[i] == "COL"]
        col_ends = [0] * len(col_starts)
        col_lens = [0] * len(col_starts)
        for i, pos in enumerate(col_starts):
            if i == len(col_starts) - 1:
                col_lens[i] = len(tokens) - pos
                col_ends[i] = len(tokens) - 1
            else:
                col_lens[i] = col_starts[i + 1] - pos
                col_ends[i] = col_starts[i + 1] - 1

            if tokens[col_ends[i]] == "[SEP]":
                col_ends[i] -= 1
                col_lens[i] -= 1
                break
        candidates = [i for i, le in enumerate(col_lens) if le > 0]
        if len(candidates) < 2:
            return tokens, labels
        else:
            idx1, idx2 = random.sample(candidates, k=2)
            start1, end1 = col_starts[idx1], col_ends[idx1]
            sub_tokens = tokens[start1 : end1 + 1]
            sub_labels = labels[start1 : end1 + 1]
            val_pos = 0
            for i, token in enumerate(sub_tokens):
                if token == "VAL":
                    val_pos = i + 1
                    break
            sub_tokens = sub_tokens[val_pos:]
            sub_labels = sub_labels[val_pos:]

            end2 = col_ends[idx2]
            new_tokens = []
            new_labels = []
            for i in range(len(tokens)):
                if start1 <= i <= end1:
                    continue
                new_tokens.append(tokens[i])
                new_labels.append(labels[i])
                if i == end2:
                    new_tokens += sub_tokens
                    new_labels += sub_labels
            return new_tokens, new_labels

    def _insert(self, labels, tokens):
        pos = self.sample_position(labels)
        symbol = random.choice("-*.,#&")
        new_tokens = tokens[:pos] + [symbol] + tokens[pos:]
        new_labels = labels[:pos] + ["O"] + labels[pos:]
        return new_tokens, new_labels

    @staticmethod
    def _drop_token(labels, tokens):
        new_tokens, new_labels = [], []
        for token, label in zip(tokens, labels):
            if label != "O" or random.randint(0, 4) != 0:
                new_tokens.append(token)
                new_labels.append(label)
        return new_labels, new_tokens

    @staticmethod
    def _drop_same(labels, tokens):
        left_token = set([])
        right_token = set([])
        left = True
        for token, label in zip(tokens, labels):
            if label == "O":
                token = token.lower()
                if left:
                    left_token.add(token)
                else:
                    right_token.add(token)
            if token == "[SEP]":
                left = False
        same = left_token & right_token
        targets = random.choices(list(same), k=1)
        new_tokens, new_labels = [], []
        for token, label in zip(tokens, labels):
            if token.lower() not in targets or label != "O":
                new_tokens.append(token)
                new_labels.append(label)
        return new_labels, new_tokens

    @staticmethod
    def _drop_symbols(labels, tokens):
        dropped_tokens = [alphanumeric(token) for token in tokens]
        new_tokens = []
        new_labels = []
        for token, d_token, label in zip(tokens, dropped_tokens, labels):
            if random.randint(0, 4) != 0 or label != "O":
                new_tokens.append(token)
                new_labels.append(label)
            else:
                if d_token != "":
                    new_tokens.append(d_token)
                    new_labels.append(label)
        return new_labels, new_tokens

    def _swap(self, labels, tokens):
        span_len = random.randint(2, 4)
        pos1, pos2 = self.sample_span(tokens, labels, span_len=span_len)
        if pos1 < 0:
            return tokens, labels
        sub_arr = tokens[pos1 : pos2 + 1]
        random.shuffle(sub_arr)
        new_tokens = tokens[:pos1] + sub_arr + tokens[pos2 + 1 :]
        new_labels = tokens[:pos1] + ["O"] * (pos2 - pos1 + 1) + labels[pos2 + 1 :]
        return new_tokens, new_labels

    def _del(self, labels, tokens):
        span_len = random.randint(1, 2)
        pos1, pos2 = self.sample_span(tokens, labels, span_len=span_len)
        if pos1 < 0:
            return tokens, labels
        new_tokens = tokens[:pos1] + tokens[pos2 + 1 :]
        new_labels = tokens[:pos1] + labels[pos2 + 1 :]
        return new_tokens, new_labels

    @staticmethod
    def _drop_len(labels, tokens):
        all_lens = [len(token) for token, label in zip(tokens, labels) if label == "O"]
        if len(all_lens) == 0:
            return tokens, labels
        target_lens = random.choices(all_lens, k=1)
        new_tokens = []
        new_labels = []

        for token, label in zip(tokens, labels):
            if label != "O" or len(token) not in target_lens:
                new_tokens.append(token)
                new_labels.append(label)
        return new_tokens, new_labels

    def augment_sent(self, text, op="all"):
        """Performs data augmentation on a classification example.

        Similar to augment(tokens, labels) but works for sentences
        or sentence-pairs.

        Args:
            text (str): the input sentence
            op (str, optional): a string encoding of the operator to be applied

        Returns:
            str: the augmented sentence
        """
        # 50% of chance of flipping
        if " [SEP] " in text and random.randint(0, 1) == 0:
            left, right = text.split(" [SEP] ")
            text = right + " [SEP] " + left

        # tokenize the sentence
        tokens = text.split(" ")

        # avoid the special tokens
        labels = []
        for token in tokens:
            if token in ["COL", "VAL"]:
                labels.append("HD")
            elif token in ["[CLS]", "[SEP]"]:
                labels.append("<SEP>")
            else:
                labels.append("O")

        if op == "all":
            # RandAugment: https://arxiv.org/pdf/1909.13719.pdf
            N = 3
            ops = ["del", "swap", "drop_col", "append_col"]
            for op in random.choices(ops, k=N):
                tokens, labels = self.augment(tokens, labels, op=op)
        else:
            tokens, labels = self.augment(tokens, labels, op=op)
        results = " ".join(tokens)
        return results

    @staticmethod
    def _process_index(
        idx: int, span_len: int, labels: Sequence[str]
    ) -> Generator[tuple[int, int]]:
        span_end = idx + span_len
        if span_end - 1 >= len(labels):
            return
        span_text = "".join(labels[idx:span_end])
        pattern = "^O{%d}$" % span_len
        if not re.match(pattern, span_text, re.S):
            return
        yield idx, span_end - 1

    def _enumerate_spans(
        self, tokens: Sequence[str], labels: Sequence[str], span_len: int = 3
    ) -> Generator[tuple[int, int]]:
        yield from (
            t
            for idx, _ in enumerate(tokens)
            for t in self._process_index(idx, span_len, labels)
        )

    def sample_span(
        self, tokens: Sequence[str], labels: Sequence[str], span_len: int = 3
    ) -> tuple[int, int]:
        candidates = list(self._enumerate_spans(tokens, labels, span_len))
        if len(candidates) == 0:
            return -1, -1
        return random.choice(candidates)

    @staticmethod
    def sample_position(labels):
        candidates = list(
            map(lambda x: x[0], filter(lambda x: x[1] == "O", enumerate(labels)))
        )
        return random.choice(candidates) if len(candidates) > 0 else -1


class DittoDataset(Dataset):
    """EM dataset"""

    def __init__(
        self,
        comparison_space: BinaryComparisonSpace,
        id_table: IdTable,
        ground_truth: set[tuple[EntityReferenceIdentifier, EntityReferenceIdentifier]],
        tokenizer: PreTrainedTokenizerFast,
        max_len=256,
        size=None,
        augmentations=None,
        left_cols: tuple | None = None,
        right_cols: tuple | None = None,
    ):
        self.__tokenizer = tokenizer
        id_pairs = list(itertools.islice(comparison_space, size))
        self.__pairs = list(tuple(map(id_table.get_all, id_pairs)))
        self.__labels = list(map(lambda pair: int(pair in ground_truth), id_pairs))
        self.__max_len = max_len
        self.__size = size
        self.__left_cols = left_cols
        self.__right_cols = right_cols
        self.da = augmentations
        if augmentations is not None:
            self.augmenter = Augmenter()
        else:
            self.augmenter = None

    def __len__(self):
        """Return the size of the dataset."""
        return len(self.__pairs)

    def __getitem__(self, idx):
        """Return a tokenized item of the dataset.

        Args:
            idx (int): the index of the item

        Returns:
            List of int: token ID's of the two entities
            List of int: token ID's of the two entities augmented (if da is set)
            int: the label of the pair (0: unmatch, 1: match)
        """
        left, right = self.__pairs[idx]
        left_text = to_text(left, self.__left_cols)
        right_text = to_text(right, self.__right_cols)
        x = self.__tokenizer.encode(
            text=left_text,
            text_pair=right_text,
            max_length=self.__max_len,
            truncation=True,
        )

        # augment if da is set
        if self.da is not None:
            augmented = self.augmenter.augment_sent(
                f"{left_text} [SEP] {right_text}", self.da
            )
            left_aug, right_aug = augmented.split(" [SEP] ")
            x_aug = self.__tokenizer.encode(
                text=left_aug,
                text_pair=right_aug,
                max_length=self.__max_len,
                truncation=True,
            )
            return x, x_aug, self.__labels[idx]
        else:
            return x, self.__labels[idx]

    @staticmethod
    def __pad(x: Sequence, total_length: int) -> torch.LongTensor:
        tensor_data = list(
            map(
                lambda vec: list(
                    itertools.chain(
                        vec, itertools.repeat(0, max(total_length - len(vec), 0))
                    )
                ),
                x,
            )
        )
        return torch.LongTensor(tensor_data)

    @staticmethod
    def __zero_padded(batch: list[tuple]) -> tuple[torch.LongTensor, ...]:
        if len(batch[0]) == 3:
            x1, x2, y = zip(*batch)

            n = max(map(len, x1 + x2))
            x1 = DittoDataset.__pad(x1, n)
            x2 = DittoDataset.__pad(x2, n)
            return x1, x2, torch.LongTensor(y)
        else:
            x, y = zip(*batch)
            n = max(map(len, x))
            x = DittoDataset.__pad(x, n)
            return x, torch.LongTensor(y)

    def split(
        self,
        train_proportion: float,
        xv_proportion: float,
        test_proportion: float,
        batch_size: int = 32,
    ) -> tuple[DataLoader, DataLoader, DataLoader]:
        proportion_total = train_proportion + xv_proportion + test_proportion
        train_proportion /= proportion_total
        xv_proportion /= proportion_total
        test_proportion = 1 - train_proportion - xv_proportion
        assert train_proportion + xv_proportion + test_proportion == 1
        total_size = len(self.__pairs)

        indices = np.arange(len(self.__pairs))
        np.random.shuffle(indices)

        train_size = int(train_proportion * total_size)
        xv_size = int(xv_proportion * total_size)

        # Step 5: Split the dataset using the shuffled indices
        train_indices = indices[:train_size]
        xv_indices = indices[train_size : train_size + xv_size]
        test_indices = indices[train_size + xv_size :]

        train_dataset = Subset(self, train_indices)
        xv_dataset = Subset(self, xv_indices)
        test_dataset = Subset(self, test_indices)
        return (
            DataLoader(
                train_dataset,
                batch_size=batch_size,
                shuffle=True,
                num_workers=0,
                collate_fn=DittoDataset.__zero_padded,
            ),
            DataLoader(
                xv_dataset,
                batch_size=batch_size * 16,
                shuffle=False,
                num_workers=0,
                collate_fn=DittoDataset.__zero_padded,
            ),
            DataLoader(
                test_dataset,
                batch_size=batch_size * 16,
                shuffle=False,
                num_workers=0,
                collate_fn=DittoDataset.__zero_padded,
            ),
        )
