from collections import OrderedDict
import torch
import torch.nn as nn
from torch.nn import functional as F

class Embedding(nn.Module):
    def __init__(self, num_ent, parameter):
        super(Embedding, self).__init__()
        # self.device = torch.device('cuda:0')
        self.device = torch.device(parameter['device'])
        self.es = parameter['embed_dim']
        
        self.embedding = nn.Embedding(num_ent + 1, self.es)
        nn.init.xavier_uniform_(self.embedding.weight)

    def forward(self, triples):
        idx = [[[t[0], t[2]] for t in batch] for batch in triples]
        idx = torch.LongTensor(idx).to(self.device)
        return self.embedding(idx)

class MetaLearner(nn.Module):
    def __init__(self, K, embed_size=100, num_hidden1=500, num_hidden2=200, out_size=100, dropout_p=0.5):
        super(MetaLearner, self).__init__()
        self.embed_size = embed_size
        self.K = K
        # self.out_size = out_size
        # self.hidden_size = out_size
        self.out_size = embed_size
        self.hidden_size = embed_size
        self.rnn = nn.LSTM(embed_size,self.hidden_size,2,dropout=0.2)
        # nn.init.xavier_normal_(self.rnn.all_weights)

    def forward(self, inputs):
        size = inputs.shape
        x = torch.stack([inputs[:,0,0,:],inputs[:,0,1,:],inputs[:,1,1,:]],dim=1)
        x = x.transpose(0,1)

        _,(x,c) = self.rnn(x)
        x = x[-1]
        x = x.squeeze(0)

        return x.view(size[0], 1, 1, self.out_size)

class EmbeddingLearner(nn.Module):
    def __init__(self):
        super(EmbeddingLearner, self).__init__()

    def forward(self, h, t, r, pos_num):
        score = -torch.norm(h + r - t, 2, -1).squeeze(2)
        p_score = score[:, :pos_num]
        n_score = score[:, pos_num:]
        return p_score, n_score

def bpr_loss(p_scores, n_values,device):
    ratio = int(n_values.shape[1] / p_scores.shape[1])
    temp_pvalues = torch.tensor([],device=device)
    for i in range(p_scores.shape[1]):
        temp_pvalues = torch.cat((temp_pvalues, p_scores[:, i, None].expand(-1, ratio)), dim=1)

    d = torch.sub(temp_pvalues,n_values)
    t = F.logsigmoid(d)
    loss = -1 * (1.0/n_values.shape[1]) * t.sum(dim=1)
    loss = loss.sum(dim=0)
    return loss

def bpr_max_loss(p_scores, n_values,device):
    s = F.softmax(n_values,dim=1)
    ratio = int(n_values.shape[1] / p_scores.shape[1])
    temp_pvalues = torch.tensor([],device=device)
    for i in range(p_scores.shape[1]):
        temp_pvalues = torch.cat((temp_pvalues,p_scores[:,i,None].expand(-1,ratio)),dim=1)

    d = torch.sigmoid(torch.sub(temp_pvalues,n_values))
    t = torch.mul(s,d)
    loss = -1 * torch.log(t.sum(dim=1))
    loss = loss.sum()
    return loss

def top_loss(p_scores, n_values,device):
    ratio = int(n_values.shape[1] / p_scores.shape[1])
    temp_pvalues = torch.tensor([],device=device)
    for i in range(p_scores.shape[1]):
        temp_pvalues = torch.cat((temp_pvalues, p_scores[:, i, None].expand(-1, ratio)), dim=1)

    t1 = torch.sigmoid(torch.sub(n_values , temp_pvalues))
    t2 = torch.sigmoid(torch.pow(n_values,2))
    t = torch.add(t1,t2)
    t = t.sum(dim=1)
    loss = t / n_values.shape[1]
    loss = loss.sum(dim=0)
    return loss


class MetaTL(nn.Module):
    def __init__(self, itemnum, parameter):
        super(MetaTL, self).__init__()
        # self.device = torch.device(parameter['device'])
        self.device = parameter['device']
        self.beta = parameter['beta']
        # self.dropout_p = parameter['dropout_p']
        self.embed_dim = parameter['embed_dim']
        self.margin = parameter['margin']
        self.embedding = Embedding(itemnum, parameter)

        self.relation_learner = MetaLearner(parameter['K'] - 1, embed_size=self.embed_dim, num_hidden1=500,
                                                    num_hidden2=200, out_size=100, dropout_p=0)

        self.embedding_learner = EmbeddingLearner()
        # self.loss_func = nn.MarginRankingLoss(self.margin)
        self.loss_func = bpr_loss

        self.rel_q_sharing = dict()

    def split_concat(self, positive, negative):
        pos_neg_e1 = torch.cat([positive[:, :, 0, :],
                                negative[:, :, 0, :]], 1).unsqueeze(2)
        pos_neg_e2 = torch.cat([positive[:, :, 1, :],
                                negative[:, :, 1, :]], 1).unsqueeze(2)
        return pos_neg_e1, pos_neg_e2

    def forward(self, task, iseval=False, curr_rel=''):
        # transfer task string into embedding
        support, support_negative, query, negative = [self.embedding(t) for t in task]

        K = support.shape[1]              # num of K
        num_sn = support_negative.shape[1]  # num of support negative
        num_q = query.shape[1]              # num of query
        num_n = negative.shape[1]           # num of query negative

        rel = self.relation_learner(support)
        rel.retain_grad()

        rel_s = rel.expand(-1, K+num_sn, -1, -1)

        if iseval and curr_rel != '' and curr_rel in self.rel_q_sharing.keys():
            rel_q = self.rel_q_sharing[curr_rel]
        else:
            sup_neg_e1, sup_neg_e2 = self.split_concat(support, support_negative)

            p_score, n_score = self.embedding_learner(sup_neg_e1, sup_neg_e2, rel_s, K)

            # y = torch.Tensor([1]).to(self.device)
            self.zero_grad()

            # sorted,indecies = torch.sort(n_score, descending=True,dim=1)
            # n_values = sorted[:,0:p_score.shape[1]]

            # loss = self.loss_func(p_score, n_values, y)
            loss = self.loss_func(p_score,n_score,device=self.device)
            loss.backward(retain_graph=True)

            grad_meta = rel.grad
            rel_q = rel - self.beta*grad_meta
            self.rel_q_sharing[curr_rel] = rel_q

        rel_q = rel_q.expand(-1, num_q + num_n, -1, -1)

        que_neg_e1, que_neg_e2 = self.split_concat(query, negative)  
        p_score, n_score = self.embedding_learner(que_neg_e1, que_neg_e2, rel_q, num_q)

        return p_score, n_score