model.py

import torch
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import scipy.sparse as sp
import random
import math
from toTsne import save_user_embeddings_as_tsne
from utils import *

init = nn.init.xavier_uniform_

class D2ASR(nn.Module):
    def __init__(self, config, args):
        super(D2ASR, self).__init__()

        
        self.args = args
        self.user_num = config['user_num']
        self.item_num = config['item_num']
        self.embed_dim = args.embed_dim
        self.n_layers = args.n_layers
        self.n_layer_sn = args.n_layer_sn
        self.n_head = args.n_head
        self.dropout_rate = args.dropout_rate
        self.device = args.device
        self.kd_temp = args.cl_temp
        
        """ ********* Initialize Model Parameters ********* """
        self.user_embedding = nn.Parameter(init(torch.empty(self.user_num, self.embed_dim)))
        self.item_embedding = nn.Parameter(init(torch.empty(self.item_num, self.embed_dim)))
        
        """ ********* Load U-U Social Network and U-I Interaction Graph ********* """
        self.adj_graph = config['adj_graph']
        self.social_graph = config['social_graph']
        self.social_standard = self.generate_norm_social(self.social_graph)
        
        """ ********* Initialize Graph Generator ********* """
        """ self.userlinner1 = nn.Linear(self.embed_dim, self.embed_dim)
        self.userlinner2 = nn.Linear(self.embed_dim, self.embed_dim) """
        self.userlinner1 = Extractor([64, 32, 16, 64])
        self.userlinner2 = Extractor([64, 32, 16, 64])
        self.popitemlinner1 = nn.Linear(in_features=2*self.embed_dim, out_features=self.embed_dim)
        self.gcnLayers = nn.Sequential(*[GCNLayer() for i in range(self.n_layers)])
        self.socialgcn = SocialGCN(self.n_layer_sn, self.dropout_rate)
        self.socialgt = SocialGT(self.embed_dim, self.device, self.n_head, self.user_num, self.gcnLayers)
        
        """ ********* Loss Function ********* """
        self.kd_criterion = KDLoss(self.kd_temp)
        
        
    def get_item_embedding(self):
        return self.item_embedding

    def get_user_embedding(self):
        return self.user_embedding
    
    def get_embeddings(self):
        return torch.cat([self.user_embedding, self.item_embedding], dim=0)
    
    def _normalize_adj_standard(self, mat):
        degree = np.array(mat.sum(axis=-1))
        dInvSqrt = np.reshape(np.power(degree, -0.5), [-1])
        dInvSqrt[np.isinf(dInvSqrt)] = 0.0
        dInvSqrtMat = sp.diags(dInvSqrt)
        return mat.dot(dInvSqrtMat).transpose().dot(dInvSqrtMat).tocoo()
    
    def plain_to_norm_social(self, torch_mat):
        rows = torch_mat[0].cpu().numpy()
        cols = torch_mat[1].cpu().numpy()
        values = np.ones_like(rows)
        mat = sp.coo_matrix((values, (rows, cols)), shape=(self.user_num, self.user_num))
        mat = (mat + sp.eye(mat.shape[0])) * 1.0
        mat = self._normalize_adj_standard(mat)
        # make cuda tensor
        idxs = torch.from_numpy(np.vstack([mat.row, mat.col]).astype(np.int64)) # torch.Size([2, 434248])
        vals = torch.from_numpy(mat.data.astype(np.float32))
        shape = torch.Size(mat.shape)
        return torch.sparse.FloatTensor(idxs, vals, shape).to(self.device)
    
    def generate_norm_social(self, mat):
        mat = (mat != 0) * 1.0
        mat = (mat + sp.eye(mat.shape[0])) * 1.0
        mat = self._normalize_adj_standard(mat)
        # make cuda tensor
        idxs = torch.from_numpy(np.vstack([mat.row, mat.col]).astype(np.int64)) # torch.Size([2, 434248])
        vals = torch.from_numpy(mat.data.astype(np.float32))
        shape = torch.Size(mat.shape)
        return torch.sparse.FloatTensor(idxs, vals, shape).to(self.device)
    
    
    def generate_norm_adj(self, mat, mean=False):
        # make ui adj
        a = sp.csr_matrix((mat.shape[0], mat.shape[0]))
        b = sp.csr_matrix((mat.shape[1], mat.shape[1]))
        # #([42712, 42712] || [42712, 26822]) === ([26822, 42712] || [26822, 26822]) ——> [42712, 69534] === [26822, 69534] ——> [69534, 69534]
        mat = sp.vstack([sp.hstack([a, mat]), sp.hstack([mat.transpose(), b])]) 
        mat = (mat != 0) * 1.0
        mat = (mat + sp.eye(mat.shape[0])) * 1.0
        if mean:
            mat = self._normalize_adj_mean(mat)
        else:
            mat = self._normalize_adj_standard(mat)
        # make cuda tensor
        idxs = torch.from_numpy(np.vstack([mat.row, mat.col]).astype(np.int64)) # torch.Size([2, 434248])
        vals = torch.from_numpy(mat.data.astype(np.float32))
        shape = torch.Size(mat.shape)
        return torch.sparse.FloatTensor(idxs, vals, shape).to(self.device)
    
    def edge_dropout(self, sp_adj, social=False):
        """Input: a sparse user-item adjacency matrix and a dropout rate."""
        adj_shape = sp_adj.get_shape()
        edge_count = sp_adj.count_nonzero()
        row_idx, col_idx = sp_adj.nonzero()
        keep_idx = random.sample(range(edge_count), int(edge_count * (1 - self.dropout_rate)))
        user_np = np.array(row_idx)[keep_idx]
        item_np = np.array(col_idx)[keep_idx]
        edges = np.ones_like(user_np, dtype=np.float32)
        X = sp.csr_matrix((edges, (user_np, item_np)), shape=adj_shape)
        if social:
            dropout_adj = self.generate_norm_social(X)
        else:
            dropout_adj = self.generate_norm_adj(X)
        return dropout_adj
    
    def lightgcn(self, adj, embeds):
        embedsLst = [embeds]
        for gcn in self.gcnLayers:
            embeds = gcn(adj, embedsLst[-1])
            embedsLst.append(embeds)
        embeds = torch.stack(embedsLst, dim=1)
        light_out = torch.mean(embeds, dim=1)
        return torch.split(light_out, [self.user_num, self.item_num])
        
    def forward(self, ui_user_idx, ui_pos_idx, ui_neg_idx, 
                uu_user_idx, pos_user_idx, neg_user_idx,
                pop_embedding_inter, pop_embedding_social,
                norm_adj, dropped_adj1, dropped_adj2, 
                social_norm_aug2, social_norm_aug1, 
                aug_1, aug_2, step):
        
        embeds = torch.concat([self.user_embedding, self.item_embedding], dim=0)
        user_embedding, item_embedding = self.lightgcn(norm_adj, embeds)
        user_embedding_aug1, item_embedding_aug1 = self.lightgcn(dropped_adj1, embeds)
        user_embedding_aug2, item_embedding_aug2 = self.lightgcn(dropped_adj2, embeds)
        
        user_aug_linner1 = F.relu(self.userlinner1(user_embedding_aug1))
        user_aug_linner2 = F.relu(self.userlinner2(user_embedding_aug2))
        
        if args.use_diff:
            user_social_embedding = self.socialgcn(self.social_standard, self.get_user_embedding())
            #save_user_embeddings_as_tsne(user_social_embedding, "user_social", args.dataset)
            user_social_embedding_diff = social_norm_aug2
            #save_user_embeddings_as_tsne(user_social_embedding, "user_diff", args.dataset)
        else:
            user_social_embedding = self.socialgcn(social_norm_aug1, self.get_user_embedding())
            user_social_embedding_diff = self.socialgcn(social_norm_aug2, self.get_user_embedding())
        #user_social_embedding_diff = self.socialgcn(social_norm_aug2, self.get_user_embedding())
        
        user_social_embedding_aug1 = aug_1(user_social_embedding)
        user_social_embedding_diff_aug2 = aug_2(user_social_embedding_diff)
        
        ui_u_embeddings = user_embedding[ui_user_idx]
        pos_embeddings = item_embedding[ui_pos_idx]
        neg_embeddings = item_embedding[ui_neg_idx]
        
        
        uu_u_embeddings = user_social_embedding[uu_user_idx]
        pos_u_embeddings = user_social_embedding[pos_user_idx]
        neg_u_embeddings = user_social_embedding[neg_user_idx]
        
        ui_u_embeddings_T = user_social_embedding[ui_user_idx]
        pos_embeddings_T = self.get_item_embedding()[ui_pos_idx]
        neg_embeddings_T = self.get_item_embedding()[ui_neg_idx]
        
        pos_scores_ui_S = torch.mul(ui_u_embeddings, pos_embeddings).sum(dim=1)
        neg_scores_ui_S = torch.mul(ui_u_embeddings, neg_embeddings).sum(dim=1)
        pos_scores_ui_T = torch.mul(ui_u_embeddings_T, pos_embeddings_T).sum(dim=1)
        neg_scores_ui_T = torch.mul(ui_u_embeddings_T, neg_embeddings_T).sum(dim=1)
        
        kd_loss_ui_S2T = self.kd_criterion(pos_scores_ui_T, pos_scores_ui_S, neg_scores_ui_T)
        kd_loss_ui_T2S = self.kd_criterion(pos_scores_ui_S, pos_scores_ui_T, neg_scores_ui_S)
        
        batch_bpr_loss = pairPredict(ui_u_embeddings, pos_embeddings, neg_embeddings)
        
        batch_social_loss = pairPredict(uu_u_embeddings, pos_u_embeddings, neg_u_embeddings) * args.social_reg
        batch_reg_loss = calcRegLoss(self.parameters()) * args.reg
        batch_cl_loss = cal_cl_loss(args.cl_temp, [ui_user_idx, ui_pos_idx], user_embedding_aug1, user_embedding_aug2, item_embedding_aug1, item_embedding_aug2) * args.ssl_reg
        batch_social_cl_loss = cal_cl_loss(args.cl_temp, [uu_user_idx, pos_user_idx], user_social_embedding_aug1, user_social_embedding_diff_aug2) * args.social_cl_reg
        batch_cross_cl_loss = cal_cross_cl_loss(args.cl_temp, [uu_user_idx, pos_user_idx], user_aug_linner1, user_aug_linner2, user_social_embedding_aug1, user_social_embedding_diff_aug2) * args.cross_cl_reg
        batch_kd_loss = (args.kd_reg_S2T * kd_loss_ui_S2T + args.kd_reg_T2S * kd_loss_ui_T2S) * args.kd_reg
        
        return batch_bpr_loss, batch_reg_loss,\
                batch_social_loss, batch_social_cl_loss,  \
                batch_cl_loss, batch_cross_cl_loss, \
                batch_kd_loss
    
    def predict(self, adj, need_score=False, users=None, trainMask=None):
        embeds = self.get_embeddings()
        user_embedding_final, item_embedding_final = self.lightgcn(adj, embeds)
        if need_score:
            batch_ratings = torch.mm(user_embedding_final[users], torch.transpose(item_embedding_final, 1, 0)) * (1 - trainMask) - trainMask * 1e8
            return batch_ratings
        return user_embedding_final, item_embedding_final
            
    
class GCNLayer(nn.Module):
    def __init__(self):
        super(GCNLayer, self).__init__()

    def forward(self, adj, embeds):
        return torch.spmm(adj, embeds)

class SocialGCN(nn.Module):
    def __init__(self, n_layer, dropout_rate):
        super(SocialGCN, self).__init__()
        self.n_layer = n_layer
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, norm_social, user_embedding):
        embeds = [user_embedding]
        for i in range(self.n_layer):
            embeds.append(torch.spmm(norm_social, embeds[i]))
        lightout = torch.mean(torch.stack(embeds, dim=1), dim=1)
        return lightout
    
class SocialGT(nn.Module):
    def __init__(self, embed_dim, device, n_head, user_num, gcn_layers):
        super(SocialGT, self).__init__()
        init = nn.init.xavier_uniform_
        self.embed_dim = embed_dim
        self.qTrans = nn.Parameter(init(torch.empty(self.embed_dim, self.embed_dim)))
        self.kTrans = nn.Parameter(init(torch.empty(self.embed_dim, self.embed_dim)))
        self.vTrans = nn.Parameter(init(torch.empty(self.embed_dim, self.embed_dim)))
        self.user_num = user_num
        self.device = device
        self.n_head = n_head
        self.gcn_layers = gcn_layers
        
    def forward(self, norm_adj, social_adj, user_embeddings, item_embeddings):
        """ LightGCN but not Pooling"""
        embeds = torch.concat([user_embeddings, item_embeddings], dim=0)
        embedsLst = [embeds]
        for i, gcn in enumerate(self.gcn_layers):
            embeds = gcn(norm_adj, embedsLst[-1])
        #embeds = sum(embedsLst) # torch.Size([69534, 32])
        embeds = embedsLst[-1]
        
        indices = social_adj._indices()
        rows, cols = indices[0, :], indices[1, :]
        embeds = embeds[:self.user_num]
        rowEmbeds = embeds[rows] # torch.Size([769765, 64])
        colEmbeds = embeds[cols] # torch.Size([769765, 64])

        qEmbeds = (rowEmbeds @ self.qTrans).view([-1, self.n_head, self.embed_dim // self.n_head])
        kEmbeds = (colEmbeds @ self.kTrans).view([-1, self.n_head, self.embed_dim // self.n_head])
        vEmbeds = (colEmbeds @ self.vTrans).view([-1, self.n_head, self.embed_dim // self.n_head])

        att = torch.einsum('ehd, ehd -> eh', qEmbeds, kEmbeds)
        att = torch.clamp(att, -10.0, 10.0)
        expAtt = torch.exp(att)
        tem = torch.zeros([social_adj.shape[0], self.n_head]).to(self.device)
        attNorm = (tem.index_add_(0, rows, expAtt))[rows]
        att = expAtt / (attNorm + 1e-8)

        resEmbeds = torch.einsum('eh, ehd -> ehd', att, vEmbeds).view([-1, self.embed_dim])
        tem = torch.zeros([social_adj.shape[0], self.embed_dim]).to(self.device) # torch.Size([45919, 64])
        resEmbeds = tem.index_add_(0, rows, resEmbeds)  # nd rows = torch.Size([769765])
        return resEmbeds, att # torch.Size([45919, 64]) torch.Size([769765, 4])
    
class Denoise(nn.Module):
    def __init__(self, device, n_layer_sn, dropout_rate, social_graph, in_dims, out_dims, z_dim, emb_size, norm=False, dropout=0.5):
        super(Denoise, self).__init__()
        self.device = device
        self.in_dims = in_dims
        self.out_dims = out_dims
        self.time_emb_dim = emb_size
        self.norm = norm
        self.n_layer_sn = n_layer_sn
        self.dropout_rate = dropout_rate
        self.social_graph = social_graph
        self.z_dim = z_dim
        self.social_standard = self.generate_norm_social(self.social_graph)
        self.gcn_layers = nn.Sequential(*[GCNLayer() for i in range(3)])

        self.emb_layer = nn.Linear(self.time_emb_dim, self.time_emb_dim)
        
        self.socialgcn = SocialGCN(self.n_layer_sn, self.dropout_rate)

        in_dims_temp = [self.time_emb_dim + self.z_dim] + self.in_dims[1:]

        out_dims_temp = self.out_dims

        self.in_layers = nn.ModuleList([nn.Linear(d_in, d_out) for d_in, d_out in zip(in_dims_temp[:-1], in_dims_temp[1:])])
        self.out_layers = nn.ModuleList([nn.Linear(d_in, d_out) for d_in, d_out in zip(out_dims_temp[:-1], out_dims_temp[1:])])

        self.drop = nn.Dropout(dropout)
        self.init_weights()
        
        # Generate epsilon from standard normal distribution
        
        # VAE parameters
        self.A = nn.Parameter(torch.triu(torch.randn(self.z_dim, self.z_dim), 1))  # Upper triangular matrix A
        
        self.f_e = nn.Linear(self.in_dims[0], self.z_dim)
        
        # Nonlinear functions g_i (using simple MLPs here as an example)
        self.g_i = nn.ModuleList([nn.Sequential(nn.Linear(self.z_dim, self.z_dim), nn.ReLU(), nn.Linear(self.z_dim, 1)) for _ in range(self.z_dim)])
    
    def get_social_embedding(self, user_embedding):
        return self.socialgcn(self.social_standard, user_embedding)
    
    def generate_norm_social(self, mat):
        mat = (mat != 0) * 1.0
        mat = (mat + sp.eye(mat.shape[0])) * 1.0
        mat = self._normalize_adj_standard(mat)
        # make cuda tensor
        idxs = torch.from_numpy(np.vstack([mat.row, mat.col]).astype(np.int64)) # torch.Size([2, 434248])
        vals = torch.from_numpy(mat.data.astype(np.float32))
        shape = torch.Size(mat.shape)
        return torch.sparse.FloatTensor(idxs, vals, shape).to(self.device)
    
    def _normalize_adj_standard(self, mat):
        degree = np.array(mat.sum(axis=-1))
        dInvSqrt = np.reshape(np.power(degree, -0.5), [-1])
        dInvSqrt[np.isinf(dInvSqrt)] = 0.0
        dInvSqrtMat = sp.diags(dInvSqrt)
        return mat.dot(dInvSqrtMat).transpose().dot(dInvSqrtMat).tocoo()

    def init_weights(self):
        for layer in self.in_layers:
            size = layer.weight.size()
            if size[0] + size[1] == 0:  # 检测是否为零维度
                std = 0.00000001  # 设置一个小的默认值
            else:
                std = np.sqrt(2.0 / (size[0] + size[1]))
            layer.weight.data.normal_(0.0, std)
            layer.bias.data.normal_(0.0, 0.001)
    
        for layer in self.out_layers:
            size = layer.weight.size()
            if size[0] + size[1] == 0:  # 检测是否为零维度
                std = 0.00000001  # 设置一个小的默认值
            else:
                std = np.sqrt(2.0 / (size[0] + size[1]))
            layer.weight.data.normal_(0.0, std)
            layer.bias.data.normal_(0.0, 0.001)
    
        size = self.emb_layer.weight.size()
        if size[0] + size[1] == 0:  # 检测是否为零维度
            std = 0.00000001  # 设置一个小的默认值
        else:
            std = np.sqrt(2.0 / (size[0] + size[1]))
        self.emb_layer.weight.data.normal_(0.0, std)
        self.emb_layer.bias.data.normal_(0.0, 0.001)

    """ def forward(self, x, timesteps, mess_dropout=True):
        freqs = torch.exp(-math.log(10000) * torch.arange(start=0, end=self.time_emb_dim//2, dtype=torch.float32) / (self.time_emb_dim//2)).to(self.device)
        temp = timesteps[:, None].float() * freqs[None]
        time_emb = torch.cat([torch.cos(temp), torch.sin(temp)], dim=-1)
        if self.time_emb_dim % 2:
            time_emb = torch.cat([time_emb, torch.zeros_like(time_emb[:, :1])], dim=-1)
        emb = self.emb_layer(time_emb)
        if self.norm:
            x = F.normalize(x.to(torch.float32))
        if mess_dropout:
            x = self.drop(x)
        h = torch.cat([x, emb], dim=-1)
        for i, layer in enumerate(self.in_layers):
            h = layer(h)
            h = torch.tanh(h)
        for i, layer in enumerate(self.out_layers):
            h = layer(h)
            if i != len(self.out_layers) - 1:
                h = torch.tanh(h)

        return h """
    
    def forward(self, x, timesteps, mess_dropout=True):
        # Calculate z using the formula z = (I - A^T)^(-1) * epsilon
        epsilon = F.relu(self.f_e(x.type(torch.float32)))
        I = torch.eye(self.z_dim).to(self.device)
        A_T = self.A.T
        inverse_term = torch.inverse(I - A_T)
        z = torch.matmul(inverse_term, epsilon.T).T
        
        # Apply nonlinear functions g_i to z
        z_nonlinear = torch.zeros_like(z)
        for i in range(self.z_dim):
            z_nonlinear[:, i] = self.g_i[i](z).squeeze(-1)
        
        # Concatenate z_nonlinear with input x
        #x = torch.cat([x, z_nonlinear], dim=1)
        x = z_nonlinear
        
        #save_user_embeddings_as_tsne(x, "user_causal", args.dataset)
        # Time embedding calculations
        freqs = torch.exp(-math.log(10000) * torch.arange(start=0, end=self.time_emb_dim // 2, dtype=torch.float32) / (self.time_emb_dim // 2)).to(self.device)
        temp = timesteps[:, None].float() * freqs[None]
        time_emb = torch.cat([torch.cos(temp), torch.sin(temp)], dim=-1)
        
        if self.time_emb_dim % 2:
            time_emb = torch.cat([time_emb, torch.zeros_like(time_emb[:, :1])], dim=-1)
            
        emb = self.emb_layer(time_emb)
        
        if self.norm:
            x = F.normalize(x.to(torch.float32))
            
        if mess_dropout:
            x = self.drop(x)
            
        h = torch.cat([x, emb], dim=-1)
        
        for i, layer in enumerate(self.in_layers):
            h = layer(h)
            h = torch.tanh(h)
            
        for i, layer in enumerate(self.out_layers):
            h = layer(h)
            if i != len(self.out_layers) - 1:
                h = torch.tanh(h)
        
        return h

class GaussianDiffusion(nn.Module):
    def __init__(self, device, item_num, noise_scale, noise_min, noise_max, steps, beta_fixed=True):
        super(GaussianDiffusion, self).__init__()

        self.noise_scale = noise_scale
        self.noise_min = noise_min
        self.noise_max = noise_max
        self.steps = steps
        self.item_num = item_num
        self.device = device

        if noise_scale != 0:
            self.betas = torch.tensor(self.get_betas(), dtype=torch.float64).to(self.device)
            if beta_fixed:
                self.betas[0] = 0.0001

            self.calculate_for_diffusion()

    def get_betas(self):
        start = self.noise_scale * self.noise_min
        end = self.noise_scale * self.noise_max
        variance = np.linspace(start, end, self.steps, dtype=np.float64)
        alpha_bar = 1 - variance
        betas = []
        betas.append(1 - alpha_bar[0])
        for i in range(1, self.steps):
            betas.append(min(1 - alpha_bar[i] / alpha_bar[i-1], 0.999))
        return np.array(betas)
    
    def calculate_for_diffusion(self):
        alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(alphas, axis=0).to(self.device)
        self.alphas_cumprod_prev = torch.cat([torch.tensor([1.0]).to(self.device), self.alphas_cumprod[:-1]]).to(self.device)
        self.alphas_cumprod_next = torch.cat([self.alphas_cumprod[1:], torch.tensor([0.0]).to(self.device)]).to(self.device)

        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
        self.log_one_minus_alphas_cumprod = torch.log(1.0 - self.alphas_cumprod)
        self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)

        self.posterior_variance = (
            self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_log_variance_clipped = torch.log(torch.cat([self.posterior_variance[1].unsqueeze(0), self.posterior_variance[1:]]))
        self.posterior_mean_coef1 = (self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod))
        self.posterior_mean_coef2 = ((1.0 - self.alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - self.alphas_cumprod))

    def p_sample(self, model, x_start, steps):
        if steps == 0:
            x_t = x_start
        else:
            t = torch.tensor([steps-1] * x_start.shape[0]).to(self.device)
            x_t = self.q_sample(x_start, t)
        
        indices = list(range(self.steps))[::-1]

        for i in indices:
            t = torch.tensor([i] * x_t.shape[0]).to(self.device)
            model_mean, model_log_variance = self.p_mean_variance(model, x_t, t)
            x_t = model_mean
        return x_t
    
    def q_sample(self, x_start, t, noise=None):
        # print(self.betas)
        if noise is None:
            noise = torch.randn_like(x_start)
            # noise = torch.randn_like(x_start) / 100
        sqrt_alphas_cumprod_t = self.extract(self.sqrt_alphas_cumprod, t, x_start.shape)
        sqrt_one_minus_alphas_cumprod_t = self.extract(
            self.sqrt_one_minus_alphas_cumprod, t, x_start.shape
        )
        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
    
    def extract(self, a, t, x_shape):
        batch_size = t.shape[0]
        out = a.gather(-1, t)
        return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)
    
    """ def q_sample(self, x_start, t, noise=None):
        if noise is None:
            noise = torch.randn_like(x_start)
        return self._extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start + self._extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise """
    
    def _extract_into_tensor(self, arr, timesteps, broadcast_shape):
        arr = arr.to(self.device)
        res = arr[timesteps].float()
        while len(res.shape) < len(broadcast_shape):
            res = res[..., None]
        return res.expand(broadcast_shape)
    
    def p_mean_variance(self, model, x, t):
        model_output = model(x, t, False)

        model_variance = self.posterior_variance
        model_log_variance = self.posterior_log_variance_clipped

        model_variance = self._extract_into_tensor(model_variance, t, x.shape)
        model_log_variance = self._extract_into_tensor(model_log_variance, t, x.shape)

        model_mean = (self._extract_into_tensor(self.posterior_mean_coef1, t, x.shape) * model_output + self._extract_into_tensor(self.posterior_mean_coef2, t, x.shape) * x)
        
        return model_mean, model_log_variance

    def training_losses(self, model, x_start, ui_matrix, user_embedding, item_embedding, batch_index=None):
        batch_size = x_start.size(0)
        ts = torch.randint(0, self.steps, (batch_size,)).long().to(self.device)
        noise = torch.randn_like(x_start)
        if self.noise_scale != 0:
            x_t = self.q_sample(x_start, ts, noise)
        else:
            x_t = x_start

        model_output = model(x_t, ts)
        mse = self.mean_flat((x_start - model_output) ** 2)

        weight = self.SNR(ts - 1) - self.SNR(ts)
        weight = torch.where((ts == 0), 1.0, weight)

        h_a = self._h_A(model.A, model.A.size()[0])
        diff_loss = weight * mse + 0.03*h_a + 0.005*h_a*h_a

        #item_user_matrix = torch.spmm(ui_matrix.t(), model_output.t()).t()
        #user_embedding_social = torch.mm(item_user_matrix, item_embedding)
        ukgc_loss = self.mean_flat((model_output - user_embedding[batch_index]) ** 2)
        ukgc_loss = torch.zeros_like(diff_loss).to(device)

        return diff_loss.mean(), ukgc_loss.mean()
        
    def mean_flat(self, tensor):
        return tensor.mean(dim=list(range(1, len(tensor.shape))))
    
    def SNR(self, t):
        self.alphas_cumprod = self.alphas_cumprod.to(self.device)
        return self.alphas_cumprod[t] / (1 - self.alphas_cumprod[t])
    
    def _h_A(self, A, m):
        expm_A = self.matrix_poly(A*A, m)
        h_A = torch.trace(expm_A) - m
        return h_A
    
    def matrix_poly(self, matrix, d):
        x = torch.eye(d).to(device)+ torch.div(matrix.to(device), d).to(device)
        return torch.matrix_power(x, d)

class Normalize(nn.Module):
    def __init__(self, power=2):
        super(Normalize, self).__init__()
        self.power = power

    def forward(self, x):
        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)
        out = x.div(norm)
        return out

class Extractor(nn.Module):
    def __init__(self,layers):
        super(Extractor, self).__init__()
        self.dense_1=nn.Linear(layers[0],layers[1])
        self.dense_2=nn.Linear(layers[1],layers[2])
        self.dense_3=nn.Linear(layers[2],layers[3])
        self.actfunction = F.relu
        self.normal=Normalize()
    def forward(self,input): # [B H]
        # layer 1
        output=self.dense_1(input)
        if self.actfunction!=None:
            output=self.actfunction(output)
        # layer 2
        output=self.dense_2(output)
        if self.actfunction != None:
            output = self.actfunction(output)
        # layer 3
        output=self.dense_3(output)
        output=self.normal(output)
        return output