import json
import re
import torch
import numpy as np
from collections import defaultdict

class summarization_data:
    def __init__(self, raw_data, tokenizer, domain_adaption=False, wwm_prob=0.1):
        data = self.select_data(raw_data)
        self.data = self.data_clean(data)
        self.len = len(self.data)
        self.tokenizer = tokenizer
        self.padding_max_len = "max_length"
        self.domain_adaption = domain_adaption
        self.wwm_prob = wwm_prob

    def select_data(self, raw_data):
        data = list()
        for item in raw_data:
            del item['pub_time']
            del item['labels']
            data.append(item)
        return data

    def data_clean(self, data):
        for item in data:
            item["text"] = re.sub(r"[?|!+?|:|(|)]|\\|-|/.*?/|http\S+", "", item["text"].lower())
            item["title"] = re.sub(r"[?|!+?|:|(|)]|\\|-|/.*?/|http\S+", "", item["title"].lower())
        return data
    

    def __getitem__(self, index):
        if self.domain_adaption:
            tokenized_text = self.tokenizer(
                self.data[index]["text"],
                add_special_tokens=True,
                padding="max_length",
                return_token_type_ids=True,
                truncation=True,
            )
            text_mask = tokenized_text['attention_mask']

            input_ids, labels = self._word_masking(self.tokenizer, tokenized_text, self.wwm_prob)
            return {
                'input_ids': torch.tensor(input_ids),
                'attention_mask': torch.tensor(text_mask),
                'labels': torch.tensor(labels)
            }

        else:
            tokenized_text = self.tokenizer(
                "summarize:"+self.data[index]["text"],
                add_special_tokens=True,
                padding="max_length",
                return_token_type_ids=True,
                truncation=True,
            )
            text_ids = tokenized_text['input_ids']
            text_mask = tokenized_text['attention_mask']

            tokenized_title = self.tokenizer(
                self.data[index]["title"],
                padding="max_length",
                return_token_type_ids=True,
                truncation=True,
            )
            title_ids = tokenized_title['input_ids']
            return {
                'input_ids': torch.tensor(text_ids),
                'attention_mask': torch.tensor(text_mask),
                'labels': torch.tensor(title_ids)
            }

    def _word_masking(self, tokenizer, tokenized_inputs, wwm_prob):
        # randomly mask words
        input_ids = tokenized_inputs["input_ids"]
        mask = np.random.binomial(1, wwm_prob, (len(input_ids),))

        labels = list()

        for idx in np.where(mask == 1)[0]:
            #add special sentinel tokens
            sentinel_token = tokenizer.additional_special_tokens[input_ids[idx] % 100]

            labels.append(tokenizer(sentinel_token).input_ids[0])
            labels.append(input_ids[idx])
            input_ids[idx] = tokenizer(sentinel_token).input_ids[0]

        return input_ids, labels

    def __len__(self):
        return self.len


# from tqdm.auto import tqdm
# from transformers import T5TokenizerFast, T5ForConditionalGeneration
# from transformers import DataCollatorForSeq2Seq
# from torch.utils.data import DataLoader, random_split
# from FTE_NLP.model.summarization_dataset_v1 import *
# from torch.optim import AdamW
# from transformers import get_scheduler
# from FTE_NLP.utils.post_process import *
#
# json_filename = '../data/raw_EDT/Trading_benchmark/evaluate_news_test.json'
# with open(json_filename) as data_file:
#     test_data = json.loads(data_file.read())
#
# model_checkpoint = "t5-small"
# tokenizer = T5TokenizerFast.from_pretrained(model_checkpoint, model_max_length=512)
# model = T5ForConditionalGeneration.from_pretrained(model_checkpoint)
#
# all_dataset = summarization_data(test_data, tokenizer, domain_adaption=True, wwm_prob=0.1)
# train_dataset, eval_dataset = random_split(all_dataset, [7, 3], generator=torch.Generator().manual_seed(42))
#
# # data_collator = DataCollatorForSeq2Seq(tokenizer, model,label_pad_token_id=0)
# data_collator = DataCollatorForSeq2Seq(tokenizer,model)
# # pass data to dataloader
# train_params = {'batch_size': 2, 'shuffle': False, 'num_workers': 0}
# train_loader = DataLoader(train_dataset, collate_fn=data_collator, **train_params)
#
# eval_params = {'batch_size': 2, 'shuffle': False, 'num_workers': 0}
# eval_loader = DataLoader(eval_dataset, collate_fn=data_collator, **train_params)
#
#
# for item in train_loader:
#     print(item)
#     print("input id  numbers:",item["input_ids"][0])
#     print("input id: ",tokenizer.decode(item["input_ids"][0]))
#     print("labels id number:",item["labels"][0])
#     print("labels:",tokenizer.decode(item["labels"][0],skip_special_tokens=False))
#     break