import torch
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel

from torch.nn.modules.container import ModuleList
import copy

# from dataset import load_data
from transformers import BertTokenizer

from torch.utils.data import DataLoader
from torch.utils.data import Dataset

from tqdm import tqdm
import numpy as np
import src.Multi2OIE.utils.bio as bio


# from extract import extract

def _get_clones(module, n):
    return ModuleList([copy.deepcopy(module) for _ in range(n)])


def _get_position_idxs(pred_mask, input_ids):
    position_idxs = torch.zeros(pred_mask.shape, dtype=int, device=pred_mask.device)
    for mask_idx, cur_mask in enumerate(pred_mask):
        position_idxs[mask_idx, :] += 2
        cur_nonzero = (cur_mask == 0).nonzero()
        start = torch.min(cur_nonzero).item()
        end = torch.max(cur_nonzero).item()
        position_idxs[mask_idx, start:end + 1] = 1
        pad_start = max(input_ids[mask_idx].nonzero()).item() + 1
        position_idxs[mask_idx, pad_start:] = 0
    return position_idxs


def _get_pred_feature(pred_hidden, pred_mask):
    B, L, D = pred_hidden.shape
    pred_features = torch.zeros((B, L, D), device=pred_mask.device)
    for mask_idx, cur_mask in enumerate(pred_mask):
        pred_position = (cur_mask == 0).nonzero().flatten()
        pred_feature = torch.mean(pred_hidden[mask_idx, pred_position], dim=0)
        pred_feature = torch.cat(L * [pred_feature.unsqueeze(0)])
        pred_features[mask_idx, :, :] = pred_feature
    return pred_features


def _get_activation_fn(activation):
    if activation == "relu":
        return F.relu
    elif activation == "gelu":
        return F.gelu
    else:
        raise RuntimeError("activation should be relu/gelu, not %s." % activation)


class ArgExtractorLayer(nn.Module):
    def __init__(self,
                 d_model=768,
                 n_heads=8,
                 d_feedforward=2048,
                 dropout=0.1,
                 activation='relu'):
        """
        A layer similar to Transformer decoder without decoder self-attention.
        (only encoder-decoder multi-head attention followed by feed-forward layers)

        :param d_model: model dimensionality (default=768 from BERT-base)
        :param n_heads: number of heads in multi-head attention layer
        :param d_feedforward: dimensionality of point-wise feed-forward layer
        :param dropout: drop rate of all layers
        :param activation: activation function after first feed-forward layer
        """
        super(ArgExtractorLayer, self).__init__()
        self.multihead_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
        self.linear1 = nn.Linear(d_model, d_feedforward)
        self.dropout1 = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_feedforward, d_model)

        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)
        self.activation = _get_activation_fn(activation)

    def forward(self, target, source, key_mask=None):
        """
        Single Transformer Decoder layer without self-attention

        :param target: a tensor which takes a role as a query
        :param source: a tensor which takes a role as a key & value
        :param key_mask: key mask tensor with the shape of (batch_size, sequence_length)
        """
        # Multi-head attention layer (+ add & norm)
        attended = self.multihead_attn(
            target, source, source,
            key_padding_mask=key_mask)[0]
        skipped = target + self.dropout1(attended)
        normed = self.norm1(skipped)

        # Point-wise feed-forward layer (+ add & norm)
        projected = self.linear2(self.dropout2(self.activation(self.linear1(normed))))
        skipped = normed + self.dropout1(projected)
        normed = self.norm2(skipped)
        return normed


class ArgModule(nn.Module):
    def __init__(self, arg_layer, n_layers):
        """
        Module for extracting arguments based on given encoder output and predicates.
        It uses ArgExtractorLayer as a base block and repeat the block N('n_layers') times

        :param arg_layer: an instance of the ArgExtractorLayer() class (required)
        :param n_layers: the number of sub-layers in the ArgModule (required).
        """
        super(ArgModule, self).__init__()
        self.layers = _get_clones(arg_layer, n_layers)
        self.n_layers = n_layers

    def forward(self, encoded, predicate, pred_mask=None):
        """
        :param encoded: output from sentence encoder with the shape of (L, B, D),
            where L is the sequence length, B is the batch size, D is the embedding dimension
        :param predicate: output from predicate module with the shape of (L, B, D)
        :param pred_mask: mask that prevents attention to tokens which are not predicates
            with the shape of (B, L)
        :return: tensor like Transformer Decoder Layer Output
        """
        output = encoded
        for layer_idx in range(self.n_layers):
            output = self.layers[layer_idx](
                target=output, source=predicate, key_mask=pred_mask)
        return output


class Multi2OIE(nn.Module):
    def __init__(self,
                 bert_config,
                 mh_dropout=0.1,
                 pred_clf_dropout=0.,
                 arg_clf_dropout=0.3,
                 n_arg_heads=8,
                 n_arg_layers=4,
                 pos_emb_dim=64,
                 pred_n_labels=3,
                 arg_n_labels=9):
        super(Multi2OIE, self).__init__()
        self.pred_n_labels = pred_n_labels
        self.arg_n_labels = arg_n_labels

        self.bert = BertModel.from_pretrained(
            bert_config,
            output_hidden_states=True)
        d_model = self.bert.config.hidden_size
        self.pred_dropout = nn.Dropout(pred_clf_dropout)
        self.pred_classifier = nn.Linear(d_model, self.pred_n_labels)

        self.position_emb = nn.Embedding(3, pos_emb_dim, padding_idx=0)
        d_model += (d_model + pos_emb_dim)
        arg_layer = ArgExtractorLayer(
            d_model=d_model,
            n_heads=n_arg_heads,
            dropout=mh_dropout)
        self.arg_module = ArgModule(arg_layer, n_arg_layers)
        self.arg_dropout = nn.Dropout(arg_clf_dropout)
        self.arg_classifier = nn.Linear(d_model, arg_n_labels)

    def forward(self,
                input_ids,
                attention_mask,
                predicate_mask=None,
                predicate_hidden=None,
                total_pred_labels=None,
                arg_labels=None):

        # predicate extraction
        bert_hidden = self.bert(input_ids, attention_mask)[0]
        pred_logit = self.pred_classifier(self.pred_dropout(bert_hidden))

        # predicate loss
        if total_pred_labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            active_loss = attention_mask.view(-1) == 1
            active_logits = pred_logit.view(-1, self.pred_n_labels)
            active_labels = torch.where(
                active_loss, total_pred_labels.view(-1),
                torch.tensor(loss_fct.ignore_index).type_as(total_pred_labels))
            pred_loss = loss_fct(active_logits, active_labels)

        # inputs for argument extraction
        pred_feature = _get_pred_feature(bert_hidden, predicate_mask)
        position_vectors = self.position_emb(_get_position_idxs(predicate_mask, input_ids))
        bert_hidden = torch.cat([bert_hidden, pred_feature, position_vectors], dim=2)
        bert_hidden = bert_hidden.transpose(0, 1)

        # argument extraction
        arg_hidden = self.arg_module(bert_hidden, bert_hidden, predicate_mask)
        arg_hidden = arg_hidden.transpose(0, 1)
        arg_logit = self.arg_classifier(self.arg_dropout(arg_hidden))

        # argument loss
        if arg_labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            active_loss = attention_mask.view(-1) == 1
            active_logits = arg_logit.view(-1, self.arg_n_labels)
            active_labels = torch.where(
                active_loss, arg_labels.view(-1),
                torch.tensor(loss_fct.ignore_index).type_as(arg_labels))
            arg_loss = loss_fct(active_logits, active_labels)

        # total loss
        batch_loss = pred_loss + arg_loss
        outputs = (batch_loss, pred_loss, arg_loss)
        return outputs

    def extract_predicate(self,
                          input_ids,
                          attention_mask):
        bert_hidden = self.bert(input_ids, attention_mask)[0]
        pred_logit = self.pred_classifier(bert_hidden)
        return pred_logit, bert_hidden

    def extract_argument(self,
                         input_ids,
                         predicate_hidden,
                         predicate_mask):
        pred_feature = _get_pred_feature(predicate_hidden, predicate_mask)
        position_vectors = self.position_emb(_get_position_idxs(predicate_mask, input_ids))
        arg_input = torch.cat([predicate_hidden, pred_feature, position_vectors], dim=2)
        arg_input = arg_input.transpose(0, 1)
        arg_hidden = self.arg_module(arg_input, arg_input, predicate_mask)
        arg_hidden = arg_hidden.transpose(0, 1)
        return self.arg_classifier(arg_hidden)


class OieEvalDataset(Dataset):
    def __init__(self, sentences, max_len, tokenizer_config):
        self.sentences = sentences
        self.tokenizer = BertTokenizer.from_pretrained(tokenizer_config)
        self.vocab = self.tokenizer.vocab
        self.max_len = max_len

        self.pad_idx = self.vocab['[PAD]']
        self.cls_idx = self.vocab['[CLS]']
        self.sep_idx = self.vocab['[SEP]']
        self.mask_idx = self.vocab['[MASK]']

    def add_pad(self, token_ids):
        diff = self.max_len - len(token_ids)
        if diff > 0:
            token_ids += [self.pad_idx] * diff
        else:
            token_ids = token_ids[:self.max_len - 1] + [self.sep_idx]
        return token_ids

    def idx2mask(self, token_ids):
        return [token_id != self.pad_idx for token_id in token_ids]

    def __len__(self):
        return len(self.sentences)

    def __getitem__(self, idx):
        token_ids = self.add_pad(self.tokenizer.encode(self.sentences[idx]))
        att_mask = self.idx2mask(token_ids)
        token_strs = self.tokenizer.convert_ids_to_tokens(token_ids)
        sentence = self.sentences[idx]

        assert len(token_ids) == self.max_len
        assert len(att_mask) == self.max_len
        assert len(token_strs) == self.max_len
        batch = [
            torch.tensor(token_ids),
            torch.tensor(att_mask),
            token_strs,
            sentence
        ]
        return batch


def extract(model, loader, device, tokenizer):

    # model.eval()
    result = []
    for step, batch in tqdm(enumerate(loader), desc='eval_steps', total=len(loader)):
        token_strs = [[word for word in sent] for sent in np.asarray(batch[-2]).T]
        sentences = batch[-1]
        print(sentences)
        token_ids, att_mask = map(lambda x: x.to(device), batch[:-2])

        with torch.no_grad():
            """
            We will iterate B(batch_size) times
            because there are more than one predicate in one batch.
            In feeding to argument extractor, # of predicates takes a role as batch size.

            pred_logit: (B, L, 3)
            pred_hidden: (B, L, D)
            pred_tags: (B, P, L) ~ list of tensors, where P is # of predicate in each batch
            """
            pred_logit, pred_hidden = model.extract_predicate(
                input_ids=token_ids, attention_mask=att_mask)
            pred_tags = torch.argmax(pred_logit, 2)
            pred_tags = bio.filter_pred_tags(pred_tags, token_strs)
            pred_tags = bio.get_single_predicate_idxs(pred_tags)
            pred_probs = torch.nn.Softmax(2)(pred_logit)

            # iterate B times (one iteration means extraction for one sentence)
            for cur_pred_tags, cur_pred_hidden, cur_att_mask, cur_token_id, cur_pred_probs, token_str, sentence \
                    in zip(pred_tags, pred_hidden, att_mask, token_ids, pred_probs, token_strs, sentences):

                # generate temporary batch for this sentence and feed to argument module
                cur_pred_masks = bio.get_pred_mask(cur_pred_tags).to(device)
                n_predicates = cur_pred_masks.shape[0]
                if n_predicates == 0:
                    continue  # if there is no predicate, we cannot extract.
                cur_pred_hidden = torch.cat(n_predicates * [cur_pred_hidden.unsqueeze(0)])
                cur_token_id = torch.cat(n_predicates * [cur_token_id.unsqueeze(0)])
                cur_arg_logit = model.extract_argument(
                    input_ids=cur_token_id,
                    predicate_hidden=cur_pred_hidden,
                    predicate_mask=cur_pred_masks)

                # filter and get argument tags with highest probability
                cur_arg_tags = torch.argmax(cur_arg_logit, 2)
                cur_arg_probs = torch.nn.Softmax(2)(cur_arg_logit)
                cur_arg_tags = bio.filter_arg_tags(cur_arg_tags, cur_pred_tags, token_str)

                # get string tuples and write results
                cur_extractions, cur_extraction_idxs = bio.get_tuple(sentence, cur_pred_tags, cur_arg_tags, tokenizer)
                cur_confidences = bio.get_confidence_score(cur_pred_probs, cur_arg_probs, cur_extraction_idxs)
                for extraction, confidence in zip(cur_extractions, cur_confidences):
                    # print('\n')
                    # print("\t".join([sentence] + [str(1.0)] + extraction[:3]))
                    res_dict = {
                        'score': str(confidence),
                        'relation': extraction[0],
                        'arg1': extraction[1],
                        'arg2': extraction[2]
                    }
                    result.append(res_dict)
    # print("\nExtraction Done.\n")
    return result


class OIE:
    def __init__(self, model_path, BERT, batch_size):
        if torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        self.model = Multi2OIE(bert_config=BERT).to(self.device)
        model_weights = torch.load(model_path, map_location=torch.device(self.device))
        self.model.load_state_dict(model_weights)
        self.max_len = 64
        self.batch_size = batch_size
        self.tokenizer_config = BERT
        self.tokenizer = BertTokenizer.from_pretrained(BERT)

    def predict(self, sentences):
        loader = DataLoader(
            dataset=OieEvalDataset(
                sentences,
                self.max_len,
                self.tokenizer_config),
            batch_size=self.batch_size,
            num_workers=4,
            pin_memory=True)
        return extract(model=self.model, loader=loader, tokenizer=self.tokenizer, device=self.device)
