import numpy as np

import pandas as pd

import torch

import re

from torch import nn

from torch.utils.data import Dataset

from torch.utils.data import DataLoader

from transformers import BertTokenizer,BertForTokenClassification

import random

import nltk

from nltk.tokenize import sent_tokenize

nltk.download('punkt')

def shuffle_sentences_and_entities(text, entities):

    sentences = sent_tokenize(text)

    entity_tokens = entities.split() # align with words in text

    # identify start and end indices of sentences in terms of word counts

    word_counts = [len(sentence.split()) for sentence in sentences]

    start_indices = [sum(word_counts[:i]) for i in range(len(word_counts))]

    end_indices = [sum(word_counts[:i+1]) for i in range(len(word_counts))]

    # split entities into groups (corresponding to sentences)

    sentence_entities = [entity_tokens[start:end] for start, end in zip(start_indices, end_indices)]

    # shuffle sentence-entities pairs

    combined = list(zip(sentences, sentence_entities))

    random.seed(42)

    random.shuffle(combined)

    shuffled_sentences, shuffled_sentence_entities = zip(*combined)

    # reconstruction

    augmented_text = ' '.join(shuffled_sentences)

    augmented_entities = ' '.join([' '.join(entity_group) for entity_group in shuffled_sentence_entities])

    return augmented_text, augmented_entities

class Dataloader():

    """

    Dataloader used for loading the dataset used in this project. Also provides a framework for automatic

    tokenization of the data.

    """

    def __init__(self, label_to_ids, ids_to_label, transfer_learning, max_tokens, type):

        self.label_to_ids = label_to_ids

        self.ids_to_label = ids_to_label

        self.max_tokens = max_tokens

        self.transfer_learning = transfer_learning

        self.type = type

    def load_dataset(self, full = False, augment = False):

        """

        Loads the dataset and automatically initialized a tokenizer for the Custom_Dataset initialization.

        Parameters:

        full (bool): Whether the function should return the whole dataset or not - will return a train-val-test split

                     according to the Pareto principle (80:20).

        augment (bool): Whether the existing dataset should be extended via augmented data. Augmentation in this sense

                        means that the dataset will be extended via instances where the sentences are randomly switched around.

        Returns:

        if full:

            dataset (Custom_Dataset): the full dataset in one.

        else:

            tuple:

                - train_dataset (Custom_Dataset): Dataset used for training.

                - val_dataset (Custom_Dataset): Dataset used for validation.

                - test_dataset (Custom_Dataset): Dataset sued for testing.

        """

        if self.transfer_learning:

            data = pd.read_csv("../datasets/labelled_data/MEDCOND/all.csv", names=['text', 'entity'], header=None, sep="|")

            tokenizer = BertTokenizer.from_pretrained('alvaroalon2/biobert_diseases_ner')

            data['entity'] = data['entity'].apply(lambda x: x.replace('B-MEDCOND', 'B-DISEASE'))

            data['entity'] = data['entity'].apply(lambda x: x.replace('I-MEDCOND', 'I-DISEASE'))

        else:

            data = pd.read_csv(f"../datasets/labelled_data/{self.type}/all.csv", names=['text', 'entity'], header=None, sep="|")

            tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

            tokenizer.add_tokens(['B-' + self.type, 'I-' + self.type, 'O'])

        if not full:

            #train_data = data.sample((int) (len(data)*0.8), random_state=7).reset_index(drop=True)

            #test_data = data.drop(train_data.index).reset_index(drop=True)

            train_data = data.sample(frac=0.7, random_state=7).reset_index(drop=True)

            remaining_data = data.drop(train_data.index).reset_index(drop=True)

            val_data = remaining_data.sample(frac=0.2857, random_state=7).reset_index(drop=True)

            test_data = remaining_data.drop(val_data.index).reset_index(drop=True)

            if augment:

                augmented_rows = [shuffle_sentences_and_entities(text, entities) for text, entities in zip(train_data['text'], train_data['entity'])]

                augmented_texts, augmented_entities = zip(*augmented_rows)

                augmented_data = pd.DataFrame({'text': augmented_texts, 'entity': augmented_entities})

                train_data = pd.concat([train_data, augmented_data]).reset_index(drop=True)

            train_dataset = Custom_Dataset(train_data, tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

            val_dataset = Custom_Dataset(val_data, tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

            test_dataset = Custom_Dataset(test_data, tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

            return train_dataset, val_dataset, test_dataset

        else:

            dataset = Custom_Dataset(data, tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

            return dataset

    def load_custom(self, data):

        """

        Loads the dataset, but with entities swapped from MEDCOND to DISEASE (if transfer learning

        is enabled).

        Parameters:

        data (dataframe): Data extracted from csv file.

        Returns:

        dataset (Custom_Dataset): Dataset changed accordingly.

        """

        if self.transfer_learning:

            tokenizer = BertTokenizer.from_pretrained('alvaroalon2/biobert_diseases_ner')

            data['entity'] = data['entity'].apply(lambda x: x.replace('B-MEDCOND', 'B-DISEASE'))

            data['entity'] = data['entity'].apply(lambda x: x.replace('I-MEDCOND', 'I-DISEASE'))

        else:

            tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

            tokenizer.add_tokens(['B-' + self.type, 'I-' + self.type, 'O'])

        dataset = Custom_Dataset(data, tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

        return dataset

    def convert_id_to_label(self, ids):

        return [self.ids_to_label.get(x) for x in ids.numpy()[0]]

def tokenize_and_preserve_labels(sentence, text_labels, tokenizer, label_to_ids, ids_to_label, max_tokens):

    """

    Tokenizes each word separately. This may take longer, but increases accuracy. Preserves the labels

    of each word, adhereing to B and I prefixes.

    Parameters:

    sentence (string): Sentence to be tokenized.

    text_labels (numpy.array): Contains the labels of the sentence.

    tokenizer (BertTokenizer): Tokenizer used for tokenizing sentences.

    label_to_ids (dict): Dictionary containing label-id mappings.

    ids_to_label (dict): Dictionary containing id-label mappings.

    max_tokens (int): The maximum tokens allowed (input size of BERT model).

    Returns:

        tuple:

            - tokenized_sentence (numpy.array): Array containing all tokens of the give sentence.

            - labels (numpy.array): Array containing the corresponding labels of the tokens.

    """

    tokenized_sentence = []

    labels = []

    for word, label in zip(sentence, text_labels):

        tokenized_word = tokenizer.tokenize(word)

        n_subwords = len(tokenized_word)

        if(len(tokenized_sentence)>=max_tokens): #truncate

            return tokenized_sentence, labels

        tokenized_sentence.extend(tokenized_word)

        if label.startswith("B-"):

            labels.extend([label])

            labels.extend([ids_to_label.get(label_to_ids.get(label)+1)]*(n_subwords-1))

        else:

            labels.extend([label] * n_subwords)

    return tokenized_sentence, labels

class Custom_Dataset(Dataset):

    """

    Dataset used for loading and tokenizing sentences on-the-fly.

    """

    def __init__(self, data, tokenizer, label_to_ids, ids_to_label, max_tokens):

        self.data = data

        self.tokenizer = tokenizer

        self.label_to_ids = label_to_ids

        self.ids_to_label = ids_to_label

        self.max_tokens = max_tokens

    def __len__(self):

        return len(self.data)

    def __getitem__(self, idx):

        """

        Takes the current sentence with its labels and tokenizes it on-the-fly.

        Returns:

        item (torch.tensor): Tensor which can be fed into model.

        """

        sentence = re.findall(r"\w+|\w+(?='s)|'s|['\".,!?;]", self.data['text'][idx].strip(), re.UNICODE)

        word_labels = self.data['entity'][idx].split(" ")

        t_sen, t_labl = tokenize_and_preserve_labels(sentence, word_labels, self.tokenizer, self.label_to_ids, self.ids_to_label, self.max_tokens)

        sen_code = self.tokenizer.encode_plus(t_sen,

            add_special_tokens=True, # adds [CLS] and [SEP]

            max_length = self.max_tokens, # maximum tokens of a sentence

            padding='max_length',

            return_attention_mask=True, # generates the attention mask

            truncation = True

            )

        #shift labels (due to [CLS] and [SEP])

        labels = [-100]*self.max_tokens #-100 is ignore token

        for i, tok in enumerate(t_labl):

            if tok != None and i < self.max_tokens-1:

                labels[i+1]=self.label_to_ids.get(tok)

        item = {key: torch.as_tensor(val) for key, val in sen_code.items()}

        item['entity'] = torch.as_tensor(labels)

        return item