import logging
import os

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn

from hopwise.model.abstract_recommender import KnowledgeRecommender
from hopwise.utils import InputType


class GraphAttentionLayer(nn.Module):
    def __init__(self, in_features, out_features, dropout, alpha, concat=True):
        super().__init__()
        self.dropout = dropout
        self.in_features = in_features
        self.out_features = out_features
        self.alpha = alpha
        self.concat = concat

        self.W = nn.Parameter(torch.empty(size=(in_features, out_features)))
        nn.init.xavier_uniform_(self.W.data, gain=1.414)
        self.a = nn.Parameter(torch.empty(size=(2 * out_features, 1)))
        nn.init.xavier_uniform_(self.a.data, gain=1.414)
        self.fc = nn.Linear(2 * out_features, out_features)

        self.leakyrelu = nn.LeakyReLU(self.alpha)

    def forward_relation(self, item_embs, entity_embs, relations, adj):
        # item_embs: N, dim
        # entity_embs: N, e_num, dim
        # relations: N, e_num, r_dim
        # adj: N, e_num

        # N, e_num, dim
        Wh = item_embs.unsqueeze(1).expand(entity_embs.size())
        # N, e_num, dim
        We = entity_embs
        a_input = torch.cat((Wh, We), dim=-1)  # (N, e_num, 2*dim)
        # N,e,2dim -> N,e,dim
        e_input = torch.multiply(self.fc(a_input), relations).sum(-1)  # N,e
        e = self.leakyrelu(e_input)  # (N, e_num)

        zero_vec = -9e15 * torch.ones_like(e)
        attention = torch.where(adj > 0, e, zero_vec)
        attention = F.softmax(attention, dim=1)
        attention = F.dropout(attention, self.dropout, training=self.training)  # N, e_num
        # (N, 1, e_num) * (N, e_num, out_features) -> N, out_features
        entity_emb_weighted = torch.bmm(attention.unsqueeze(1), entity_embs).squeeze()
        h_prime = entity_emb_weighted + item_embs

        if self.concat:
            return F.elu(h_prime)
        else:
            return h_prime

    def forward(self, item_embs, entity_embs, adj):
        Wh = torch.mm(item_embs, self.W)  # h.shape: (N, in_features), Wh.shape: (N, out_features)
        We = torch.matmul(
            entity_embs, self.W
        )  # entity_embs: (N, e_num, in_features), We.shape: (N, e_num, out_features)
        a_input = self._prepare_cat(Wh, We)  # (N, e_num, 2*out_features)
        e = self.leakyrelu(torch.matmul(a_input, self.a).squeeze(2))  # (N, e_num)

        zero_vec = -9e15 * torch.ones_like(e)
        attention = torch.where(adj > 0, e, zero_vec)
        attention = F.softmax(attention, dim=1)
        attention = F.dropout(attention, self.dropout, training=self.training)  # N, e_num
        # (N, 1, e_num) * (N, e_num, out_features) -> N, out_features
        entity_emb_weighted = torch.bmm(attention.unsqueeze(1), entity_embs).squeeze()
        h_prime = entity_emb_weighted + item_embs

        if self.concat:
            return F.elu(h_prime)
        else:
            return h_prime

    def _prepare_cat(self, Wh, We):
        Wh = Wh.unsqueeze(1).expand(We.size())  # (N, e_num, out_features)
        return torch.cat((Wh, We), dim=-1)  # (N, e_num, 2*out_features)


class KGEncoder(nn.Module):
    def __init__(self, config, dataset, kg_dataset):
        super().__init__()

        self.user_history_matrix = dataset.history_item_matrix()[0].to(config["device"])

        self.maxhis = config["maxhis"]
        self.kgcn = config["kgcn"]
        self.dropout = config["dropout"]
        self.keep_prob = 1 - self.dropout  # Added
        self.A_split = config["A_split"]
        self.device = config["device"]

        self.latent_dim = config["latent_dim_rec"]
        self.n_layers = config["lightGCN_n_layers"]
        self.max_entities_per_user = config["max_entities_per_user"]
        self.kg_dataset = kg_dataset
        self.gat = GAT(self.latent_dim, self.latent_dim, dropout=0.4, alpha=0.2).train()

        self.inter_feat = dataset.inter_feat
        self.num_users = dataset.user_num
        self.num_items = dataset.item_num

        self.__init_weight(dataset)
        self.config = config

    def __init_weight(self, dataset):
        self.entity_count = dataset.entity_num
        self.relation_count = dataset.relation_num

        self.embedding_user = torch.nn.Embedding(num_embeddings=self.num_users, embedding_dim=self.latent_dim)
        # item and kg entity
        self.embedding_item = torch.nn.Embedding(num_embeddings=self.num_items, embedding_dim=self.latent_dim)
        self.embedding_entity = torch.nn.Embedding(num_embeddings=self.entity_count + 1, embedding_dim=self.latent_dim)
        self.embedding_relation = torch.nn.Embedding(
            num_embeddings=self.relation_count + 1, embedding_dim=self.latent_dim
        )
        # relation weights
        self.W_R = nn.Parameter(torch.Tensor(self.relation_count, self.latent_dim, self.latent_dim))
        nn.init.xavier_uniform_(self.W_R, gain=nn.init.calculate_gain("relu"))

        nn.init.normal_(self.embedding_user.weight, std=0.1)
        nn.init.normal_(self.embedding_item.weight, std=0.1)
        nn.init.normal_(self.embedding_entity.weight, std=0.1)
        nn.init.normal_(self.embedding_relation.weight, std=0.1)

        self.f = nn.Sigmoid()
        self.Graph = dataset.norm_adjacency_matrix().coalesce().to(self.device)
        self.kg_dict, self.item2relations = self.get_kg_dict(self.num_items)

    def get_kg_dict(self, item_num):
        i2es = dict()
        i2rs = dict()
        for item in range(item_num):
            rts = self.kg_dataset.get(item, False)
            if rts:
                tails = list(set([ent for tail_list in rts.values() for ent in tail_list]))
                relations = list(rts.keys())
                if len(tails) > self.max_entities_per_user:
                    i2es[item] = torch.IntTensor(tails).to(self.device)[: self.max_entities_per_user]
                    i2rs[item] = torch.IntTensor(relations).to(self.device)[: self.max_entities_per_user]
                else:
                    # last embedding pos as padding idx
                    tails.extend([self.dataset.entity_count] * (self.max_entities_per_user - len(tails)))
                    relations.extend([self.dataset.relation_count] * (self.max_entities_per_user - len(relations)))
                    i2es[item] = torch.IntTensor(tails).to(self.device)
                    i2rs[item] = torch.IntTensor(relations).to(self.device)
            else:
                i2es[item] = torch.IntTensor([self.num_items] * self.max_entities_per_user).to(self.device)
                i2rs[item] = torch.IntTensor([self.relation_count] * self.max_entities_per_user).to(self.device)
        return i2es, i2rs

    def computer(self):
        with torch.no_grad():
            users_emb = self.embedding_user.weight
            items_emb = self.cal_item_embedding_from_kg(self.kg_dict)
            all_emb = torch.cat([users_emb, items_emb])
            embs = [all_emb]
            if self.dropout:
                if self.training:
                    g_droped = self.__dropout(self.keep_prob)
                else:
                    g_droped = self.Graph
            else:
                g_droped = self.Graph

            for layer in range(self.n_layers):
                all_emb = torch.sparse.mm(g_droped, all_emb)
                embs.append(all_emb)

            embs = torch.stack(embs, dim=1)
            light_out = torch.mean(embs, dim=1)
            users, items = torch.split(light_out, [self.num_users, self.num_items])
            return users, items

    def __dropout_x(self, x, keep_prob):
        size = x.size()
        index = x.indices().t()
        values = x.values()
        random_index = torch.rand(len(values)) + keep_prob
        random_index = random_index.int().bool()
        index = index[random_index]
        values = values[random_index] / keep_prob
        g = torch.sparse_coo_tensor(index.t(), values, size)
        return g

    def __dropout(self, keep_prob):
        if self.A_split:
            graph = []
            for g in self.Graph:
                graph.append(self.__dropout_x(g, keep_prob))
        else:
            graph = self.__dropout_x(self.Graph, keep_prob)
        return graph

    def cal_item_embedding_from_kg(self, kg: dict):
        if kg is None:
            kg = self.kg_dict

        if self.kgcn == "GAT":
            return self.cal_item_embedding_gat(kg)
        elif self.kgcn == "RGAT":
            return self.cal_item_embedding_rgat(kg)
        elif self.kgcn == "MEAN":
            raise NotImplementedError("The 'MEAN' option for kgcn is not yet implemented.")
        elif self.kgcn == "NO":
            return self.embedding_item.weight

    def cal_item_embedding_gat(self, kg: dict):
        item_embs = self.embedding_item(torch.IntTensor(list(kg.keys())).to(self.device))  # item_num, emb_dim
        # item_num, entity_num_each
        item_entities = torch.stack(list(kg.values()))
        # item_num, entity_num_each, emb_dim
        entity_embs = self.embedding_entity(item_entities)
        # item_num, entity_num_each
        padding_mask = torch.where(
            item_entities != self.entity_count, torch.ones_like(item_entities), torch.zeros_like(item_entities)
        ).float()
        return self.gat(item_embs, entity_embs, padding_mask)

    def cal_item_embedding_rgat(self, kg: dict):
        item_embs = self.embedding_item(torch.IntTensor(list(kg.keys())).to(self.device))  # item_num, emb_dim
        # item_num, entity_num_each
        item_entities = torch.stack(list(kg.values()))
        item_relations = torch.stack(list(self.item2relations.values()))
        # item_num, entity_num_each, emb_dim
        entity_embs = self.embedding_entity(item_entities)
        relation_embs = self.embedding_relation(item_relations)  # item_num, entity_num_each, emb_dim
        # w_r = self.W_R[relation_embs] # item_num, entity_num_each, emb_dim, emb_dim
        # item_num, entity_num_each
        padding_mask = torch.where(
            item_entities != self.entity_count, torch.ones_like(item_entities), torch.zeros_like(item_entities)
        ).float()
        return self.gat.forward_relation(item_embs, entity_embs, relation_embs, padding_mask)


class KGLRR(KnowledgeRecommender):
    """
    KGLRR: Reinforced logical reasoning over KGs for interpretable recommendation system
    """

    input_type = InputType.PAIRWISE

    def __init__(self, config, dataset) -> None:
        super().__init__(config, dataset)

        self.kg_dataset = dataset.ckg_dict_graph()

        self.encoder = KGEncoder(config, dataset, self.kg_dataset)
        self.latent_dim = config["latent_dim_rec"]

        self.r_logic = config["r_logic"]
        self.r_length = config["r_length"]
        self.layers = config["layers"]
        self.sim_scale = config["sim_scale"]
        self.loss_sum = config["loss_sum"]
        self.l2s_weight = config["l2_loss"]
        self.is_explain = config["explain"]

        self.num_items = dataset.item_num

        self._init_weights()
        self.bceloss = nn.BCEWithLogitsLoss()

    def _init_weights(self):
        self.true = torch.nn.Parameter(
            torch.from_numpy(np.random.uniform(0, 1, size=[1, self.latent_dim]).astype(np.float32)),
            requires_grad=False,
        )

        self.and_layer = torch.nn.Linear(self.latent_dim * 2, self.latent_dim)
        for i in range(self.layers):
            setattr(self, "and_layer_%d" % i, torch.nn.Linear(self.latent_dim * 2, self.latent_dim * 2))

        self.or_layer = torch.nn.Linear(self.latent_dim * 2, self.latent_dim)
        for i in range(self.layers):
            setattr(self, "or_layer_%d" % i, torch.nn.Linear(self.latent_dim * 2, self.latent_dim * 2))

    def logic_or(self, vector1, vector2, train=False):
        vector1, vector2 = self.uniform_size(vector1, vector2, train)
        vector = torch.cat((vector1, vector2), dim=-1)
        for i in range(self.layers):
            vector = F.relu(getattr(self, "or_layer_%d" % i)(vector))
        vector = self.or_layer(vector)
        return vector

    def logic_and(self, vector1, vector2, train=False):
        vector1, vector2 = self.uniform_size(vector1, vector2, train)
        vector = torch.cat((vector1, vector2), dim=-1)
        for i in range(self.layers):
            vector = F.relu(getattr(self, "and_layer_%d" % i)(vector))
        vector = self.and_layer(vector)
        return vector

    def logic_regularizer(self, train: bool, check_list: list, constraint, constraint_valid):
        # This function calculates the gap between logical expressions and the real world

        # length
        r_length = constraint.norm(dim=2).sum()
        check_list.append(("r_length", r_length))

        # and
        r_and_true = 1 - self.similarity(self.logic_and(constraint, self.true, train=train), constraint)
        r_and_true = (r_and_true * constraint_valid).sum()
        check_list.append(("r_and_true", r_and_true))

        r_and_self = 1 - self.similarity(self.logic_and(constraint, constraint, train=train), constraint)
        r_and_self = (r_and_self * constraint_valid).sum()
        check_list.append(("r_and_self", r_and_self))

        # or
        r_or_true = 1 - self.similarity(self.logic_or(constraint, self.true, train=train), self.true)
        r_or_true = (r_or_true * constraint_valid).sum()
        check_list.append(("r_or_true", r_or_true))

        r_or_self = 1 - self.similarity(self.logic_or(constraint, constraint, train=train), constraint)
        r_or_self = (r_or_self * constraint_valid).sum()
        check_list.append(("r_or_self", r_or_self))

        r_loss = r_and_true + r_and_self + r_or_true + r_or_self

        if self.r_logic > 0:
            r_loss = r_loss * self.r_logic
        else:
            r_loss = torch.from_numpy(np.array(0.0, dtype=np.float32)).to(self.device)
            r_loss.requires_grad = True

        r_loss += r_length * self.r_length
        check_list.append(("r_loss", r_loss))
        return r_loss

    def similarity(self, vector1, vector2, sigmoid=True):
        result = F.cosine_similarity(vector1, vector2, dim=-1)
        result = result * self.sim_scale
        if sigmoid:
            return result.sigmoid()
        return result

    def uniform_size(self, vector1, vector2, train=False):
        # Removed vector size normalization
        if len(vector1.size()) < len(vector2.size()):
            vector1 = vector1.expand_as(vector2)
        elif vector2.size() != vector1.size():
            vector2 = vector2.expand_as(vector1)
        if train:
            r12 = torch.Tensor(vector1.size()[:-1]).to(self.device).uniform_(0, 1).bernoulli()
            r12 = r12.unsqueeze(-1)
            new_v1 = r12 * vector1 + (-r12 + 1) * vector2
            new_v2 = r12 * vector2 + (-r12 + 1) * vector1
            return new_v1, new_v2
        return vector1, vector2

    def predict(self, interaction):
        users = interaction[self.USER_ID]

        history = self.encoder.user_history_matrix[users, : self.encoder.maxhis]  # B * H
        item_embed = self.encoder.computer()[1]  # item_num * V

        his_valid = history.ge(0).float()  # B * H

        maxlen = int(his_valid.sum(dim=1).max().item())

        elements = item_embed[history] * his_valid.unsqueeze(-1)  # B * H * V

        tmp_o = None
        for i in range(maxlen):
            tmp_o_valid = his_valid[:, i].unsqueeze(-1)
            if tmp_o is None:
                tmp_o = elements[:, i, :] * tmp_o_valid  # B * V
            else:
                # Only perform OR operation if valid; otherwise, if the history is not that long (not valid),
                # keep the original content unchanged
                tmp_o = self.logic_or(tmp_o, elements[:, i, :]) * tmp_o_valid + tmp_o * (-tmp_o_valid + 1)  # B * V
        or_vector = tmp_o  # B * V
        left_valid = his_valid[:, 0].unsqueeze(-1)  # B * 1

        prediction = []
        for i in range(users.size(0)):
            sent_vector = (
                left_valid[i] * self.logic_and(or_vector[i].unsqueeze(0).repeat(self.num_items, 1), item_embed)
                + (-left_valid[i] + 1) * item_embed
            )  # item_size * V
            ithpred = self.similarity(sent_vector, self.true, sigmoid=True)  # item_size
            prediction.append(ithpred)

        prediction = torch.stack(prediction).to(self.device)  # [B, item_size]

        if self.is_explain:
            topk_items = torch.topk(prediction, k=1, dim=1).indices  # shape: [B, 1]
            topk_embeds = self.encoder.computer()[1][topk_items.squeeze(1)]  # shape: [B, V]
            explanation = self.explain(users, history, topk_embeds)
            return prediction, explanation
        return prediction

    def explain(self, users, history, items):
        bs = users.size(0)
        _, item_embed = self.encoder.computer()  # user_num/item_num * V

        his_valid = history.ge(0).float()  # B * H
        elements = item_embed[history.abs()] * his_valid.unsqueeze(-1)  # B * H * V

        similarity_rlt = []
        for i in range(bs):
            tmp_a_valid = his_valid[i, :].unsqueeze(-1)  # H
            tmp_item = items[i].unsqueeze(0).expand(elements[i].size(0), -1)  # [H, V]
            tmp_a = self.logic_and(tmp_item, elements[i]) * tmp_a_valid
            similarity_rlt.append(self.similarity(tmp_a, self.true))

        return torch.stack(similarity_rlt).to(self.device)  # [H, V]

    def full_sort_predict(self, interaction):
        r"""Full sort prediction function.
        Given users, calculate the scores between users and all candidate items.

        Args:
            interaction (Interaction): Interaction class of the batch.

        Returns:
            torch.Tensor: Predicted scores for given users and all candidate items,
            shape: [n_batch_users, n_candidate_items]
        """
        # The predict function already does what is needed (users vs all items)
        prediction = self.predict(interaction)
        return prediction

    def predict_or_and(self, users, pos, neg, history):
        # Store content for checking: logic regularization
        # Compute L2 regularization on embeddings
        check_list = []
        bs = users.size(0)
        users_embed, item_embed = self.encoder.computer()

        # Each item in the history is marked as positive, but the latter part of the history may be -1,
        # indicating it is not that long
        his_valid = history.ge(0).float()  # B * H
        maxlen = int(his_valid.sum(dim=1).max().item())
        elements = item_embed[history.abs()] * his_valid.unsqueeze(-1)  # B * H * V

        # For later validation, each vector should satisfy the corresponding constraint in the logical
        # expression; 'valid' indicates the validity of the corresponding element in the constraint vector
        constraint = [elements.view([bs, -1, self.latent_dim])]  # B * H * V
        constraint_valid = [his_valid.view([bs, -1])]  # B * H

        tmp_o = None
        for i in range(maxlen):
            tmp_o_valid = his_valid[:, i].unsqueeze(-1)
            if tmp_o is None:
                tmp_o = elements[:, i, :] * tmp_o_valid  # B * V
            else:
                # Only perform OR operation if valid; otherwise, if the history is not that long (not valid),
                # keep the original content unchanged
                tmp_o = self.logic_or(tmp_o, elements[:, i, :]) * tmp_o_valid + tmp_o * (-tmp_o_valid + 1)  # B * V
                constraint.append(tmp_o.view([bs, 1, self.latent_dim]))  # B * 1 * V
                constraint_valid.append(tmp_o_valid)  # B * 1
        or_vector = tmp_o  # B * V
        left_valid = his_valid[:, 0].unsqueeze(-1)  # B * 1

        right_vector_true = item_embed[pos]  # B * V
        right_vector_false = item_embed[neg]  # B * V

        constraint.append(right_vector_true.view([bs, 1, self.latent_dim]))  # B * 1 * V
        constraint_valid.append(
            torch.ones((bs, 1)).to(self.device)
        )  # B * 1   # Indicates that all items to be judged are valid
        constraint.append(right_vector_false.view([bs, 1, self.latent_dim]))  # B * 1 * V
        constraint_valid.append(torch.ones((bs, 1)).to(self.device))  # B * 1

        sent_vector = (
            self.logic_and(or_vector, right_vector_true) * left_valid + (-left_valid + 1) * right_vector_true
        )  # B * V
        constraint.append(sent_vector.view([bs, 1, self.latent_dim]))  # B * 1 * V
        constraint_valid.append(left_valid)  # B * 1
        prediction_true = self.similarity(sent_vector, self.true, sigmoid=False).view([-1])  # B
        check_list.append(("prediction_true", prediction_true))

        sent_vector = (
            self.logic_and(or_vector, right_vector_false) * left_valid + (-left_valid + 1) * right_vector_false
        )  # B * V
        constraint.append(sent_vector.view([bs, 1, self.latent_dim]))  # B * 1 * V
        constraint_valid.append(left_valid)  # B * 1
        prediction_false = self.similarity(sent_vector, self.true, sigmoid=False).view([-1])  # B
        check_list.append(("prediction_false", prediction_false))

        constraint = torch.cat(tuple(constraint), dim=1)
        constraint_valid = torch.cat(tuple(constraint_valid), dim=1)

        return prediction_true, prediction_false, check_list, constraint, constraint_valid

    def calculate_loss(self, interaction):
        """
        Calculates the total loss by combining:
        - BCE Loss (rloss)
        - Entropy Loss (tloss)
        - L2 Loss (l2loss)
        """
        # Extraction of tensors from the interaction dictionary
        batch_users = interaction[self.USER_ID]
        batch_pos = interaction[self.ITEM_ID]
        batch_neg = interaction[self.NEG_ITEM_ID]

        # Build the history in the same way as in predict
        batch_history = self.encoder.user_history_matrix[batch_users, : self.encoder.maxhis]

        # Forward of the model with the 3 loss components
        rloss, tloss, l2loss = self.forward(False, 0, batch_users, batch_pos, batch_neg, batch_history)

        # Combination of the 3 components
        total_loss = rloss + tloss + l2loss

        return total_loss

    def triple_loss(self, TItemScore, FItemScore):
        bce_loss = self.bceloss(TItemScore.sigmoid(), torch.ones_like(TItemScore)) + self.bceloss(
            FItemScore.sigmoid(), torch.zeros_like(FItemScore)
        )
        # Input positive and negative example scores, maximizing the score difference
        if self.loss_sum:
            loss = torch.sum(F.softplus(-(TItemScore - FItemScore)))
        else:
            loss = torch.mean(F.softplus(-(TItemScore - FItemScore)))
        return (loss + bce_loss) * 0.5

    def l2_loss(self, users, pos, neg, history):
        users_embed, item_embed = self.encoder.computer()
        users_emb = users_embed[users]
        pos_emb = item_embed[pos]
        neg_emb = item_embed[neg]
        his_valid = history.ge(0).float()  # B * H
        elements = item_embed[history.abs()] * his_valid.unsqueeze(-1)  # B * H * V
        # L2 regularization loss
        reg_loss = (
            (1 / 2)
            * (users_emb.norm(2).pow(2) + pos_emb.norm(2).pow(2) + neg_emb.norm(2).pow(2) + elements.norm(2).pow(2))
            / float(len(users))
        )
        if not self.loss_sum:
            reg_loss /= users.size(0)
        return reg_loss * self.l2s_weight

    def check(self, check_list):
        """
        Logs the shape and contents of tensors in the provided check_list.
        Each element in check_list is expected to be a tuple where the first item is a string (label)
        and the second item is a tensor. For each tuple, this function:
          - Converts the tensor to a NumPy array after detaching it from the computation graph and moving it to CPU.
          - Logs the label and the shape of the array.
          - Logs the array contents, with a threshold of 20 elements for display.
        Args:
            check_list (list of tuple): List of (label, tensor) pairs to be logged for inspection.
        """

        logging.info(os.linesep)
        for t in check_list:
            d = np.array(t[1].detach().cpu())
            logging.info(os.linesep.join([t[0] + "\t" + str(d.shape), np.array2string(d, threshold=20)]) + os.linesep)

    def forward(self, print_check: bool, return_pred: bool, *args, **kwards):
        prediction1, prediction0, check_list, constraint, constraint_valid = self.predict_or_and(*args, **kwards)
        rloss = self.logic_regularizer(False, check_list, constraint, constraint_valid)
        tloss = self.triple_loss(prediction1, prediction0)
        l2loss = self.l2_loss(*args, **kwards)

        if print_check:
            self.check(check_list)

        if return_pred:
            return prediction1, rloss + tloss + l2loss
        return rloss, tloss, l2loss


class GAT(nn.Module):
    def __init__(self, nfeat, nhid, dropout, alpha):
        """Dense version of GAT."""
        super().__init__()
        self.dropout = dropout

        self.layer = GraphAttentionLayer(nfeat, nhid, dropout=dropout, alpha=alpha, concat=False)

    def forward(self, item_embs, entity_embs, adj):
        x = F.dropout(item_embs, self.dropout, training=self.training)
        y = F.dropout(entity_embs, self.dropout, training=self.training)
        x = self.layer(x, y, adj)
        x = F.dropout(x, self.dropout, training=self.training)
        return x

    def forward_relation(self, item_embs, entity_embs, w_r, adj):
        x = F.dropout(item_embs, self.dropout, training=self.training)
        y = F.dropout(entity_embs, self.dropout, training=self.training)
        x = self.layer.forward_relation(x, y, w_r, adj)
        x = F.dropout(x, self.dropout, training=self.training)
        return x
