data.py

"""
数据层：生成合成用户-物品交互数据，模拟真实推荐场景
包含用户特征、物品特征、交互记录（点击/购买等）
"""
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader


def generate_synthetic_data(
    num_users=10000,
    num_items=5000,
    num_interactions=200000,
    num_user_features=16,
    num_item_features=32,
    seed=42,
):
    """
    生成合成推荐数据
    - 用户特征：年龄段、性别、消费等级等 one-hot 编码
    - 物品特征：类别、价格档位、品牌等 one-hot 编码
    - 交互记录：user_id, item_id, label(0/1), 位置用于后续特征
    """
    np.random.seed(seed)
    torch.manual_seed(seed)

    # 用户画像
    user_features = torch.randn(num_users, num_user_features)
    # 让用户有聚类结构（模拟真实人群分组）
    for i in range(5):
        center = torch.randn(num_user_features) * 3
        start = i * (num_users // 5)
        end = (i + 1) * (num_users // 5)
        user_features[start:end] += center

    # 物品画像
    item_features = torch.randn(num_items, num_item_features)
    for i in range(10):
        center = torch.randn(num_item_features) * 3
        start = i * (num_items // 10)
        end = (i + 1) * (num_items // 10)
        item_features[start:end] += center

    # 生成交互记录：正样本 + 负样本
    # 正样本：用户和物品特征相似度高的更容易交互
    user_ids = np.random.randint(0, num_users, num_interactions)
    item_ids = np.random.randint(0, num_items, num_interactions)

    # 用特征点积 + 噪声模拟真实交互概率
    scores = []
    for u, i in zip(user_ids, item_ids):
        # 取用户前 min 个维度与物品特征做点积
        dim = min(num_user_features, num_item_features)
        score = (user_features[u, :dim] * item_features[i, :dim]).sum().item()
        scores.append(score)

    scores = np.array(scores)
    # sigmoid 转概率，加噪声，生成标签
    probs = 1 / (1 + np.exp(-scores / dim))
    labels = (np.random.random(num_interactions) < probs).astype(np.float32)

    interactions = {
        "user_ids": torch.tensor(user_ids, dtype=torch.long),
        "item_ids": torch.tensor(item_ids, dtype=torch.long),
        "labels": torch.tensor(labels, dtype=torch.float),
    }

    return {
        "num_users": num_users,
        "num_items": num_items,
        "num_user_features": num_user_features,
        "num_item_features": num_item_features,
        "user_features": user_features,
        "item_features": item_features,
        "interactions": interactions,
    }


class InteractionDataset(Dataset):
    """交互数据集，供 DataLoader 使用"""

    def __init__(self, interactions, user_features, item_features):
        self.user_ids = interactions["user_ids"]
        self.item_ids = interactions["item_ids"]
        self.labels = interactions["labels"]
        self.user_features = user_features
        self.item_features = item_features

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, idx):
        u = self.user_ids[idx]
        i = self.item_ids[idx]
        return {
            "user_id": u,
            "item_id": i,
            "user_feat": self.user_features[u],
            "item_feat": self.item_features[i],
            "label": self.labels[idx],
        }


def split_data(interactions, train_ratio=0.8, batch_size=256):
    """划分训练/测试集，返回 DataLoader"""
    n = len(interactions["labels"])
    perm = torch.randperm(n)
    split = int(n * train_ratio)

    train_inter = {k: v[perm[:split]] for k, v in interactions.items()}
    test_inter = {k: v[perm[split:]] for k, v in interactions.items()}

    return train_inter, test_inter


def get_all_items_for_user(user_id, num_items, item_features, exclude=None):
    """
    为一个用户生成所有候选物品（用于召回阶段）
    exclude: 已交互过的 item_id 集合
    """
    all_items = torch.arange(num_items)
    if exclude is not None:
        mask = ~torch.tensor([i.item() in exclude for i in all_items])
        all_items = all_items[mask]
    return all_items