import warnings

import numpy as np
import torch
import torch.nn.functional as F
from sklearn import metrics
from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
from tqdm import tqdm

from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, \
    convert_examples_to_hierarchical_features
from utils.preprocessing import pad_input_matrix
from utils.tokenization import BertTokenizer
from utils.emotion import Emotion

# Suppress warnings from sklearn.metrics
warnings.filterwarnings('ignore')


class BertEvaluator(object):
    def __init__(self, model, processor, args, split='dev'):
        self.args = args
        self.model = model
        self.processor = processor
        self.tokenizer = BertTokenizer.from_pretrained(args.model, is_lowercase=args.is_lowercase)
        self.emotioner = Emotion(args.nrc_path, args.max_em_len, args.emotion_filters)
        if split == 'test':
            self.eval_examples = self.processor.get_test_examples(args.data_dir, args.test_name)
        elif split == 'dev':
            self.eval_examples = self.processor.get_dev_examples(args.data_dir, args.dev_name)
        else:
            self.eval_examples = self.processor.get_any_examples(args.data_dir, split)

    def get_scores(self, silent=False, return_indices=False):
        all_indices = []
        if self.args.is_hierarchical:
            eval_features = convert_examples_to_hierarchical_features(
                self.eval_examples, self.args.max_seq_length, self.tokenizer)
        else:
            eval_features = convert_examples_to_features_with_emotion(
                self.eval_examples, self.args.max_seq_length, self.tokenizer, self.emotioner)

        unpadded_input_ids = [f.input_ids for f in eval_features]
        unpadded_input_mask = [f.input_mask for f in eval_features]
        unpadded_segment_ids = [f.segment_ids for f in eval_features]
        unpadded_emotion_scores = [f.sentiment_scores for f in eval_features]


        if self.args.is_hierarchical:
            pad_input_matrix(unpadded_input_ids, self.args.max_doc_length)
            pad_input_matrix(unpadded_input_mask, self.args.max_doc_length)
            pad_input_matrix(unpadded_segment_ids, self.args.max_doc_length)

        padded_input_ids = torch.tensor(unpadded_input_ids, dtype=torch.long)
        padded_input_mask = torch.tensor(unpadded_input_mask, dtype=torch.long)
        padded_segment_ids = torch.tensor(unpadded_segment_ids, dtype=torch.long)
        padded_emotion_ids = torch.tensor(unpadded_emotion_scores, dtype=torch.long)
        label_ids = torch.tensor([f.label_id for f in eval_features], dtype=torch.long)

        eval_data = TensorDataset(padded_input_ids, padded_input_mask, padded_segment_ids, padded_emotion_ids, label_ids)
        eval_sampler = SequentialSampler(eval_data)
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=self.args.batch_size)

        self.model.eval()

        total_loss = 0
        nb_eval_steps, nb_eval_examples = 0, 0
        predicted_labels, target_labels = list(), list()

        for input_ids, input_mask, segment_ids, emotion_ids, label_ids in tqdm(eval_dataloader, desc="Evaluating", disable=silent):
            input_ids = input_ids.to(self.args.device)
            input_mask = input_mask.to(self.args.device)
            segment_ids = segment_ids.to(self.args.device)
            emotion_ids = emotion_ids.to(self.args.device)
            label_ids = label_ids.to(self.args.device)

            with torch.no_grad():
                if return_indices:
                    outs = self.model(input_ids, segment_ids, input_mask, emotion_ids=emotion_ids, return_indices=return_indices)
                else:
                    outs = self.model(input_ids, segment_ids, input_mask, emotion_ids=emotion_ids)
                    if isinstance(outs, tuple):
                        outs, _ = outs

            if return_indices:
                logits, indices = outs
                all_indices.extend(indices.cpu().detach().numpy())
            else:
                logits = outs

            if self.args.is_multilabel:
                predicted_labels.extend(F.sigmoid(logits).round().long().cpu().detach().numpy())
                target_labels.extend(label_ids.cpu().detach().numpy())
                loss = F.binary_cross_entropy_with_logits(logits, label_ids.float(), size_average=False)
                average, average_mac = 'micro', 'macro'

            else:
                predicted_labels.extend(torch.argmax(logits, dim=1).cpu().detach().numpy())
                target_labels.extend(torch.argmax(label_ids, dim=1).cpu().detach().numpy())
                loss = F.cross_entropy(logits, torch.argmax(label_ids, dim=1))
                average, average_mac = 'binary', 'binary'

            if self.args.n_gpu > 1:
                loss = loss.mean()
            if self.args.gradient_accumulation_steps > 1:
                loss = loss / self.args.gradient_accumulation_steps
            total_loss += loss.item()

            nb_eval_examples += input_ids.size(0)
            nb_eval_steps += 1

        predicted_labels, target_labels = np.array(predicted_labels), np.array(target_labels)
        accuracy = metrics.accuracy_score(target_labels, predicted_labels)
        precision = metrics.precision_score(target_labels, predicted_labels, average=average)
        recall = metrics.recall_score(target_labels, predicted_labels, average=average)
        avg_loss = total_loss / nb_eval_steps

        hamming_loss = metrics.hamming_loss(target_labels, predicted_labels)
        jaccard_score = metrics.jaccard_score(target_labels, predicted_labels, average=average)
        f1_micro = metrics.f1_score(target_labels, predicted_labels, average=average)
        f1_macro = metrics.f1_score(target_labels, predicted_labels, average=average_mac)

        if return_indices:
            return [accuracy, precision, recall, f1_micro, avg_loss, f1_macro, hamming_loss, jaccard_score, predicted_labels, target_labels, all_indices],\
               ['accuracy', 'precision', 'recall', 'f1_micro', 'avg_loss', 'f1_macro', 'hamming_loss', 'jaccard', 'predicted_labels', 'target_labels', 'all_indices']
        else:

            return [accuracy, precision, recall, f1_micro, avg_loss, f1_macro, hamming_loss, jaccard_score, predicted_labels, target_labels],\
               ['accuracy', 'precision', 'recall', 'f1_micro', 'avg_loss', 'f1_macro', 'hamming_loss', 'jaccard', 'predicted_labels', 'target_labels']


    def get_bert_layers(self, silent=False, last_bert_layers=-1):
        if self.args.is_hierarchical:
            eval_features = convert_examples_to_hierarchical_features(
                self.eval_examples, self.args.max_seq_length, self.tokenizer)
        else:
            eval_features = convert_examples_to_features_with_emotion(
                self.eval_examples, self.args.max_seq_length, self.tokenizer, self.emotioner)

        unpadded_input_ids = [f.input_ids for f in eval_features]
        unpadded_input_mask = [f.input_mask for f in eval_features]
        unpadded_segment_ids = [f.segment_ids for f in eval_features]
        unpadded_emotion_ids = [f.emotioniment_scores for f in eval_features]


        if self.args.is_hierarchical:
            pad_input_matrix(unpadded_input_ids, self.args.max_doc_length)
            pad_input_matrix(unpadded_input_mask, self.args.max_doc_length)
            pad_input_matrix(unpadded_segment_ids, self.args.max_doc_length)

        padded_input_ids = torch.tensor(unpadded_input_ids, dtype=torch.long)
        padded_input_mask = torch.tensor(unpadded_input_mask, dtype=torch.long)
        padded_segment_ids = torch.tensor(unpadded_segment_ids, dtype=torch.long)
        padded_emotion_ids = torch.tensor(unpadded_emotion_ids, dtype=torch.long)
        label_ids = torch.tensor([f.label_id for f in eval_features], dtype=torch.long)

        eval_data = TensorDataset(padded_input_ids, padded_input_mask, padded_segment_ids, padded_emotion_ids, label_ids)
        eval_sampler = SequentialSampler(eval_data)
        eval_dataloader = DataLoader(eval_data, sampler=eval_sampler, batch_size=self.args.batch_size)

        self.model.eval()

        bert_layers_l, label_ids_l = [], []

        for input_ids, input_mask, segment_ids, emotion_ids, label_ids in tqdm(eval_dataloader, desc="Evaluating", disable=silent):
            input_ids = input_ids.to(self.args.device)
            input_mask = input_mask.to(self.args.device)
            segment_ids = segment_ids.to(self.args.device)
            emotion_ids = emotion_ids.to(self.args.device)
            label_ids = label_ids.to(self.args.device)

            with torch.no_grad():
                bert_layers = self.model.get_bert_embedding(input_ids, segment_ids, input_mask, emotion_ids=emotion_ids, last_bert_layers=last_bert_layers)
                label_ids = torch.argmax(label_ids, dim=1).cpu().detach().numpy()
                bert_layers_l.extend(bert_layers)
                label_ids_l.extend(label_ids)
        bert_layers_l = torch.stack(bert_layers_l, dim=0)
        return bert_layers_l, label_ids_l




[(17, 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), False, 'import warnings\n'), (25, 'utils.tokenization.BertTokenizer.from_pretrained', 'BertTokenizer.from_pretrained', (['args.model'], {'is_lowercase': 'args.is_lowercase'}), False, 'from utils.tokenization import BertTokenizer\n'), (26, 'utils.emotion.Emotion', 'Emotion', (['args.nrc_path', 'args.max_em_len', 'args.emotion_filters'], {}), False, 'from utils.emotion import Emotion\n'), (54, 'torch.tensor', 'torch.tensor', (['unpadded_input_ids'], {'dtype': 'torch.long'}), False, 'import torch\n'), (55, 'torch.tensor', 'torch.tensor', (['unpadded_input_mask'], {'dtype': 'torch.long'}), False, 'import torch\n'), (56, 'torch.tensor', 'torch.tensor', (['unpadded_segment_ids'], {'dtype': 'torch.long'}), False, 'import torch\n'), (57, 'torch.tensor', 'torch.tensor', (['unpadded_emotion_scores'], {'dtype': 'torch.long'}), False, 'import torch\n'), (58, 'torch.tensor', 'torch.tensor', (['[f.label_id for f in eval_features]'], {'dtype': 'torch.long'}), False, 'import torch\n'), (60, 'torch.utils.data.TensorDataset', 'TensorDataset', (['padded_input_ids', 'padded_input_mask', 'padded_segment_ids', 'padded_emotion_ids', 'label_ids'], {}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (61, 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['eval_data'], {}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (62, 'torch.utils.data.DataLoader', 'DataLoader', (['eval_data'], {'sampler': 'eval_sampler', 'batch_size': 'self.args.batch_size'}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (70, 'tqdm.tqdm', 'tqdm', (['eval_dataloader'], {'desc': '"""Evaluating"""', 'disable': 'silent'}), False, 'from tqdm import tqdm\n'), (113, 'sklearn.metrics.accuracy_score', 'metrics.accuracy_score', (['target_labels', 'predicted_labels'], {}), False, 'from sklearn import metrics\n'), (114, 'sklearn.metrics.precision_score', 'metrics.precision_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), False, 'from sklearn import metrics\n'), (115, 'sklearn.metrics.recall_score', 'metrics.recall_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), False, 'from sklearn import metrics\n'), (118, 'sklearn.metrics.hamming_loss', 'metrics.hamming_loss', (['target_labels', 'predicted_labels'], {}), False, 'from sklearn import metrics\n'), (119, 'sklearn.metrics.jaccard_score', 'metrics.jaccard_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), False, 'from sklearn import metrics\n'), (120, 'sklearn.metrics.f1_score', 'metrics.f1_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), False, 'from sklearn import metrics\n'), (121, 'sklearn.metrics.f1_score', 'metrics.f1_score', (['target_labels', 'predicted_labels'], {'average': 'average_mac'}), False, 'from sklearn import metrics\n'), (151, 'torch.tensor', 'torch.tensor', (['unpadded_input_ids'], {'dtype': 'torch.long'}), False, 'import torch\n'), (152, 'torch.tensor', 'torch.tensor', (['unpadded_input_mask'], {'dtype': 'torch.long'}), False, 'import torch\n'), (153, 'torch.tensor', 'torch.tensor', (['unpadded_segment_ids'], {'dtype': 'torch.long'}), False, 'import torch\n'), (154, 'torch.tensor', 'torch.tensor', (['unpadded_emotion_ids'], {'dtype': 'torch.long'}), False, 'import torch\n'), (155, 'torch.tensor', 'torch.tensor', (['[f.label_id for f in eval_features]'], {'dtype': 'torch.long'}), False, 'import torch\n'), (157, 'torch.utils.data.TensorDataset', 'TensorDataset', (['padded_input_ids', 'padded_input_mask', 'padded_segment_ids', 'padded_emotion_ids', 'label_ids'], {}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (158, 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['eval_data'], {}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (159, 'torch.utils.data.DataLoader', 'DataLoader', (['eval_data'], {'sampler': 'eval_sampler', 'batch_size': 'self.args.batch_size'}), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), (165, 'tqdm.tqdm', 'tqdm', (['eval_dataloader'], {'desc': '"""Evaluating"""', 'disable': 'silent'}), False, 'from tqdm import tqdm\n'), (177, 'torch.stack', 'torch.stack', (['bert_layers_l'], {'dim': '(0)'}), False, 'import torch\n'), (37, 'datasets.bert_processors.abstract_processor.convert_examples_to_hierarchical_features', 'convert_examples_to_hierarchical_features', (['self.eval_examples', 'self.args.max_seq_length', 'self.tokenizer'], {}), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), (40, 'datasets.bert_processors.abstract_processor.convert_examples_to_features_with_emotion', 'convert_examples_to_features_with_emotion', (['self.eval_examples', 'self.args.max_seq_length', 'self.tokenizer', 'self.emotioner'], {}), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), (50, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_ids', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (51, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_mask', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (52, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_segment_ids', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (112, 'numpy.array', 'np.array', (['predicted_labels'], {}), True, 'import numpy as np\n'), (112, 'numpy.array', 'np.array', (['target_labels'], {}), True, 'import numpy as np\n'), (134, 'datasets.bert_processors.abstract_processor.convert_examples_to_hierarchical_features', 'convert_examples_to_hierarchical_features', (['self.eval_examples', 'self.args.max_seq_length', 'self.tokenizer'], {}), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), (137, 'datasets.bert_processors.abstract_processor.convert_examples_to_features_with_emotion', 'convert_examples_to_features_with_emotion', (['self.eval_examples', 'self.args.max_seq_length', 'self.tokenizer', 'self.emotioner'], {}), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), (147, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_ids', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (148, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_mask', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (149, 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_segment_ids', 'self.args.max_doc_length'], {}), False, 'from utils.preprocessing import pad_input_matrix\n'), (77, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (172, 'torch.no_grad', 'torch.no_grad', ([], {}), False, 'import torch\n'), (100, 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), False, 'import torch\n'), (174, 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), False, 'import torch\n'), (98, 'torch.argmax', 'torch.argmax', (['logits'], {'dim': '(1)'}), False, 'import torch\n'), (99, 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), False, 'import torch\n'), (92, 'torch.nn.functional.sigmoid', 'F.sigmoid', (['logits'], {}), True, 'import torch.nn.functional as F\n')]