import torch
from modules import MSA, BiLSTM, GraphTrans  
from utlis import *
from torch import nn
import dgl
            

class GraphWriter(nn.Module):
    def __init__(self, args):
        super(GraphWriter, self).__init__()
        self.args = args
        if args.title:
            self.title_emb = nn.Embedding(len(args.title_vocab), args.nhid, padding_idx=0)
            self.title_enc = BiLSTM(args, enc_type='title')
            self.title_attn = MSA(args)
        self.ent_emb = nn.Embedding(len(args.ent_text_vocab), args.nhid, padding_idx=0)
        self.tar_emb = nn.Embedding(len(args.text_vocab), args.nhid, padding_idx=0)
        if args.title:
            nn.init.xavier_normal_(self.title_emb.weight)
        nn.init.xavier_normal_(self.ent_emb.weight)
        self.rel_emb = nn.Embedding(len(args.rel_vocab), args.nhid, padding_idx=0)
        nn.init.xavier_normal_(self.rel_emb.weight)
        self.decode_lstm = nn.LSTMCell(args.dec_ninp, args.nhid)
        self.ent_enc = BiLSTM(args, enc_type='entity')
        self.graph_enc = GraphTrans(args)
        self.ent_attn = MSA(args)
        self.copy_attn = MSA(args, mode='copy')
        self.copy_fc = nn.Linear(args.dec_ninp, 1)
        self.pred_v_fc = nn.Linear(args.dec_ninp, len(args.text_vocab))

    def enc_forward(self, batch, ent_mask, ent_text_mask, ent_len, rel_mask, title_mask):
        title_enc = None
        if self.args.title:
            title_enc = self.title_enc(self.title_emb(batch['title']), title_mask)
        ent_enc = self.ent_enc(self.ent_emb(batch['ent_text']), ent_text_mask, ent_len = batch['ent_len'])
        rel_emb = self.rel_emb(batch['rel'])
        g_ent, g_root = self.graph_enc(ent_enc, ent_mask, ent_len, rel_emb, rel_mask, batch['graph'])
        return g_ent, g_root, title_enc, ent_enc 

    def forward(self, batch, beam_size=-1):
        ent_mask = len2mask(batch['ent_len'], self.args.device)
        ent_text_mask = batch['ent_text']==0
        rel_mask = batch['rel']==0 # 0 means the <PAD>
        title_mask = batch['title']==0
        g_ent, g_root, title_enc, ent_enc = self.enc_forward(batch, ent_mask, ent_text_mask, batch['ent_len'], rel_mask, title_mask)

        _h, _c = g_root, g_root.clone().detach()
        ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
        if self.args.title:
            attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
            ctx = torch.cat([ctx, attn], 1)
        if beam_size<1:
            # training
            outs = []
            tar_inp = self.tar_emb(batch['text'].transpose(0,1))
            for t, xt in enumerate(tar_inp):
                _xt = torch.cat([ctx, xt], 1)
                _h, _c = self.decode_lstm(_xt, (_h, _c))
                ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
                if self.args.title:
                    attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
                    ctx = torch.cat([ctx, attn], 1)
                outs.append(torch.cat([_h, ctx], 1)) 
            outs = torch.stack(outs, 1)
            copy_gate = torch.sigmoid(self.copy_fc(outs))
            EPSI = 1e-6
            # copy
            pred_v = torch.log(copy_gate+EPSI) + torch.log_softmax(self.pred_v_fc(outs), -1)
            pred_c = torch.log((1. - copy_gate)+EPSI) + torch.log_softmax(self.copy_attn(outs, ent_enc, mask=ent_mask), -1)
            pred = torch.cat([pred_v, pred_c], -1)
            return pred
        else:
            if beam_size==1:
                # greedy
                device = g_ent.device
                B = g_ent.shape[0]
                ent_type = batch['ent_type'].view(B, -1)
                seq = (torch.ones(B,).long().to(device) * self.args.text_vocab('<BOS>')).unsqueeze(1)
                for t in range(self.args.beam_max_len):
                    _inp = replace_ent(seq[:,-1], ent_type, len(self.args.text_vocab))
                    xt = self.tar_emb(_inp)
                    _xt = torch.cat([ctx, xt], 1)
                    _h, _c = self.decode_lstm(_xt, (_h, _c))
                    ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
                    if self.args.title:
                        attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
                        ctx = torch.cat([ctx, attn], 1)
                    _y = torch.cat([_h, ctx], 1)
                    copy_gate = torch.sigmoid(self.copy_fc(_y))
                    pred_v = torch.log(copy_gate) + torch.log_softmax(self.pred_v_fc(_y), -1)
                    pred_c = torch.log((1. - copy_gate)) + torch.log_softmax(self.copy_attn(_y.unsqueeze(1), ent_enc, mask=ent_mask).squeeze(1), -1)
                    pred = torch.cat([pred_v, pred_c], -1).view(B,-1)
                    for ban_item in ['<BOS>', '<PAD>', '<UNK>']:
                        pred[:, self.args.text_vocab(ban_item)] = -1e8
                    _, word = pred.max(-1)
                    seq = torch.cat([seq, word.unsqueeze(1)], 1)
                return seq
            else:
                # beam search
                device = g_ent.device
                B = g_ent.shape[0]
                BSZ = B * beam_size
                _h = _h.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                _c = _c.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                ent_mask = ent_mask.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                if self.args.title:
                    title_mask = title_mask.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                    title_enc = title_enc.view(B, 1, title_enc.size(1), -1).repeat(1, beam_size, 1, 1).view(BSZ, title_enc.size(1), -1)
                ctx = ctx.view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                ent_type = batch['ent_type'].view(B, 1, -1).repeat(1, beam_size, 1).view(BSZ, -1)
                g_ent = g_ent.view(B, 1, g_ent.size(1), -1).repeat(1, beam_size, 1, 1).view(BSZ, g_ent.size(1), -1)
                ent_enc = ent_enc.view(B, 1, ent_enc.size(1), -1).repeat(1, beam_size, 1, 1).view(BSZ, ent_enc.size(1), -1)

                beam_best = torch.zeros(B).to(device) - 1e9
                beam_best_seq = [None] * B 
                beam_seq = (torch.ones(B, beam_size).long().to(device) * self.args.text_vocab('<BOS>')).unsqueeze(-1)
                beam_score = torch.zeros(B, beam_size).to(device)
                done_flag = torch.zeros(B, beam_size)
                for t in range(self.args.beam_max_len):
                    _inp = replace_ent(beam_seq[:,:,-1].view(-1), ent_type, len(self.args.text_vocab))
                    xt = self.tar_emb(_inp)
                    _xt = torch.cat([ctx, xt], 1)
                    _h, _c = self.decode_lstm(_xt, (_h, _c))
                    ctx = _h + self.ent_attn(_h, g_ent, mask=ent_mask)
                    if self.args.title:
                        attn = _h + self.title_attn(_h, title_enc, mask=title_mask)
                        ctx = torch.cat([ctx, attn], 1)
                    _y = torch.cat([_h, ctx], 1)
                    copy_gate = torch.sigmoid(self.copy_fc(_y))
                    pred_v = torch.log(copy_gate) + torch.log_softmax(self.pred_v_fc(_y), -1)
                    pred_c = torch.log((1. - copy_gate)) + torch.log_softmax(self.copy_attn(_y.unsqueeze(1), ent_enc, mask=ent_mask).squeeze(1), -1)
                    pred = torch.cat([pred_v, pred_c], -1).view(B, beam_size, -1)
                    for ban_item in ['<BOS>', '<PAD>', '<UNK>']:
                        pred[:, :, self.args.text_vocab(ban_item)] = -1e8
                    if t==self.args.beam_max_len-1: # force ending 
                        tt = pred[:, :, self.args.text_vocab('<EOS>')]
                        pred = pred*0-1e8
                        pred[:, :, self.args.text_vocab('<EOS>')] = tt
                    cum_score = beam_score.view(B,beam_size,1) + pred
                    score, word = cum_score.topk(dim=-1, k=beam_size) # B, beam_size, beam_size
                    score, word = score.view(B,-1), word.view(B,-1)
                    eos_idx = self.args.text_vocab('<EOS>')
                    if beam_seq.size(2)==1:
                        new_idx = torch.arange(beam_size).to(word)
                        new_idx = new_idx[None,:].repeat(B,1)
                    else:
                        _, new_idx = score.topk(dim=-1, k=beam_size)
                    new_src, new_score, new_word, new_done = [], [], [], []
                    LP = beam_seq.size(2) ** self.args.lp
                    for i in range(B):
                        for j in range(beam_size):
                            tmp_score = score[i][new_idx[i][j]]
                            tmp_word = word[i][new_idx[i][j]]
                            src_idx = new_idx[i][j]//beam_size
                            new_src.append(src_idx)
                            if tmp_word == eos_idx:
                                new_score.append(-1e8)
                            else:
                                new_score.append(tmp_score)
                            new_word.append(tmp_word)

                            if tmp_word == eos_idx and done_flag[i][src_idx]==0 and tmp_score/LP>beam_best[i]:
                                beam_best[i] = tmp_score/LP
                                beam_best_seq[i] = beam_seq[i][src_idx]
                            if tmp_word == eos_idx:
                                new_done.append(1)
                            else:
                                new_done.append(done_flag[i][src_idx])
                    new_score = torch.Tensor(new_score).view(B,beam_size).to(beam_score)
                    new_word = torch.Tensor(new_word).view(B,beam_size).to(beam_seq)
                    new_src = torch.LongTensor(new_src).view(B,beam_size).to(device)
                    new_done = torch.Tensor(new_done).view(B,beam_size).to(done_flag)
                    beam_score = new_score
                    done_flag = new_done
                    beam_seq = beam_seq.view(B,beam_size,-1)[torch.arange(B)[:,None].to(device), new_src]
                    beam_seq = torch.cat([beam_seq, new_word.unsqueeze(2)], 2)
                    _h = _h.view(B,beam_size,-1)[torch.arange(B)[:,None].to(device), new_src].view(BSZ,-1)
                    _c = _c.view(B,beam_size,-1)[torch.arange(B)[:,None].to(device), new_src].view(BSZ,-1)
                    ctx = ctx.view(B,beam_size,-1)[torch.arange(B)[:,None].to(device), new_src].view(BSZ,-1)

                return beam_best_seq