图神经网络推荐¶

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from scipy.sparse import coo_matrix

class GCNLayer(nn.Module):  # 继承nn.Module定义网络层
    """
    简单的图卷积层实现（纯 PyTorch，不依赖 PyG）

    实现公式：H^{(l+1)} = σ(Ã H^{(l)} W^{(l)})
    其中 Ã = D^{-1/2} A D^{-1/2}（含自环）
    """
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()  # super()调用父类方法
        self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
        if bias:
            self.bias = nn.Parameter(torch.FloatTensor(out_features))
        else:
            self.bias = None
        self.reset_parameters()

    def reset_parameters(self):
        """Xavier 初始化"""
        nn.init.xavier_uniform_(self.weight)
        if self.bias is not None:
            nn.init.zeros_(self.bias)

    def forward(self, x: torch.Tensor, adj_norm: torch.Tensor) -> torch.Tensor:
        """
        前向传播
        Args:
            x: 节点特征矩阵 [num_nodes, in_features]
            adj_norm: 归一化邻接矩阵 Ã [num_nodes, num_nodes]（稀疏张量）
        Returns:
            输出特征矩阵 [num_nodes, out_features]
        """
        # 特征变换：H W
        support = torch.mm(x, self.weight)
        # 邻域聚合：Ã (H W)
        output = torch.spmm(adj_norm, support)
        if self.bias is not None:
            output += self.bias
        return output

class SimpleGCN(nn.Module):
    """
    两层 GCN 模型
    用于节点分类或生成节点 Embedding
    """
    def __init__(self, n_features: int, n_hidden: int, n_classes: int, dropout: float = 0.5):
        super().__init__()
        self.gc1 = GCNLayer(n_features, n_hidden)
        self.gc2 = GCNLayer(n_hidden, n_classes)
        self.dropout = dropout

    def forward(self, x: torch.Tensor, adj_norm: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: 节点特征 [num_nodes, n_features]
            adj_norm: 归一化邻接矩阵
        """
        x = F.relu(self.gc1(x, adj_norm))  # F.xxx PyTorch函数式API
        x = F.dropout(x, self.dropout, training=self.training)
        x = self.gc2(x, adj_norm)
        return x

def build_normalized_adj(edge_index: np.ndarray, num_nodes: int) -> torch.Tensor:
    """
    从边列表构建归一化邻接矩阵（含自环）

    Args:
        edge_index: 边列表 [2, num_edges]，每列是 (src, dst)
        num_nodes: 节点总数
    Returns:
        归一化后的稀疏邻接矩阵 Ã = D^{-1/2} (A + I) D^{-1/2}
    """
    # 构建邻接矩阵（加自环）
    row = np.concatenate([edge_index[0], edge_index[1], np.arange(num_nodes)])
    col = np.concatenate([edge_index[1], edge_index[0], np.arange(num_nodes)])
    data = np.ones(len(row), dtype=np.float32)
    adj = coo_matrix((data, (row, col)), shape=(num_nodes, num_nodes))

    # 计算度矩阵 D^{-1/2}
    degree = np.array(adj.sum(axis=1)).flatten()  # np.array创建NumPy数组
    d_inv_sqrt = np.power(degree, -0.5)
    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0

    # 对称归一化：D^{-1/2} A D^{-1/2}
    row_norm = d_inv_sqrt[adj.row]
    col_norm = d_inv_sqrt[adj.col]
    adj.data = adj.data * row_norm * col_norm

    # 转为 PyTorch 稀疏张量
    indices = torch.LongTensor(np.vstack([adj.row, adj.col]))
    values = torch.FloatTensor(adj.data)
    shape = torch.Size(adj.shape)
    return torch.sparse_coo_tensor(indices, values, shape)

# ---- 使用示例 ----
if __name__ == "__main__":
    # 构造一个简单的 5 节点图
    edge_index = np.array([[0, 0, 1, 2, 3],
                           [1, 2, 2, 3, 4]])  # 5 条边
    num_nodes = 5

    # 构建归一化邻接矩阵
    adj_norm = build_normalized_adj(edge_index, num_nodes)

    # 随机节点特征
    x = torch.randn(num_nodes, 16)

    # 创建并运行 GCN
    model = SimpleGCN(n_features=16, n_hidden=32, n_classes=8)
    output = model(x, adj_norm)
    print(f"输入形状: {x.shape}")        # [5, 16]
    print(f"输出形状: {output.shape}")    # [5, 8]
    print(f"输出: {output}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from scipy.sparse import coo_matrix, eye as sp_eye
from collections import defaultdict
import random
import time

class LightGCN(nn.Module):
    """
    LightGCN 推荐模型（纯 PyTorch 实现，不依赖 PyG）

    论文：LightGCN: Simplifying and Powering Graph Convolution Network
          for Recommendation (SIGIR 2020)

    核心特点：
    1. 移除了 GCN 中的特征变换和非线性激活
    2. 使用层组合（Layer Combination）融合多层信息
    3. 仅需学习初始 Embedding
    """
    def __init__(self, num_users: int, num_items: int, embedding_dim: int = 64,
                 n_layers: int = 3, reg_weight: float = 1e-4):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.n_layers = n_layers
        self.reg_weight = reg_weight

        # 初始 Embedding（唯一需要学习的参数）
        self.user_embedding = nn.Embedding(num_users, embedding_dim)
        self.item_embedding = nn.Embedding(num_items, embedding_dim)

        # Xavier 初始化
        nn.init.xavier_uniform_(self.user_embedding.weight)
        nn.init.xavier_uniform_(self.item_embedding.weight)

        # 归一化邻接矩阵（在 set_graph 中设置）
        self.adj_norm = None

    def set_graph(self, user_item_pairs: np.ndarray):
        """
        构建并设置归一化邻接矩阵

        Args:
            user_item_pairs: [N, 2] 数组，每行是 (user_id, item_id)
        """
        n_nodes = self.num_users + self.num_items

        # 用户-物品二部图邻接矩阵
        # 用户节点编号: 0 ~ num_users-1
        # 物品节点编号: num_users ~ num_users+num_items-1
        users = user_item_pairs[:, 0]
        items = user_item_pairs[:, 1] + self.num_users  # 物品编号偏移

        # 构建对称邻接矩阵 (无自环)
        row = np.concatenate([users, items])
        col = np.concatenate([items, users])
        data = np.ones(len(row), dtype=np.float32)
        adj = coo_matrix((data, (row, col)), shape=(n_nodes, n_nodes))

        # D^{-1/2} A D^{-1/2} 归一化
        degree = np.array(adj.sum(axis=1)).flatten()
        d_inv_sqrt = np.power(degree, -0.5)
        d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0
        adj.data = adj.data * d_inv_sqrt[adj.row] * d_inv_sqrt[adj.col]

        # 转为 PyTorch 稀疏张量
        indices = torch.LongTensor(np.vstack([adj.row, adj.col]))
        values = torch.FloatTensor(adj.data)
        self.adj_norm = torch.sparse_coo_tensor(indices, values,
                                                 torch.Size([n_nodes, n_nodes]))

    def compute_all_embeddings(self) -> tuple:
        """
        LightGCN 的核心：多层图卷积 + 层组合

        Returns:
            (user_embeddings, item_embeddings) 最终表征
        """
        # 拼接用户和物品的初始 Embedding
        all_emb = torch.cat([self.user_embedding.weight,  # torch.cat沿已有维度拼接张量
                             self.item_embedding.weight], dim=0)  # [n_nodes, dim]

        emb_list = [all_emb]  # 保存每层的 Embedding

        # 多层图卷积（无特征变换，无激活）
        for layer in range(self.n_layers):
            all_emb = torch.spmm(self.adj_norm, all_emb)  # Ã · E^{(l)}
            emb_list.append(all_emb)

        # 层组合：各层均匀加权
        all_emb = torch.stack(emb_list, dim=0).mean(dim=0)  # [n_nodes, dim]  # torch.stack沿新维度拼接张量

        user_emb = all_emb[:self.num_users]
        item_emb = all_emb[self.num_users:]
        return user_emb, item_emb

    def bpr_loss(self, users: torch.Tensor, pos_items: torch.Tensor,
                 neg_items: torch.Tensor) -> torch.Tensor:
        """
        BPR 损失函数

        Args:
            users: 用户 ID [batch_size]
            pos_items: 正样本物品 ID [batch_size]
            neg_items: 负样本物品 ID [batch_size]
        Returns:
            BPR loss + L2 正则化
        """
        user_emb, item_emb = self.compute_all_embeddings()

        u_emb = user_emb[users]           # [batch, dim]
        pos_emb = item_emb[pos_items]      # [batch, dim]
        neg_emb = item_emb[neg_items]      # [batch, dim]

        # BPR Loss
        pos_scores = (u_emb * pos_emb).sum(dim=1)  # [batch]
        neg_scores = (u_emb * neg_emb).sum(dim=1)   # [batch]
        bpr_loss = -F.logsigmoid(pos_scores - neg_scores).mean()

        # L2 正则化（只对初始 Embedding 正则化）
        u_emb_0 = self.user_embedding(users)
        pos_emb_0 = self.item_embedding(pos_items)
        neg_emb_0 = self.item_embedding(neg_items)
        reg_loss = self.reg_weight * (
            u_emb_0.norm(2).pow(2) +
            pos_emb_0.norm(2).pow(2) +
            neg_emb_0.norm(2).pow(2)
        ) / users.shape[0]

        return bpr_loss + reg_loss

    def predict(self, users: torch.Tensor) -> torch.Tensor:
        """
        为指定用户生成所有物品的预测分数

        Args:
            users: 用户 ID [batch_size]
        Returns:
            预测分数矩阵 [batch_size, num_items]
        """
        user_emb, item_emb = self.compute_all_embeddings()
        u_emb = user_emb[users]  # [batch, dim]
        scores = torch.mm(u_emb, item_emb.T)  # [batch, num_items]
        return scores

def sample_negative(user_item_dict: dict, num_items: int, user: int) -> int:
    """随机负采样：采一个用户未交互过的物品"""
    while True:
        neg = random.randint(0, num_items - 1)
        if neg not in user_item_dict[user]:
            return neg

def train_lightgcn(model: LightGCN, train_pairs: np.ndarray,
                   n_epochs: int = 50, batch_size: int = 1024, lr: float = 1e-3):
    """
    LightGCN 训练循环

    Args:
        model: LightGCN 模型
        train_pairs: [N, 2] 训练交互对 (user, item)
        n_epochs: 训练轮数
        batch_size: 批大小
        lr: 学习率
    """
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    # 构建用户交互字典（用于负采样）
    user_item_dict = defaultdict(set)  # defaultdict访问不存在的键时返回默认值
    for u, i in train_pairs:
        user_item_dict[u].add(i)

    model.train()  # train()训练模式
    for epoch in range(n_epochs):
        # 随机打乱训练数据
        perm = np.random.permutation(len(train_pairs))
        total_loss = 0.0
        n_batches = 0

        for start in range(0, len(train_pairs), batch_size):
            end = min(start + batch_size, len(train_pairs))
            batch_idx = perm[start:end]
            batch_pairs = train_pairs[batch_idx]

            users = torch.LongTensor(batch_pairs[:, 0])
            pos_items = torch.LongTensor(batch_pairs[:, 1])
            neg_items = torch.LongTensor([
                sample_negative(user_item_dict, model.num_items, u)
                for u in batch_pairs[:, 0]
            ])

            loss = model.bpr_loss(users, pos_items, neg_items)

            optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新参数

            total_loss += loss.item()  # 将单元素张量转为Python数值
            n_batches += 1

        if (epoch + 1) % 10 == 0:
            avg_loss = total_loss / n_batches
            print(f"Epoch {epoch+1}/{n_epochs}, Loss: {avg_loss:.4f}")

# ---- 快速验证 ----
if __name__ == "__main__":
    # 模拟小规模数据
    num_users, num_items = 100, 200
    n_interactions = 1000

    train_pairs = np.column_stack([
        np.random.randint(0, num_users, n_interactions),
        np.random.randint(0, num_items, n_interactions)
    ])

    model = LightGCN(num_users, num_items, embedding_dim=64, n_layers=3)
    model.set_graph(train_pairs)

    print("开始训练 LightGCN...")
    train_lightgcn(model, train_pairs, n_epochs=50, batch_size=256)

    # 预测
    model.eval()
    with torch.no_grad():  # 禁用梯度计算，节省内存
        test_users = torch.LongTensor([0, 1, 2])
        scores = model.predict(test_users)
        print(f"\n预测分数矩阵形状: {scores.shape}")  # [3, 200]

        # 每个用户的 Top-5 推荐
        _, top_items = scores.topk(5, dim=1)
        for idx, uid in enumerate(test_users.tolist()):  # enumerate同时获取索引和元素
            print(f"用户 {uid} 的 Top-5 推荐: {top_items[idx].tolist()}")

import torch
import torch.nn as nn
import random
from collections import Counter

class PinSageLayer(nn.Module):
    """
    PinSage 单层聚合（简化版本）

    核心：随机游走采样 + Importance Pooling
    """
    def __init__(self, in_dim: int, out_dim: int):
        super().__init__()
        self.linear = nn.Linear(in_dim, out_dim)
        self.relu = nn.ReLU()

    def random_walk_sample(self, graph: dict, node: int,
                           walk_length: int = 10, num_walks: int = 5,
                           top_k: int = 15) -> list:
        """
        随机游走采样邻居

        Args:
            graph: 邻接表 {node: [neighbors]}
            node: 起始节点
            walk_length: 游走长度
            num_walks: 游走次数
            top_k: 选取的邻居数量
        Returns:
            [(neighbor, visit_count), ...] 按频次排序的 Top-K 邻居
        """
        visit_count = Counter()  # Counter统计元素出现次数

        for _ in range(num_walks):
            current = node
            for _ in range(walk_length):
                neighbors = graph.get(current, [])
                if not neighbors:
                    break
                current = random.choice(neighbors)
                if current != node:
                    visit_count[current] += 1

        # 返回频次最高的 Top-K 个邻居
        return visit_count.most_common(top_k)

    def forward(self, node_embs: torch.Tensor, neighbor_info: list) -> torch.Tensor:
        """
        Importance Pooling 聚合

        Args:
            node_embs: 所有节点的 Embedding [num_nodes, dim]
            neighbor_info: 每个目标节点的邻居信息
                           [[(neighbor_id, count), ...], ...]
        Returns:
            更新后的节点 Embedding [num_target_nodes, out_dim]
        """
        outputs = []
        for neighbors_with_counts in neighbor_info:
            if not neighbors_with_counts:
                # 无邻居，使用零向量
                outputs.append(torch.zeros(node_embs.shape[1]))
                continue

            ids, counts = zip(*neighbors_with_counts)  # zip按位置配对
            ids = list(ids)
            # 归一化频次作为权重
            weights = torch.FloatTensor(counts)
            weights = weights / weights.sum()

            # 加权聚合邻居 Embedding
            neighbor_emb = node_embs[ids]  # [K, dim]
            pooled = (weights.unsqueeze(1) * neighbor_emb).sum(dim=0)  # [dim]  # unsqueeze增加一个维度
            outputs.append(pooled)

        stacked = torch.stack(outputs, dim=0)  # [batch, dim]
        return self.relu(self.linear(stacked))

# ---- 使用示例 ----
if __name__ == "__main__":
    num_nodes = 1000
    dim = 64

    # 模拟邻接表
    graph = {i: random.sample(range(num_nodes), min(20, num_nodes))
             for i in range(num_nodes)}

    # 初始 Embedding
    node_embs = torch.randn(num_nodes, dim)

    # 对 batch 中的节点进行随机游走采样
    layer = PinSageLayer(dim, dim)
    target_nodes = [0, 1, 2, 3, 4]

    neighbor_info = [
        layer.random_walk_sample(graph, node, walk_length=10,
                                  num_walks=5, top_k=15)
        for node in target_nodes
    ]

    updated = layer(node_embs, neighbor_info)
    print(f"更新后 Embedding 形状: {updated.shape}")  # [5, 64]

import torch
import torch.nn as nn
import torch.nn.functional as F

class KGATLayer(nn.Module):
    """
    KGAT 注意力聚合层（简化版）

    对每个节点，根据不同关系类型计算邻居的注意力权重，
    加权聚合邻居信息。
    """
    def __init__(self, emb_dim: int, n_relations: int, dropout: float = 0.1):
        super().__init__()
        self.emb_dim = emb_dim

        # 每种关系一个变换矩阵
        self.W_r = nn.Parameter(torch.FloatTensor(n_relations, emb_dim, emb_dim))
        nn.init.xavier_uniform_(self.W_r.view(-1, emb_dim))  # 重塑张量形状

        # 注意力计算
        self.leaky_relu = nn.LeakyReLU(0.2)
        self.dropout = nn.Dropout(dropout)

    def compute_attention(self, head_emb: torch.Tensor, tail_emb: torch.Tensor,
                          relation_ids: torch.Tensor) -> torch.Tensor:
        """
        计算注意力分数

        Args:
            head_emb: 头实体 Embedding [num_edges, dim]
            tail_emb: 尾实体 Embedding [num_edges, dim]
            relation_ids: 关系类型 ID [num_edges]
        Returns:
            注意力分数 [num_edges]
        """
        # 获取对应关系的变换矩阵
        W = self.W_r[relation_ids]  # [num_edges, dim, dim]

        # 变换尾实体 Embedding
        tail_transformed = torch.bmm(W, tail_emb.unsqueeze(2)).squeeze(2)  # [num_edges, dim]  # squeeze压缩维度

        # 计算注意力分数
        scores = (head_emb * tail_transformed).sum(dim=1)  # [num_edges]
        scores = self.leaky_relu(scores)
        return scores

    def forward(self, entity_emb: torch.Tensor, edge_index: torch.Tensor,
                relation_ids: torch.Tensor) -> torch.Tensor:
        """
        前向传播

        Args:
            entity_emb: 实体 Embedding [num_entities, dim]
            edge_index: 边索引 [2, num_edges]，第 0 行是头实体，第 1 行是尾实体
            relation_ids: 关系类型 [num_edges]
        Returns:
            更新后的实体 Embedding [num_entities, dim]
        """
        heads, tails = edge_index[0], edge_index[1]
        head_emb = entity_emb[heads]  # [num_edges, dim]
        tail_emb = entity_emb[tails]  # [num_edges, dim]

        # 计算注意力分数
        attn_scores = self.compute_attention(head_emb, tail_emb, relation_ids)

        # 按头实体分组做 softmax（使用 scatter 实现）
        num_entities = entity_emb.shape[0]
        attn_weights = torch.zeros(heads.shape[0], device=entity_emb.device)

        # 简化的 softmax：对每个头节点的所有边做归一化
        max_scores = torch.zeros(num_entities, device=entity_emb.device)
        max_scores.scatter_reduce_(0, heads, attn_scores, reduce='amax', include_self=True)

        exp_scores = torch.exp(attn_scores - max_scores[heads])
        sum_exp = torch.zeros(num_entities, device=entity_emb.device)
        sum_exp.scatter_add_(0, heads, exp_scores)
        attn_weights = exp_scores / (sum_exp[heads] + 1e-8)
        attn_weights = self.dropout(attn_weights)

        # 获取对应关系的变换矩阵，变换尾实体
        W = self.W_r[relation_ids]
        tail_transformed = torch.bmm(W, tail_emb.unsqueeze(2)).squeeze(2)

        # 加权聚合
        weighted = attn_weights.unsqueeze(1) * tail_transformed  # [num_edges, dim]

        # 聚合到头实体
        output = torch.zeros(num_entities, self.emb_dim, device=entity_emb.device)
        output.scatter_add_(0, heads.unsqueeze(1).expand_as(weighted), weighted)

        return F.elu(output)

class KGAT(nn.Module):
    """
    简化的 KGAT 模型
    """
    def __init__(self, num_entities: int, num_relations: int,
                 emb_dim: int = 64, n_layers: int = 2):
        super().__init__()
        self.entity_embedding = nn.Embedding(num_entities, emb_dim)
        nn.init.xavier_uniform_(self.entity_embedding.weight)

        self.layers = nn.ModuleList([
            KGATLayer(emb_dim, num_relations) for _ in range(n_layers)
        ])

    def forward(self, edge_index: torch.Tensor,
                relation_ids: torch.Tensor) -> torch.Tensor:
        """
        Returns:
            最终实体 Embedding [num_entities, dim]
        """
        x = self.entity_embedding.weight
        emb_list = [x]

        for layer in self.layers:
            x = layer(x, edge_index, relation_ids)
            emb_list.append(x)

        # 层组合
        return torch.stack(emb_list, dim=0).mean(dim=0)

# ---- 使用示例 ----
if __name__ == "__main__":
    num_entities = 500   # 用户 + 物品 + KG 实体
    num_relations = 10   # 关系类型数
    num_edges = 2000

    # 随机生成知识图谱边
    edge_index = torch.stack([
        torch.randint(0, num_entities, (num_edges,)),
        torch.randint(0, num_entities, (num_edges,))
    ])
    relation_ids = torch.randint(0, num_relations, (num_edges,))

    model = KGAT(num_entities, num_relations, emb_dim=64, n_layers=2)
    entity_embs = model(edge_index, relation_ids)
    print(f"实体 Embedding 形状: {entity_embs.shape}")  # [500, 64]

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

class GraphContrastiveLearning(nn.Module):
    """
    图对比学习推荐框架

    支持两种模式：
    1. SGL 风格：通过边丢弃进行图增强
    2. SimGCL 风格：直接在 Embedding 上加噪声
    """
    def __init__(self, num_users: int, num_items: int, emb_dim: int = 64,
                 n_layers: int = 3, temperature: float = 0.2,
                 cl_weight: float = 0.1, noise_eps: float = 0.1,
                 mode: str = 'simgcl'):
        """
        Args:
            mode: 'sgl' 或 'simgcl'
        """
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.n_layers = n_layers
        self.temperature = temperature
        self.cl_weight = cl_weight
        self.noise_eps = noise_eps
        self.mode = mode

        self.user_embedding = nn.Embedding(num_users, emb_dim)
        self.item_embedding = nn.Embedding(num_items, emb_dim)
        nn.init.xavier_uniform_(self.user_embedding.weight)
        nn.init.xavier_uniform_(self.item_embedding.weight)

        self.adj_norm = None  # 原始归一化邻接矩阵

    def set_graph(self, adj_norm: torch.Tensor):
        """设置图的邻接矩阵"""
        self.adj_norm = adj_norm

    def lightgcn_propagate(self, adj: torch.Tensor,
                           add_noise: bool = False) -> tuple:
        """
        LightGCN 风格的图传播

        Args:
            adj: 邻接矩阵
            add_noise: 是否在每层传播后添加噪声（SimGCL）
        """
        all_emb = torch.cat([self.user_embedding.weight,
                             self.item_embedding.weight], dim=0)
        emb_list = [all_emb]

        for _ in range(self.n_layers):
            all_emb = torch.spmm(adj, all_emb)
            if add_noise:
                # SimGCL：添加均匀分布的归一化噪声
                noise = torch.rand_like(all_emb)
                noise = F.normalize(noise, dim=1) * self.noise_eps
                all_emb = all_emb + noise
            emb_list.append(all_emb)

        final_emb = torch.stack(emb_list, dim=0).mean(dim=0)
        user_emb = final_emb[:self.num_users]
        item_emb = final_emb[self.num_users:]
        return user_emb, item_emb

    def info_nce_loss(self, emb1: torch.Tensor, emb2: torch.Tensor,
                      indices: torch.Tensor) -> torch.Tensor:
        """
        InfoNCE 对比损失

        Args:
            emb1: 视图 1 的 Embedding [N, dim]
            emb2: 视图 2 的 Embedding [N, dim]
            indices: 目标节点索引 [batch_size]
        Returns:
            InfoNCE loss
        """
        z1 = F.normalize(emb1[indices], dim=1)  # [batch, dim]
        z2 = F.normalize(emb2[indices], dim=1)  # [batch, dim]

        # 正样本相似度
        pos_sim = (z1 * z2).sum(dim=1) / self.temperature  # [batch]

        # 所有节点作为负样本
        all_z2 = F.normalize(emb2, dim=1)  # [N, dim]
        neg_sim = torch.mm(z1, all_z2.T) / self.temperature  # [batch, N]

        # InfoNCE
        loss = -pos_sim + torch.logsumexp(neg_sim, dim=1)
        return loss.mean()

    def forward(self, users: torch.Tensor, pos_items: torch.Tensor,
                neg_items: torch.Tensor) -> torch.Tensor:
        """
        计算总损失 = BPR Loss + λ · CL Loss
        """
        # ---- 主任务：BPR Loss ----
        user_emb, item_emb = self.lightgcn_propagate(self.adj_norm, add_noise=False)

        u = user_emb[users]
        p = item_emb[pos_items]
        n = item_emb[neg_items]

        pos_scores = (u * p).sum(dim=1)
        neg_scores = (u * n).sum(dim=1)
        bpr_loss = -F.logsigmoid(pos_scores - neg_scores).mean()

        # ---- 对比学习：InfoNCE Loss ----
        if self.mode == 'simgcl':
            # SimGCL：两次带噪声的传播生成两个视图
            user_emb_1, item_emb_1 = self.lightgcn_propagate(
                self.adj_norm, add_noise=True)
            user_emb_2, item_emb_2 = self.lightgcn_propagate(
                self.adj_norm, add_noise=True)
        else:
            # SGL：需要两个增强图（这里简化为两次独立的 edge dropout）
            user_emb_1, item_emb_1 = self.lightgcn_propagate(
                self.adj_norm, add_noise=False)
            user_emb_2, item_emb_2 = self.lightgcn_propagate(
                self.adj_norm, add_noise=False)

        cl_loss_user = self.info_nce_loss(user_emb_1, user_emb_2, users)
        cl_loss_item = self.info_nce_loss(item_emb_1, item_emb_2, pos_items)
        cl_loss = cl_loss_user + cl_loss_item

        return bpr_loss + self.cl_weight * cl_loss

# ---- 使用示例 ----
if __name__ == "__main__":
    from scipy.sparse import coo_matrix

    num_users, num_items = 100, 200
    emb_dim = 64
    n_nodes = num_users + num_items

    # 构造随机交互图
    n_edges = 500
    users_arr = np.random.randint(0, num_users, n_edges)
    items_arr = np.random.randint(0, num_items, n_edges) + num_users

    row = np.concatenate([users_arr, items_arr])
    col = np.concatenate([items_arr, users_arr])
    data = np.ones(len(row), dtype=np.float32)
    adj = coo_matrix((data, (row, col)), shape=(n_nodes, n_nodes))

    degree = np.array(adj.sum(1)).flatten()
    d_inv_sqrt = np.power(degree + 1e-8, -0.5)
    adj.data = adj.data * d_inv_sqrt[adj.row] * d_inv_sqrt[adj.col]

    indices = torch.LongTensor(np.vstack([adj.row, adj.col]))
    values = torch.FloatTensor(adj.data)
    adj_norm = torch.sparse_coo_tensor(indices, values, torch.Size([n_nodes, n_nodes]))

    # 创建模型并训练
    model = GraphContrastiveLearning(num_users, num_items, emb_dim=emb_dim, mode='simgcl')
    model.set_graph(adj_norm)

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(20):
        users = torch.randint(0, num_users, (128,))
        pos_items = torch.randint(0, num_items, (128,))
        neg_items = torch.randint(0, num_items, (128,))

        loss = model(users, pos_items, neg_items)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

"""
LightGCN 推荐系统完整实战
数据集：MovieLens-100K

运行方式：python lightgcn_movielens.py
依赖：torch, numpy, scipy, pandas
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from scipy.sparse import coo_matrix
from collections import defaultdict
import random
import time
import os
import urllib.request
import zipfile

# ============================================================
# 第 1 部分：数据加载与预处理
# ============================================================

def download_movielens_100k(data_dir: str = './ml-100k'):
    """下载 MovieLens-100K 数据集"""
    if os.path.exists(os.path.join(data_dir, 'u.data')):
        print("数据集已存在，跳过下载。")
        return

    url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'
    zip_path = 'ml-100k.zip'

    print("正在下载 MovieLens-100K...")
    urllib.request.urlretrieve(url, zip_path)

    with zipfile.ZipFile(zip_path, 'r') as z:
        z.extractall('.')
    os.remove(zip_path)
    print("下载完成。")

def load_and_split_data(data_dir: str = './ml-100k', test_ratio: float = 0.2):
    """
    加载并划分数据集

    按时间排序后，每个用户最后 20% 的交互作为测试集

    Returns:
        train_pairs: [N_train, 2] (user_id, item_id)
        test_dict: {user_id: set(item_ids)} 测试集
        num_users, num_items
    """
    # 读取数据：user_id  item_id  rating  timestamp
    data = []
    filepath = os.path.join(data_dir, 'u.data')
    with open(filepath, 'r') as f:  # with自动管理文件关闭
        for line in f:
            parts = line.strip().split('\t')
            user, item, rating, ts = int(parts[0]), int(parts[1]), \
                                      float(parts[2]), int(parts[3])
            if rating >= 3.0:  # 只保留正向交互（评分 ≥ 3）
                data.append((user, item, ts))

    # 重新编号：从 0 开始
    users = sorted(set(d[0] for d in data))
    items = sorted(set(d[1] for d in data))
    user_map = {u: i for i, u in enumerate(users)}
    item_map = {it: i for i, it in enumerate(items)}

    num_users = len(users)
    num_items = len(items)

    # 按用户分组并按时间排序
    user_interactions = defaultdict(list)
    for user, item, ts in data:
        user_interactions[user_map[user]].append((item_map[item], ts))

    train_pairs = []
    test_dict = defaultdict(set)

    for user, interactions in user_interactions.items():
        # 按时间排序
        interactions.sort(key=lambda x: x[1])  # lambda匿名函数
        items_sorted = [it for it, _ in interactions]

        # 划分训练集和测试集
        n_test = max(1, int(len(items_sorted) * test_ratio))
        train_items = items_sorted[:-n_test]
        test_items = items_sorted[-n_test:]

        for item in train_items:
            train_pairs.append([user, item])
        for item in test_items:
            test_dict[user].add(item)

    train_pairs = np.array(train_pairs)

    print(f"用户数: {num_users}, 物品数: {num_items}")
    print(f"训练交互数: {len(train_pairs)}, 测试用户数: {len(test_dict)}")
    print(f"稀疏度: {1 - len(train_pairs) / (num_users * num_items):.4f}")

    return train_pairs, dict(test_dict), num_users, num_items

# ============================================================
# 第 2 部分：图构建
# ============================================================

def build_adj_matrix(train_pairs: np.ndarray, num_users: int,
                     num_items: int) -> torch.Tensor:
    """
    构建归一化邻接矩阵（用户-物品二部图）

    邻接矩阵结构：
    A = | 0   R |
        | R^T 0 |
    其中 R 是用户-物品交互矩阵

    归一化：Ã = D^{-1/2} A D^{-1/2}
    """
    n_nodes = num_users + num_items

    users = train_pairs[:, 0]
    items = train_pairs[:, 1] + num_users

    # 对称邻接矩阵
    row = np.concatenate([users, items])
    col = np.concatenate([items, users])
    data = np.ones(len(row), dtype=np.float32)

    adj = coo_matrix((data, (row, col)), shape=(n_nodes, n_nodes))

    # 归一化
    degree = np.array(adj.sum(axis=1)).flatten()
    d_inv_sqrt = np.power(degree, -0.5)
    d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0.0
    adj.data = adj.data * d_inv_sqrt[adj.row] * d_inv_sqrt[adj.col]

    # 转 PyTorch 稀疏张量
    indices = torch.LongTensor(np.vstack([adj.row, adj.col]))
    values = torch.FloatTensor(adj.data)
    return torch.sparse_coo_tensor(indices, values, torch.Size([n_nodes, n_nodes]))

# ============================================================
# 第 3 部分：LightGCN 模型
# ============================================================

class LightGCNFull(nn.Module):
    """带完整评估功能的 LightGCN"""

    def __init__(self, num_users: int, num_items: int,
                 emb_dim: int = 64, n_layers: int = 3, reg: float = 1e-4):
        super().__init__()
        self.num_users = num_users
        self.num_items = num_items
        self.n_layers = n_layers
        self.reg = reg

        self.user_emb = nn.Embedding(num_users, emb_dim)
        self.item_emb = nn.Embedding(num_items, emb_dim)
        nn.init.xavier_uniform_(self.user_emb.weight)
        nn.init.xavier_uniform_(self.item_emb.weight)

    def propagate(self, adj: torch.Tensor) -> tuple:
        """LightGCN 传播 + 层组合"""
        all_emb = torch.cat([self.user_emb.weight, self.item_emb.weight], dim=0)
        emb_list = [all_emb]

        for _ in range(self.n_layers):
            all_emb = torch.spmm(adj, all_emb)
            emb_list.append(all_emb)

        final = torch.stack(emb_list, dim=0).mean(dim=0)
        return final[:self.num_users], final[self.num_users:]

    def bpr_loss(self, adj: torch.Tensor, users: torch.Tensor,
                 pos: torch.Tensor, neg: torch.Tensor) -> torch.Tensor:
        user_e, item_e = self.propagate(adj)

        u = user_e[users]
        p = item_e[pos]
        n = item_e[neg]

        pos_s = (u * p).sum(dim=1)
        neg_s = (u * n).sum(dim=1)
        loss = -F.logsigmoid(pos_s - neg_s).mean()

        # L2 正则（仅初始 Embedding）
        reg = self.reg * (
            self.user_emb(users).norm(2).pow(2) +
            self.item_emb(pos).norm(2).pow(2) +
            self.item_emb(neg).norm(2).pow(2)
        ) / len(users)

        return loss + reg

# ============================================================
# 第 4 部分：评估指标
# ============================================================

def evaluate(model: LightGCNFull, adj: torch.Tensor,
             train_user_items: dict, test_dict: dict,
             top_k: int = 20) -> dict:
    """
    评估推荐模型（Recall@K, NDCG@K）

    Args:
        model: 训练好的模型
        adj: 归一化邻接矩阵
        train_user_items: {user: set(items)} 训练集交互
        test_dict: {user: set(items)} 测试集真实交互
        top_k: 评估的 K 值
    Returns:
        {'recall@K': ..., 'ndcg@K': ...}
    """
    model.eval()

    with torch.no_grad():
        user_emb, item_emb = model.propagate(adj)

    recalls = []
    ndcgs = []

    for user in test_dict:
        if user not in train_user_items:
            continue

        # 计算该用户对所有物品的分数
        scores = torch.mv(item_emb, user_emb[user])  # [num_items]

        # 屏蔽训练集中已交互的物品
        train_items = list(train_user_items[user])
        scores[train_items] = -float('inf')

        # Top-K 推荐
        _, topk_items = scores.topk(top_k)
        topk_items = topk_items.cpu().numpy().tolist()  # 保持有序列表，NDCG需要排名信息

        # 计算 Recall@K
        true_items = test_dict[user]
        hits = len(set(topk_items) & true_items)
        recall = hits / min(len(true_items), top_k)
        recalls.append(recall)

        # 计算 NDCG@K
        dcg = 0.0
        for rank, item in enumerate(topk_items):
            if item in true_items:
                dcg += 1.0 / np.log2(rank + 2)  # rank 从 0 开始

        idcg = sum(1.0 / np.log2(i + 2) for i in range(min(len(true_items), top_k)))
        ndcg = dcg / idcg if idcg > 0 else 0.0
        ndcgs.append(ndcg)

    return {
        f'Recall@{top_k}': np.mean(recalls),
        f'NDCG@{top_k}': np.mean(ndcgs)
    }

# ============================================================
# 第 5 部分：训练与评估主流程
# ============================================================

def main():
    # 超参数
    EMB_DIM = 64
    N_LAYERS = 3
    LR = 1e-3
    REG = 1e-4
    BATCH_SIZE = 1024
    N_EPOCHS = 80
    TOP_K = 20

    # 1. 数据加载
    download_movielens_100k()
    train_pairs, test_dict, num_users, num_items = load_and_split_data()

    # 2. 构建图
    adj = build_adj_matrix(train_pairs, num_users, num_items)

    # 构建训练集用户交互字典
    train_user_items = defaultdict(set)
    for u, i in train_pairs:
        train_user_items[u].add(i)

    # 3. 创建模型
    model = LightGCNFull(num_users, num_items, EMB_DIM, N_LAYERS, REG)
    optimizer = torch.optim.Adam(model.parameters(), lr=LR)

    # 4. 训练循环
    print("\n" + "=" * 60)
    print("开始训练 LightGCN")
    print("=" * 60)

    best_recall = 0.0

    for epoch in range(N_EPOCHS):
        model.train()
        perm = np.random.permutation(len(train_pairs))
        total_loss = 0.0
        n_batch = 0

        for start in range(0, len(train_pairs), BATCH_SIZE):
            end = min(start + BATCH_SIZE, len(train_pairs))
            batch = train_pairs[perm[start:end]]

            users = torch.LongTensor(batch[:, 0])
            pos_items = torch.LongTensor(batch[:, 1])
            neg_items = torch.LongTensor([
                sample_neg(train_user_items, num_items, u)
                for u in batch[:, 0]
            ])

            loss = model.bpr_loss(adj, users, pos_items, neg_items)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            n_batch += 1

        avg_loss = total_loss / n_batch

        # 每 10 个 epoch 评估一次
        if (epoch + 1) % 10 == 0:
            metrics = evaluate(model, adj, train_user_items, test_dict, TOP_K)
            recall = metrics[f'Recall@{TOP_K}']
            ndcg = metrics[f'NDCG@{TOP_K}']

            improved = " ★ New Best!" if recall > best_recall else ""
            best_recall = max(best_recall, recall)

            print(f"Epoch {epoch+1:3d} | Loss: {avg_loss:.4f} | "
                  f"Recall@{TOP_K}: {recall:.4f} | "
                  f"NDCG@{TOP_K}: {ndcg:.4f}{improved}")

    # 5. 最终评估
    print("\n" + "=" * 60)
    print("最终评估结果")
    print("=" * 60)

    for k in [5, 10, 20]:
        metrics = evaluate(model, adj, train_user_items, test_dict, k)
        print(f"  Recall@{k:2d}: {metrics[f'Recall@{k}']:.4f}  |  "
              f"NDCG@{k:2d}: {metrics[f'NDCG@{k}']:.4f}")

    # 6. 推荐示例
    print("\n" + "=" * 60)
    print("推荐示例")
    print("=" * 60)

    model.eval()
    with torch.no_grad():
        user_emb, item_emb = model.propagate(adj)

    for uid in [0, 1, 2]:
        scores = torch.mv(item_emb, user_emb[uid])
        train_items = list(train_user_items.get(uid, set()))
        scores[train_items] = -float('inf')
        _, topk = scores.topk(10)
        print(f"  用户 {uid} 的 Top-10 推荐: {topk.tolist()}")

def sample_neg(user_items: dict, num_items: int, user: int) -> int:
    """负采样"""
    while True:
        neg = random.randint(0, num_items - 1)
        if neg not in user_items[user]:
            return neg

if __name__ == "__main__":
    random.seed(42)
    np.random.seed(42)
    torch.manual_seed(42)
    main()