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, 0, 17, 33), 'warnings.filterwarnings', 'warnings.filterwarnings', ({(17, 24, 17, 32): '"""ignore"""'}, {}), "('ignore')", False, 'import warnings\n'), ((25, 25, 25, 98), 'utils.tokenization.BertTokenizer.from_pretrained', 'BertTokenizer.from_pretrained', (), '', False, 'from utils.tokenization import BertTokenizer\n'), ((26, 25, 26, 86), 'utils.emotion.Emotion', 'Emotion', ({(26, 33, 26, 46): 'args.nrc_path', (26, 48, 26, 63): 'args.max_em_len', (26, 65, 26, 85): 'args.emotion_filters'}, {}), '(args.nrc_path, args.max_em_len, args.emotion_filters)', False, 'from utils.emotion import Emotion\n'), ((54, 27, 54, 77), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((55, 28, 55, 79), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((56, 29, 56, 81), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((57, 29, 57, 84), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((58, 20, 58, 87), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((60, 20, 60, 121), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(60, 34, 60, 50): 'padded_input_ids', (60, 52, 60, 69): 'padded_input_mask', (60, 71, 60, 89): 'padded_segment_ids', (60, 91, 60, 109): 'padded_emotion_ids', (60, 111, 60, 120): 'label_ids'}, {}), '(padded_input_ids, padded_input_mask, padded_segment_ids,\n    padded_emotion_ids, label_ids)', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((61, 23, 61, 51), 'torch.utils.data.SequentialSampler', 'SequentialSampler', ({(61, 41, 61, 50): 'eval_data'}, {}), '(eval_data)', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((62, 26, 62, 102), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((70, 74, 70, 130), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((113, 19, 113, 74), 'sklearn.metrics.accuracy_score', 'metrics.accuracy_score', ({(113, 42, 113, 55): 'target_labels', (113, 57, 113, 73): 'predicted_labels'}, {}), '(target_labels, predicted_labels)', False, 'from sklearn import metrics\n'), ((114, 20, 114, 93), 'sklearn.metrics.precision_score', 'metrics.precision_score', (), '', False, 'from sklearn import metrics\n'), ((115, 17, 115, 87), 'sklearn.metrics.recall_score', 'metrics.recall_score', (), '', False, 'from sklearn import metrics\n'), ((118, 23, 118, 76), 'sklearn.metrics.hamming_loss', 'metrics.hamming_loss', ({(118, 44, 118, 57): 'target_labels', (118, 59, 118, 75): 'predicted_labels'}, {}), '(target_labels, predicted_labels)', False, 'from sklearn import metrics\n'), ((119, 24, 119, 95), 'sklearn.metrics.jaccard_score', 'metrics.jaccard_score', (), '', False, 'from sklearn import metrics\n'), ((120, 19, 120, 85), 'sklearn.metrics.f1_score', 'metrics.f1_score', (), '', False, 'from sklearn import metrics\n'), ((121, 19, 121, 89), 'sklearn.metrics.f1_score', 'metrics.f1_score', (), '', False, 'from sklearn import metrics\n'), ((151, 27, 151, 77), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((152, 28, 152, 79), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((153, 29, 153, 81), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((154, 29, 154, 81), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((155, 20, 155, 87), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((157, 20, 157, 121), 'torch.utils.data.TensorDataset', 'TensorDataset', ({(157, 34, 157, 50): 'padded_input_ids', (157, 52, 157, 69): 'padded_input_mask', (157, 71, 157, 89): 'padded_segment_ids', (157, 91, 157, 109): 'padded_emotion_ids', (157, 111, 157, 120): 'label_ids'}, {}), '(padded_input_ids, padded_input_mask, padded_segment_ids,\n    padded_emotion_ids, label_ids)', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((158, 23, 158, 51), 'torch.utils.data.SequentialSampler', 'SequentialSampler', ({(158, 41, 158, 50): 'eval_data'}, {}), '(eval_data)', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((159, 26, 159, 102), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((165, 74, 165, 130), 'tqdm.tqdm', 'tqdm', (), '', False, 'from tqdm import tqdm\n'), ((177, 24, 177, 57), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((37, 28, 38, 77), 'datasets.bert_processors.abstract_processor.convert_examples_to_hierarchical_features', 'convert_examples_to_hierarchical_features', ({(38, 16, 38, 34): 'self.eval_examples', (38, 36, 38, 60): 'self.args.max_seq_length', (38, 62, 38, 76): 'self.tokenizer'}, {}), '(self.eval_examples, self.args.\n    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, 28, 41, 93), 'datasets.bert_processors.abstract_processor.convert_examples_to_features_with_emotion', 'convert_examples_to_features_with_emotion', ({(41, 16, 41, 34): 'self.eval_examples', (41, 36, 41, 60): 'self.args.max_seq_length', (41, 62, 41, 76): 'self.tokenizer', (41, 78, 41, 92): 'self.emotioner'}, {}), '(self.eval_examples, self.args.\n    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, 12, 50, 74), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(50, 29, 50, 47): 'unpadded_input_ids', (50, 49, 50, 73): 'self.args.max_doc_length'}, {}), '(unpadded_input_ids, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((51, 12, 51, 75), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(51, 29, 51, 48): 'unpadded_input_mask', (51, 50, 51, 74): 'self.args.max_doc_length'}, {}), '(unpadded_input_mask, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((52, 12, 52, 76), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(52, 29, 52, 49): 'unpadded_segment_ids', (52, 51, 52, 75): 'self.args.max_doc_length'}, {}), '(unpadded_segment_ids, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((112, 42, 112, 68), 'numpy.array', 'np.array', ({(112, 51, 112, 67): 'predicted_labels'}, {}), '(predicted_labels)', True, 'import numpy as np\n'), ((112, 70, 112, 93), 'numpy.array', 'np.array', ({(112, 79, 112, 92): 'target_labels'}, {}), '(target_labels)', True, 'import numpy as np\n'), ((134, 28, 135, 77), 'datasets.bert_processors.abstract_processor.convert_examples_to_hierarchical_features', 'convert_examples_to_hierarchical_features', ({(135, 16, 135, 34): 'self.eval_examples', (135, 36, 135, 60): 'self.args.max_seq_length', (135, 62, 135, 76): 'self.tokenizer'}, {}), '(self.eval_examples, self.args.\n    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, 28, 138, 93), 'datasets.bert_processors.abstract_processor.convert_examples_to_features_with_emotion', 'convert_examples_to_features_with_emotion', ({(138, 16, 138, 34): 'self.eval_examples', (138, 36, 138, 60): 'self.args.max_seq_length', (138, 62, 138, 76): 'self.tokenizer', (138, 78, 138, 92): 'self.emotioner'}, {}), '(self.eval_examples, self.args.\n    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, 12, 147, 74), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(147, 29, 147, 47): 'unpadded_input_ids', (147, 49, 147, 73): 'self.args.max_doc_length'}, {}), '(unpadded_input_ids, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((148, 12, 148, 75), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(148, 29, 148, 48): 'unpadded_input_mask', (148, 50, 148, 74): 'self.args.max_doc_length'}, {}), '(unpadded_input_mask, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((149, 12, 149, 76), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', ({(149, 29, 149, 49): 'unpadded_segment_ids', (149, 51, 149, 75): 'self.args.max_doc_length'}, {}), '(unpadded_segment_ids, self.args.max_doc_length)', False, 'from utils.preprocessing import pad_input_matrix\n'), ((77, 17, 77, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((172, 17, 172, 32), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((100, 47, 100, 77), 'torch.argmax', 'torch.argmax', (), '', False, 'import torch\n'), ((174, 28, 174, 58), 'torch.argmax', 'torch.argmax', (), '', False, 'import torch\n'), ((98, 40, 98, 67), 'torch.argmax', 'torch.argmax', (), '', False, 'import torch\n'), ((99, 37, 99, 67), 'torch.argmax', 'torch.argmax', (), '', False, 'import torch\n'), ((92, 40, 92, 57), 'torch.nn.functional.sigmoid', 'F.sigmoid', ({(92, 50, 92, 56): 'logits'}, {}), '(logits)', True, 'import torch.nn.functional as F\n')]