import torch
import torch.nn as nn
import numpy as np
from config import GlobalConfig
from config.model_config import TextDecoderConfig, KnowledgeTextDecoderConfig
from transformers import BartTokenizer, BartModel, BartForConditionalGeneration, AutoConfig
from transformers.modeling_outputs import ModelOutput
from constant import SOS_ID, EOS_ID, PAD_ID
from model.modeling_bart import BartForMultiModalGeneration

import warnings

warnings.filterwarnings('ignore')


class TextDecoder(nn.Module):
    def __init__(self, config: TextDecoderConfig, emb_size=300):
        super(TextDecoder, self).__init__()
        self.tokenizer = BartTokenizer.from_pretrained(
            '/data/chenxiaolin/bart')
        # self.bart = BartForConditionalGeneration(AutoConfig.from_pretrained('/home/share/chenxiaolin/MultimodalDialogSystem/Model_data_dump5/bart'))
        self.bart = BartForMultiModalGeneration.from_pretrained(
            '/data/chenxiaolin/bart',
            fusion_layer=5,  # default=5, 'number of fusion layers'
            use_img_trans=False,  # whether or not to use VTF
            use_forget_gate=True,  # whether or not to use forget gate
            cross_attn_type=0,  # 0
            dim_common=256,  # default=256
            n_attn_heads=1)  # default=1) # only needed when cross_attn_type is 4 or 5
        self.bart_config = self.bart.config
        self.criterion = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id)
        self.linear = nn.Linear(emb_size, 1024)

        self.linear = self.linear.to(GlobalConfig.device)
        self.criterion = self.criterion.to(GlobalConfig.device)
        self.mu_mt_prior = nn.Linear(1024, 128)  # 论文appendix提到
        self.mu_mt_post = nn.Linear(2048, 128)

    def softplus(self, x):
        return torch.log(1.0 + torch.exp(x))

    def sample_z(self, mu, logvar):
        eps = torch.rand_like(
            mu)  # Returns a tensor with the same size as input that is filled with random numbers from a uniform distribution on the interval [0, 1)[0,1).
        return mu + eps * torch.exp(0.5 * logvar)

    def target_id2word(self, target, id2word):
        batch_target = []

        for i in range(target.size(0)):
            each_target = []
            for j in range(20):
                each_id = target[i][j]
                if each_id != PAD_ID and each_id != EOS_ID:
                    each_word = id2word[each_id]
                    if isinstance(each_word, str):
                        each_target.append(each_word.replace("'", ""))
                    else:
                        print(each_word)
                        print(type(each_word))
                        exit()

            each_target1 = " ".join(each_target).replace(' .', '.')
            each_target1 = each_target1.replace(' ,', ',')  # ?
            each_target1 = each_target1.replace(' ?', '?')  #
            each_target1 = each_target1.replace(' !', '!')  #
            batch_target.append(each_target1)

        return batch_target

    def shift_tokens_right(self, input_ids, token_id):

        prev_output_tokens = input_ids.clone()
        prev_output_tokens[:, 0] = torch.ones_like(prev_output_tokens[:, 0],
                                                   device=prev_output_tokens.device) * token_id
        prev_output_tokens[:, 1:] = input_ids[:, :-1]
        return prev_output_tokens

    def save_file(self, str_true_batch, str_pred_batch, output_file):
        batch_size = len(str_pred_batch)
        for j in range(batch_size):
            str_pred = [str_pred_batch[j]]
            str_true = [str_true_batch[j]]

            line = "{}\t{}".format(' '.join(str_pred), ' '.join(str_true))
            if output_file:
                output_file.write(line + '\n')
            else:
                print(line)

    def tokenize_input(self, text, max_len, tokenizer):
        tokens = tokenizer.batch_encode_plus(text, max_length=max_len, pad_to_max_length=True, truncation=True,
                                             return_tensors='pt', add_special_tokens=False)
        return tokens['input_ids'], tokens['attention_mask']

    def tokenize_visual_input(self, text, max_len, tokenizer):
        context_batch = []
        # each_text0 = text[:, 0, :]
        # tokens0 = tokenizer.batch_encode_plus(each_text0, max_length=max_len, pad_to_max_length=True, truncation=True,
        #                                      return_tensors='pt', add_special_tokens=False)
        # each_text1 = text[:, 1, :]
        # tokens1 = tokenizer.batch_encode_plus(each_text1, max_length=max_len, pad_to_max_length=True, truncation=True,
        #                                      return_tensors='pt', add_special_tokens=False)

        for context_index in range(len(text)):  # 16
            # context_visual_ids = []
            for img_index in range(len(text[0])):
                each_img_text = text[context_index][img_index]
                tokens = tokenizer.encode(each_img_text, max_length=max_len, pad_to_max_length=True, truncation=True, return_tensors='pt', add_special_tokens=False)
                if img_index == 0:
                    context_visual_ids = tokens
                else:
                    context_visual_ids = torch.cat([context_visual_ids, tokens], dim=0)
                # context_visual_ids.append(tokens)
            if context_index == 0:
                context_batch = torch.unsqueeze(context_visual_ids, 0)
            else:
                context_batch = torch.cat([context_batch, torch.unsqueeze(context_visual_ids, 0)], dim=0)
            # context_batch.append(context_visual_ids)
        return context_batch

        # return tokens['input_ids'], tokens['attention_mask']

    def tokenize_output(self, texts, max_len, tokenizer, bos_id=0, eos_id=2, pad_id=1):
        all_tokens = []
        all_mask = []
        for text in texts:
            # tokens = [bos_id] + tokenizer.encode(text, add_special_tokens=False)
            tokens = tokenizer.encode(text, add_special_tokens=False)
            if len(tokens) >= max_len:
                tokens = tokens[:max_len - 1]
            tokens = tokens + [eos_id]
            mask = [1] * len(tokens) + [0] * (max_len - len(tokens))
            tokens = tokens + [pad_id] * (max_len - len(tokens))
            all_tokens.append(tokens)
            all_mask.append(mask)
            # print(tokens)
            # print(tokenizer.decode(tokens, skip_special_tokens=True))
            # exit(0)
        return torch.tensor(all_tokens), torch.tensor(all_mask)

    def get_txt_encoder(self, input_ids, attention_mask):
        output_attentions = self.bart_config.output_attentions
        output_hidden_states = (
            self.bart_config.output_hidden_states
        )
        encoder_outputs = self.bart.model.encoder(
            input_ids=input_ids,
            attention_mask=attention_mask,
            head_mask=None,
            inputs_embeds=None,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=None,
        )
        return encoder_outputs[0]

    def get_encoder_output(self, input_ids, attention_mask):
        all_input_ids = input_ids
        all_attention_mask = attention_mask
        encoder_outputs = self.get_txt_encoder(input_ids=all_input_ids, attention_mask=all_attention_mask)
        return encoder_outputs

    def text_generation(self, context, context1, target, id2word, output_file, image_features, image_len, context_visual_know, text_visual_know, beam_size=10):

        input_ids, attention_mask = self.tokenize_input(context1, 170, self.tokenizer)
        input_ids = input_ids.to(GlobalConfig.device)
        attention_mask = attention_mask.to(GlobalConfig.device)

        visual_know_ids = self.tokenize_visual_input(context_visual_know, 70, self.tokenizer)  # 16,2,70
        visual_know_ids = visual_know_ids.to(GlobalConfig.device)
        # print('target_test:', target)

        knowledge_ids, knowledge_mask = self.tokenize_input(text_visual_know, 100, self.tokenizer)  # 16,2,70--100
        knowledge_ids = knowledge_ids.to(GlobalConfig.device)
        knowledge_mask = knowledge_mask.to(GlobalConfig.device)

        context_input_ids, context_attention_mask = self.tokenize_input(context, 70, self.tokenizer)  #170---250
        context_input_ids = context_input_ids.to(GlobalConfig.device)
        context_attention_mask = context_attention_mask.to(GlobalConfig.device)

        batch_target = self.target_id2word(target, id2word)
        simple_target = batch_target

        decoder_ids, decoder_attention_mask = self.tokenize_output(simple_target, 50, self.tokenizer)

        decoder_ids = decoder_ids.to(GlobalConfig.device)
        decoder_attention_mask = decoder_attention_mask.to(GlobalConfig.device)

        generate_ids = self.bart.generate(input_ids=input_ids,
                                          attention_mask=attention_mask,
                                          num_beams=beam_size,
                                          max_length=20,
                                          early_stopping=True,
                                          no_repeat_ngram_size=3,
                                          image_features=image_features.float(),
                                          image_len=image_len,
                                          use_cache=True,
                                          decoder_start_token_id=self.bart_config.bos_token_id,
                                          visual_know_ids=visual_know_ids,
                                          knowledge_id=knowledge_ids,
                                          knowledge_mask=knowledge_mask,
                                          context_input_ids=context_input_ids,
                                          context_attention_mask=context_attention_mask
                                          )

        samples = [self.tokenizer.decode(w, skip_special_tokens=True, clean_up_tokenization_spaces=True) for w in
                   generate_ids]

        candidate = samples

        self.save_file(simple_target, candidate, output_file)

    def forward(self, context, context1, target, id2word, output_file, image_features, image_len, context_visual_know, text_visual_know):
        # 处理encoder
        # input_ids, attention_mask = self.tokenize_input(context1, 100, self.tokenizer)
        input_ids, attention_mask = self.tokenize_input(context1, 170, self.tokenizer)  #170---250
        input_ids = input_ids.to(GlobalConfig.device)
        attention_mask = attention_mask.to(GlobalConfig.device)

        visual_know_ids = self.tokenize_visual_input(context_visual_know, 70, self.tokenizer)  # 16,2,70--100
        visual_know_ids = visual_know_ids.to(GlobalConfig.device)

        knowledge_ids, knowledge_mask = self.tokenize_input(text_visual_know, 100, self.tokenizer)  # 16,2,70--100
        knowledge_ids = knowledge_ids.to(GlobalConfig.device)
        knowledge_mask = knowledge_mask.to(GlobalConfig.device)

        context_input_ids, context_attention_mask = self.tokenize_input(context, 70, self.tokenizer)  #170---250
        context_input_ids = context_input_ids.to(GlobalConfig.device)
        context_attention_mask = context_attention_mask.to(GlobalConfig.device)


        # target id2word
        batch_target = self.target_id2word(target, id2word)
        simple_target = batch_target
        decoder_ids, decoder_attention_mask = self.tokenize_output(simple_target, 50, self.tokenizer) # 20--50

        decoder_ids = decoder_ids.to(GlobalConfig.device)
        decoder_attention_mask = decoder_attention_mask.to(GlobalConfig.device)
        decoder_input_ids = self.shift_tokens_right(decoder_ids, self.bart_config.bos_token_id)

        decoder_outputs = self.bart(input_ids=input_ids,
                                    attention_mask=attention_mask,
                                    decoder_input_ids=decoder_input_ids,
                                    decoder_attention_mask=decoder_attention_mask,
                                    image_features=image_features.float(),
                                    image_len=image_len,
                                    visual_know_ids=visual_know_ids,
                                    knowledge_id=knowledge_ids,
                                    knowledge_mask=knowledge_mask,
                                    context_input_ids=context_input_ids,
                                    context_attention_mask=context_attention_mask
                                    )  # [0]

        decoder_outputs_encoder = self.bart.model.encoder(input_ids=decoder_ids,  # 16,50,1024
                                    attention_mask=decoder_attention_mask
                                    # decoder_input_ids=decoder_input_ids,
                                    # decoder_attention_mask=decoder_attention_mask,
                                    # image_features=image_features.float(),
                                    # image_len=image_len,
                                    # visual_know_ids=visual_know_ids,
                                    # knowledge_id=knowledge_ids,
                                    # knowledge_mask=knowledge_mask,
                                    # context_input_ids=context_input_ids,
                                    # context_attention_mask=context_attention_mask
                                    )  # [0]

        encoder_outputs = decoder_outputs.encoder_last_hidden_state  # 16,170,1024
        x_rep = torch.mean(encoder_outputs, dim=1, keepdim=False)  # encoder最终表征
        mu_mt_prior, logvar_mt_prior = self.mu_mt_prior(x_rep), self.softplus(self.mu_mt_prior(x_rep))
        y_out = decoder_outputs_encoder.last_hidden_state
        y_rep = torch.mean(y_out, dim=1, keepdim=False)
        xy_rep = torch.cat((x_rep, y_rep), -1)
        mu_mt_post, logvar_mt_post = self.mu_mt_post(xy_rep), self.softplus(self.mu_mt_post(xy_rep))  # self.logvar_mt_post(xy_rep)
        z_mt = self.sample_z(mu_mt_post, logvar_mt_post)
        # kl_loss
        kld1 = self.gaussian_kld(mu_mt_post, logvar_mt_post, mu_mt_prior, logvar_mt_prior)
        kl_mt_loss = torch.mean(kld1, dim=-1, keepdim=False)


        outputs_1 = decoder_outputs.logits

        loss_ce = self.criterion(outputs_1.view(-1, outputs_1.shape[-1]), decoder_ids.view(-1))

        # prediction
        probs = outputs_1.softmax(dim=-1)
        values, predictions = probs.topk(1)
        predictions = predictions.squeeze()

        str_pred_batch = self.tokenizer.batch_decode(predictions, skip_special_tokens=True,
                                                     clean_up_tokenization_spaces=False)

        self.save_file(simple_target, str_pred_batch, output_file)

        return loss_ce

        # embed: torch.Size([32, 300])
        # enc_hiddens: torch.Size([32, 300])
        # hidden: torch.Size([32, 512])
        # embed.unsqueeze(0): torch.Size([1, 32, 300])
        # hidden.unsqueeze(0): torch.Size([1, 32, 512])
        embed1 = self.embedding(word)
        # print('embed1', np.shape(embed1))
        # print('enc_hiddens', np.shape(enc_hiddens))

        embed = self.fc_know(torch.cat((embed1, enc_hiddens), dim=1))
        # print('embed3', np.shape(embed))

        # print('embed:', np.shape(embed))
        # print('enc_hiddens:', np.shape(enc_hiddens))
        # print('hidden:', np.shape(hidden))
        # print('embed.unsqueeze(0):', np.shape(embed.unsqueeze(0)))
        # print('hidden.unsqueeze(0):', np.shape(hidden.unsqueeze(0)))

        # if encode_knowledge_func is not None:
        #     knowledge = encode_knowledge_func(hidden, batch_id)
        #     # enc_hiddens: (seq_len, batch_size, hidden_size)
        #     # hidden: (batch_size, hidden_size)
        #     attention = torch.matmul(enc_hiddens.unsqueeze(2),
        #                              hidden.unsqueeze(0).unsqueeze(3))
        #     attention = attention.squeeze(2).squeeze(2)
        #     attention = attention.transpose(0, 1)
        #     attention = self.softmax(attention)
        #     enc_hidden = torch.matmul(attention.unsqueeze(1),
        #                               enc_hiddens.transpose(0, 1)).squeeze(1)
        #     embed = self.fc(torch.cat((embed, knowledge, enc_hidden), dim=1))
        # 当前时刻输入，和上一时刻输出
        # unsqueeze--1,32,300
        # embed.unsqueeze(0): torch.Size([1, 64, 300])
        # hidden.unsqueeze(0): torch.Size([1, 64, 512])        
        # print('embed.unsqueeze(0):', np.shape(embed.unsqueeze(0)))
        # print('hidden.unsqueeze(0):', np.shape(hidden.unsqueeze(0)))
        output1, hidden1 = self.gru(embed.unsqueeze(0), hidden.unsqueeze(0))
        # output1 torch.Size([1, 64, 512])
        # hidden1 torch.Size([1, 64, 512])        
        # print('output1', np.shape(output1))
        # print('hidden1', np.shape(hidden1))
        output2 = self.linear(output1[0])
        # output2 torch.Size([64, 4892])
        # output2 torch.Size([64, 4892])        
        # print('output2', np.shape(output2))
        output3 = self.softmax(output2)
        # print('output3: ', output3)
        # print('output2', np.shape(output3))

        hidden2 = hidden1[0]
        # print('hidden', np.shape(hidden2))
        return output3, hidden2
        # return output3, hidden


class ToHidden(nn.Module):
    def __init__(self, config: TextDecoderConfig):
        super(ToHidden, self).__init__()
        self.fcs = []
        # len(config.to_hidden_fc_sizes)--2---1
        for i in range(len(config.to_hidden_fc_sizes) - 1):
            linear = nn.Linear(config.to_hidden_fc_sizes[i],
                               config.to_hidden_fc_sizes[i + 1])
            linear = linear.to(GlobalConfig.device)
            self.fcs.append(linear)

    # 其实这个ToHidden是一层的MLP
    def forward(self, x):
        for fc in self.fcs:
            # x1 = fc(x)
            x = fc(x.clone())
        return x
