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




[((539, 572), 'warnings.filterwarnings', 'warnings.filterwarnings', (['"""ignore"""'], {}), "('ignore')\n", (562, 572), False, 'import warnings\n'), ((785, 858), 'utils.tokenization.BertTokenizer.from_pretrained', 'BertTokenizer.from_pretrained', (['args.model'], {'is_lowercase': 'args.is_lowercase'}), '(args.model, is_lowercase=args.is_lowercase)\n', (814, 858), False, 'from utils.tokenization import BertTokenizer\n'), ((885, 946), 'utils.emotion.Emotion', 'Emotion', (['args.nrc_path', 'args.max_em_len', 'args.emotion_filters'], {}), '(args.nrc_path, args.max_em_len, args.emotion_filters)\n', (892, 946), False, 'from utils.emotion import Emotion\n'), ((2358, 2408), 'torch.tensor', 'torch.tensor', (['unpadded_input_ids'], {'dtype': 'torch.long'}), '(unpadded_input_ids, dtype=torch.long)\n', (2370, 2408), False, 'import torch\n'), ((2438, 2489), 'torch.tensor', 'torch.tensor', (['unpadded_input_mask'], {'dtype': 'torch.long'}), '(unpadded_input_mask, dtype=torch.long)\n', (2450, 2489), False, 'import torch\n'), ((2520, 2572), 'torch.tensor', 'torch.tensor', (['unpadded_segment_ids'], {'dtype': 'torch.long'}), '(unpadded_segment_ids, dtype=torch.long)\n', (2532, 2572), False, 'import torch\n'), ((2603, 2658), 'torch.tensor', 'torch.tensor', (['unpadded_emotion_scores'], {'dtype': 'torch.long'}), '(unpadded_emotion_scores, dtype=torch.long)\n', (2615, 2658), False, 'import torch\n'), ((2680, 2747), 'torch.tensor', 'torch.tensor', (['[f.label_id for f in eval_features]'], {'dtype': 'torch.long'}), '([f.label_id for f in eval_features], dtype=torch.long)\n', (2692, 2747), False, 'import torch\n'), ((2771, 2876), 'torch.utils.data.TensorDataset', 'TensorDataset', (['padded_input_ids', 'padded_input_mask', 'padded_segment_ids', 'padded_emotion_ids', 'label_ids'], {}), '(padded_input_ids, padded_input_mask, padded_segment_ids,\n    padded_emotion_ids, label_ids)\n', (2784, 2876), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((2897, 2925), 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['eval_data'], {}), '(eval_data)\n', (2914, 2925), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((2953, 3029), 'torch.utils.data.DataLoader', 'DataLoader', (['eval_data'], {'sampler': 'eval_sampler', 'batch_size': 'self.args.batch_size'}), '(eval_data, sampler=eval_sampler, batch_size=self.args.batch_size)\n', (2963, 3029), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((3268, 3324), 'tqdm.tqdm', 'tqdm', (['eval_dataloader'], {'desc': '"""Evaluating"""', 'disable': 'silent'}), "(eval_dataloader, desc='Evaluating', disable=silent)\n", (3272, 3324), False, 'from tqdm import tqdm\n'), ((5390, 5445), 'sklearn.metrics.accuracy_score', 'metrics.accuracy_score', (['target_labels', 'predicted_labels'], {}), '(target_labels, predicted_labels)\n', (5412, 5445), False, 'from sklearn import metrics\n'), ((5467, 5540), 'sklearn.metrics.precision_score', 'metrics.precision_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), '(target_labels, predicted_labels, average=average)\n', (5490, 5540), False, 'from sklearn import metrics\n'), ((5559, 5629), 'sklearn.metrics.recall_score', 'metrics.recall_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), '(target_labels, predicted_labels, average=average)\n', (5579, 5629), False, 'from sklearn import metrics\n'), ((5703, 5756), 'sklearn.metrics.hamming_loss', 'metrics.hamming_loss', (['target_labels', 'predicted_labels'], {}), '(target_labels, predicted_labels)\n', (5723, 5756), False, 'from sklearn import metrics\n'), ((5782, 5853), 'sklearn.metrics.jaccard_score', 'metrics.jaccard_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), '(target_labels, predicted_labels, average=average)\n', (5803, 5853), False, 'from sklearn import metrics\n'), ((5874, 5940), 'sklearn.metrics.f1_score', 'metrics.f1_score', (['target_labels', 'predicted_labels'], {'average': 'average'}), '(target_labels, predicted_labels, average=average)\n', (5890, 5940), False, 'from sklearn import metrics\n'), ((5961, 6031), 'sklearn.metrics.f1_score', 'metrics.f1_score', (['target_labels', 'predicted_labels'], {'average': 'average_mac'}), '(target_labels, predicted_labels, average=average_mac)\n', (5977, 6031), False, 'from sklearn import metrics\n'), ((7734, 7784), 'torch.tensor', 'torch.tensor', (['unpadded_input_ids'], {'dtype': 'torch.long'}), '(unpadded_input_ids, dtype=torch.long)\n', (7746, 7784), False, 'import torch\n'), ((7814, 7865), 'torch.tensor', 'torch.tensor', (['unpadded_input_mask'], {'dtype': 'torch.long'}), '(unpadded_input_mask, dtype=torch.long)\n', (7826, 7865), False, 'import torch\n'), ((7896, 7948), 'torch.tensor', 'torch.tensor', (['unpadded_segment_ids'], {'dtype': 'torch.long'}), '(unpadded_segment_ids, dtype=torch.long)\n', (7908, 7948), False, 'import torch\n'), ((7979, 8031), 'torch.tensor', 'torch.tensor', (['unpadded_emotion_ids'], {'dtype': 'torch.long'}), '(unpadded_emotion_ids, dtype=torch.long)\n', (7991, 8031), False, 'import torch\n'), ((8053, 8120), 'torch.tensor', 'torch.tensor', (['[f.label_id for f in eval_features]'], {'dtype': 'torch.long'}), '([f.label_id for f in eval_features], dtype=torch.long)\n', (8065, 8120), False, 'import torch\n'), ((8144, 8249), 'torch.utils.data.TensorDataset', 'TensorDataset', (['padded_input_ids', 'padded_input_mask', 'padded_segment_ids', 'padded_emotion_ids', 'label_ids'], {}), '(padded_input_ids, padded_input_mask, padded_segment_ids,\n    padded_emotion_ids, label_ids)\n', (8157, 8249), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((8270, 8298), 'torch.utils.data.SequentialSampler', 'SequentialSampler', (['eval_data'], {}), '(eval_data)\n', (8287, 8298), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((8326, 8402), 'torch.utils.data.DataLoader', 'DataLoader', (['eval_data'], {'sampler': 'eval_sampler', 'batch_size': 'self.args.batch_size'}), '(eval_data, sampler=eval_sampler, batch_size=self.args.batch_size)\n', (8336, 8402), False, 'from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n'), ((8556, 8612), 'tqdm.tqdm', 'tqdm', (['eval_dataloader'], {'desc': '"""Evaluating"""', 'disable': 'silent'}), "(eval_dataloader, desc='Evaluating', disable=silent)\n", (8560, 8612), False, 'from tqdm import tqdm\n'), ((9304, 9337), 'torch.stack', 'torch.stack', (['bert_layers_l'], {'dim': '(0)'}), '(bert_layers_l, dim=0)\n', (9315, 9337), False, 'import torch\n'), ((1462, 1570), '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'], {}), '(self.eval_examples, self.args.\n    max_seq_length, self.tokenizer)\n', (1503, 1570), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), ((1628, 1752), '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'], {}), '(self.eval_examples, self.args.\n    max_seq_length, self.tokenizer, self.emotioner)\n', (1669, 1752), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), ((2110, 2172), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_ids', 'self.args.max_doc_length'], {}), '(unpadded_input_ids, self.args.max_doc_length)\n', (2126, 2172), False, 'from utils.preprocessing import pad_input_matrix\n'), ((2186, 2249), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_mask', 'self.args.max_doc_length'], {}), '(unpadded_input_mask, self.args.max_doc_length)\n', (2202, 2249), False, 'from utils.preprocessing import pad_input_matrix\n'), ((2263, 2327), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_segment_ids', 'self.args.max_doc_length'], {}), '(unpadded_segment_ids, self.args.max_doc_length)\n', (2279, 2327), False, 'from utils.preprocessing import pad_input_matrix\n'), ((5318, 5344), 'numpy.array', 'np.array', (['predicted_labels'], {}), '(predicted_labels)\n', (5326, 5344), True, 'import numpy as np\n'), ((5346, 5369), 'numpy.array', 'np.array', (['target_labels'], {}), '(target_labels)\n', (5354, 5369), True, 'import numpy as np\n'), ((6838, 6946), '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'], {}), '(self.eval_examples, self.args.\n    max_seq_length, self.tokenizer)\n', (6879, 6946), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), ((7004, 7128), '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'], {}), '(self.eval_examples, self.args.\n    max_seq_length, self.tokenizer, self.emotioner)\n', (7045, 7128), False, 'from datasets.bert_processors.abstract_processor import convert_examples_to_features_with_emotion, convert_examples_to_hierarchical_features\n'), ((7486, 7548), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_ids', 'self.args.max_doc_length'], {}), '(unpadded_input_ids, self.args.max_doc_length)\n', (7502, 7548), False, 'from utils.preprocessing import pad_input_matrix\n'), ((7562, 7625), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_input_mask', 'self.args.max_doc_length'], {}), '(unpadded_input_mask, self.args.max_doc_length)\n', (7578, 7625), False, 'from utils.preprocessing import pad_input_matrix\n'), ((7639, 7703), 'utils.preprocessing.pad_input_matrix', 'pad_input_matrix', (['unpadded_segment_ids', 'self.args.max_doc_length'], {}), '(unpadded_segment_ids, self.args.max_doc_length)\n', (7655, 7703), False, 'from utils.preprocessing import pad_input_matrix\n'), ((3636, 3651), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (3649, 3651), False, 'import torch\n'), ((8924, 8939), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (8937, 8939), False, 'import torch\n'), ((4855, 4885), 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), '(label_ids, dim=1)\n', (4867, 4885), False, 'import torch\n'), ((9127, 9157), 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), '(label_ids, dim=1)\n', (9139, 9157), False, 'import torch\n'), ((4662, 4689), 'torch.argmax', 'torch.argmax', (['logits'], {'dim': '(1)'}), '(logits, dim=1)\n', (4674, 4689), False, 'import torch\n'), ((4752, 4782), 'torch.argmax', 'torch.argmax', (['label_ids'], {'dim': '(1)'}), '(label_ids, dim=1)\n', (4764, 4782), False, 'import torch\n'), ((4308, 4325), 'torch.nn.functional.sigmoid', 'F.sigmoid', (['logits'], {}), '(logits)\n', (4317, 4325), True, 'import torch.nn.functional as F\n')]