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

class DynamicRouting(nn.Module):  # 继承nn.Module定义网络层
    """胶囊网络动态路由机制，用于提取多兴趣表征"""

    def __init__(self, input_dim, num_caps, iter_rounds=3):
        """
        Args:
            input_dim: 输入 embedding 维度
            num_caps: 兴趣胶囊数量 K
            iter_rounds: 动态路由迭代轮数
        """
        super().__init__()  # super()调用父类方法
        self.input_dim = input_dim
        self.num_caps = num_caps
        self.iter_rounds = iter_rounds
        # 共享双线性变换矩阵
        self.S = nn.Parameter(torch.randn(input_dim, input_dim) * 0.01)

    @staticmethod  # @staticmethod不需要实例即可调用
    def squash(s):
        """Squash 激活函数：将向量长度压缩到 (0, 1)"""
        s_norm = torch.norm(s, dim=-1, keepdim=True)
        return (s_norm ** 2 / (1 + s_norm ** 2)) * (s / (s_norm + 1e-8))

    def forward(self, item_embs, mask=None):
        """
        Args:
            item_embs: (batch, seq_len, dim) 行为序列 embedding
            mask: (batch, seq_len) 有效 item 的 mask
        Returns:
            caps: (batch, num_caps, dim) K 个兴趣胶囊
        """
        batch, seq_len, dim = item_embs.shape
        # 变换: (batch, seq_len, dim)
        u_hat = torch.matmul(item_embs, self.S)  # (B, N, D)

        # 初始化路由 logits: (batch, seq_len, num_caps)
        b = torch.zeros(batch, seq_len, self.num_caps, device=item_embs.device)

        for t in range(self.iter_rounds):
            if mask is not None:
                # 被 mask 的位置路由权重设为极小值
                b = b.masked_fill(~mask.unsqueeze(-1), float('-inf'))  # unsqueeze增加一个维度
            # 数值稳定性：可添加 logits = logits.clamp(-10, 10)
            b = b.clamp(-10, 10)
            c = F.softmax(b, dim=-1)  # (B, N, K)  # F.xxx PyTorch函数式API
            if mask is not None:
                c = c.masked_fill(~mask.unsqueeze(-1), 0.0)

            # 加权聚合: (B, K, D)
            s = torch.einsum('bnk,bnd->bkd', c, u_hat)
            caps = self.squash(s)  # (B, K, D)

            if t < self.iter_rounds - 1:
                # 更新路由 logits
                delta = torch.einsum('bkd,bnd->bnk', caps, u_hat)
                b = b + delta

        return caps

class MIND(nn.Module):
    """MIND: Multi-Interest Network with Dynamic Routing"""

    def __init__(self, num_items, embed_dim, num_interests=4, routing_iters=3, padding_idx=0):
        super().__init__()
        self.num_interests = num_interests
        self.item_embedding = nn.Embedding(num_items, embed_dim, padding_idx=padding_idx)
        self.routing = DynamicRouting(embed_dim, num_interests, routing_iters)
        self.output_mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )

    def forward(self, behavior_seq, target_item=None):
        """
        Args:
            behavior_seq: (batch, seq_len) 行为序列 item id
            target_item: (batch,) 目标 item id（训练时使用）
        Returns:
            训练: 用户向量 (batch, dim)
            推理: K 个兴趣向量 (batch, K, dim)
        """
        mask = (behavior_seq != 0)  # padding mask
        item_embs = self.item_embedding(behavior_seq)  # (B, N, D)
        interest_caps = self.routing(item_embs, mask)   # (B, K, D)

        if target_item is not None and self.training:
            # Label-Aware Attention
            target_emb = self.item_embedding(target_item)  # (B, D)
            attn_scores = torch.einsum('bkd,bd->bk', interest_caps, target_emb)
            attn_weights = F.softmax(attn_scores, dim=-1)  # (B, K)
            user_vec = torch.einsum('bk,bkd->bd', attn_weights, interest_caps)
            return self.output_mlp(user_vec)
        else:
            # 推理：返回 K 个兴趣向量
            return interest_caps

# ========== 使用示例 ==========
def mind_example():
    model = MIND(num_items=10000, embed_dim=64, num_interests=4)

    # 模拟一个 batch
    behavior = torch.randint(1, 10000, (32, 50))  # 50 个历史行为
    target = torch.randint(1, 10000, (32,))

    # 训练模式
    model.train()  # train()训练模式
    user_vec = model(behavior, target)  # (32, 64)
    print(f"训练 - 用户向量: {user_vec.shape}")

    # 推理模式
    model.eval()
    with torch.no_grad():  # 禁用梯度计算，节省内存
        interests = model(behavior)  # (32, 4, 64)
        print(f"推理 - 多兴趣向量: {interests.shape}")

if __name__ == "__main__":
    mind_example()

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

class ComiRecSA(nn.Module):
    """ComiRec-SA: 基于多头自注意力的多兴趣提取"""

    def __init__(self, num_items, embed_dim, num_interests=4,
                 hidden_dim=64, padding_idx=0):
        super().__init__()
        self.num_interests = num_interests
        self.item_embedding = nn.Embedding(num_items, embed_dim, padding_idx=padding_idx)
        # 自注意力参数
        self.W1 = nn.Linear(embed_dim, hidden_dim, bias=False)
        self.W2 = nn.Linear(hidden_dim, num_interests, bias=False)

    def forward(self, behavior_seq, target_item=None):
        """
        Args:
            behavior_seq: (batch, seq_len) 行为序列
            target_item: (batch,) 目标 item（训练时）
        """
        mask = (behavior_seq != 0)
        H = self.item_embedding(behavior_seq)  # (B, N, D)

        # 多头注意力提取兴趣
        attn = self.W2(torch.tanh(self.W1(H)))  # (B, N, K)
        # mask padding
        attn = attn.masked_fill(~mask.unsqueeze(-1), float('-inf'))
        A = F.softmax(attn, dim=1)  # (B, N, K)
        A = A.transpose(1, 2)  # (B, K, N)

        interest_vecs = torch.bmm(A, H)  # (B, K, D)

        if target_item is not None and self.training:
            target_emb = self.item_embedding(target_item)  # (B, D)
            scores = torch.einsum('bkd,bd->bk', interest_vecs, target_emb)
            weights = F.softmax(scores, dim=-1)
            user_vec = torch.einsum('bk,bkd->bd', weights, interest_vecs)
            return user_vec
        return interest_vecs

    @staticmethod
    def controllable_select(interest_vecs, item_embs, top_n=50, lam=0.3):
        """
        可控多样性聚合
        Args:
            interest_vecs: (K, D) 单个用户的 K 个兴趣向量
            item_embs: (M, D) 候选 item embedding（多兴趣合并后）
            top_n: 最终选择数量
            lam: 多样性控制参数 [0, 1]
        Returns:
            selected: 选中的 item 索引列表
        """
        K, D = interest_vecs.shape
        M = item_embs.shape[0]

        # 每个 item 与最相关兴趣的相似度
        relevance = torch.mm(item_embs, interest_vecs.T).max(dim=1)[0]  # (M,)

        selected = []
        selected_embs = []

        for _ in range(min(top_n, M)):
            if len(selected_embs) == 0:
                diversity = torch.zeros(M, device=item_embs.device)
            else:
                stacked = torch.stack(selected_embs)  # (|S|, D)  # torch.stack沿新维度拼接张量
                diversity = torch.mm(item_embs, stacked.T).max(dim=1)[0]

            score = (1 - lam) * relevance - lam * diversity
            # 已选中的设为 -inf
            for idx in selected:
                score[idx] = float('-inf')

            best = score.argmax().item()  # 将单元素张量转为Python数值
            selected.append(best)
            selected_embs.append(item_embs[best])

        return selected

class ConceptPool(nn.Module):
    """SINE 全局兴趣概念池"""

    def __init__(self, num_concepts, embed_dim, top_k=4):
        super().__init__()
        self.num_concepts = num_concepts
        self.top_k = top_k
        # 可学习的概念原型
        self.concepts = nn.Parameter(torch.randn(num_concepts, embed_dim) * 0.02)

    def forward(self, user_repr):
        """
        Args:
            user_repr: (batch, dim) 用户粗粒度表征
        Returns:
            activated: (batch, top_k, dim) 激活的概念向量
            indices: (batch, top_k) 激活的概念索引
        """
        # 计算与所有概念的相似度
        logits = torch.mm(user_repr, self.concepts.T)  # (B, P)
        # 稀疏选择 top-K
        topk_vals, topk_ids = torch.topk(logits, self.top_k, dim=-1)
        # Gumbel-Softmax 或直接 softmax
        weights = F.softmax(topk_vals, dim=-1)  # (B, K)
        # 取出激活的概念
        activated = self.concepts[topk_ids]  # (B, K, D)
        activated = activated * weights.unsqueeze(-1)
        return activated, topk_ids

def multi_interest_retrieval(user_interests, ann_index, top_n_per_interest=100,
                              final_top_n=300):
    """
    多兴趣向量 ANN 召回流程

    Args:
        user_interests: list of np.array, K 个兴趣向量
        ann_index: Faiss / Milvus ANN 索引
        top_n_per_interest: 每个兴趣检索数量
        final_top_n: 最终返回数量
    """
    all_candidates = {}

    # Step 1: 每个兴趣向量独立检索
    for i, vec in enumerate(user_interests):  # enumerate同时获取索引和元素
        distances, ids = ann_index.search(
            vec.reshape(1, -1), top_n_per_interest  # 重塑张量形状
        )
        for dist, item_id in zip(distances[0], ids[0]):  # zip按位置配对
            if item_id not in all_candidates or dist < all_candidates[item_id]['dist']:
                all_candidates[item_id] = {
                    'dist': dist,
                    'interest_id': i,
                    'score': float(-dist)  # 距离越小越好
                }

    # Step 2: 去重 + 按分数排序
    sorted_items = sorted(all_candidates.items(), key=lambda x: x[1]['score'], reverse=True)  # lambda匿名函数

    # Step 3: 打散（确保多兴趣都有曝光）
    result = interest_aware_shuffle(sorted_items, num_interests=len(user_interests))

    return result[:final_top_n]

def interest_aware_shuffle(sorted_items, num_interests, window_size=10):
    """
    兴趣打散策略：在每个窗口内确保不同兴趣来源的候选都有出现
    """
    result = []
    remaining = list(sorted_items)

    while remaining and len(result) < len(sorted_items):
        window = remaining[:window_size * num_interests]
        # 按兴趣来源分组
        by_interest = {}
        for item_id, info in window:
            iid = info['interest_id']
            if iid not in by_interest:
                by_interest[iid] = []
            by_interest[iid].append((item_id, info))

        # 轮流从每个兴趣取一个
        added = set()
        for _ in range(window_size):
            for iid in range(num_interests):
                if iid in by_interest and by_interest[iid]:
                    item = by_interest[iid].pop(0)
                    if item[0] not in added:
                        result.append(item)
                        added.add(item[0])

        # 移除已添加的
        remaining = [(iid, info) for iid, info in remaining if iid not in added]

    return result

import numpy as np

def estimate_propensity_naive(interaction_matrix):
    """
    基于流行度的简单倾向得分估计
    Args:
        interaction_matrix: (num_users, num_items) 交互矩阵，1表示有交互
    Returns:
        propensity: (num_items,) 每个 item 的倾向得分
    """
    item_pop = interaction_matrix.sum(axis=0)  # 每个 item 的交互次数
    # 归一化到 (0, 1]
    propensity = item_pop / item_pop.max()
    # 裁剪避免极端值
    propensity = np.clip(propensity, 0.01, 1.0)
    return propensity

import torch
import torch.nn as nn

class IPSWeightedLoss(nn.Module):
    """IPS 加权损失函数"""

    def __init__(self, clip_min=0.01, clip_max=1.0):
        super().__init__()
        self.clip_min = clip_min
        self.clip_max = clip_max

    def forward(self, predictions, labels, propensity_scores):
        """
        Args:
            predictions: (batch,) 模型预测值
            labels: (batch,) 真实标签
            propensity_scores: (batch,) 倾向得分
        """
        # 裁剪倾向得分，防止方差爆炸
        clipped_ps = torch.clamp(propensity_scores, self.clip_min, self.clip_max)
        # IPS 权重
        weights = 1.0 / clipped_ps
        # 归一化权重（Self-Normalized IPS，降低方差）
        weights = weights / weights.sum() * len(weights)
        # 加权 BCE 损失
        bce = nn.functional.binary_cross_entropy_with_logits(
            predictions, labels, reduction='none'
        )
        return (weights * bce).mean()

# 使用示例
def train_with_ips(model, optimizer, dataloader, propensity_fn):
    """IPS 加权训练循环"""
    ips_loss_fn = IPSWeightedLoss(clip_min=0.05)
    model.train()

    for batch in dataloader:
        user_ids, item_ids, labels = batch
        predictions = model(user_ids, item_ids).squeeze()  # squeeze压缩维度

        # 获取倾向得分
        ps = propensity_fn(item_ids)

        loss = ips_loss_fn(predictions, labels.float(), ps)
        optimizer.zero_grad()  # 清零梯度
        loss.backward()  # 反向传播计算梯度
        optimizer.step()  # 更新参数

import torch
import torch.nn as nn

class DoublyRobustLoss(nn.Module):
    """Doubly Robust 损失函数"""

    def __init__(self, clip_min=0.05, clip_max=1.0):
        super().__init__()
        self.clip_min = clip_min
        self.clip_max = clip_max

    def forward(self, predictions, labels, propensity_scores,
                imputed_labels, observed_mask):
        """
        Args:
            predictions: (batch,) 模型预测
            labels: (batch,) 真实标签（仅 observed 部分有效）
            propensity_scores: (batch,) 倾向得分
            imputed_labels: (batch,) 填补模型的预测
            observed_mask: (batch,) 是否被观测到 (0/1)
        """
        clipped_ps = torch.clamp(propensity_scores, self.clip_min, self.clip_max)

        # Direct Method 部分：对所有样本
        loss_dm = nn.functional.binary_cross_entropy_with_logits(
            predictions, imputed_labels, reduction='none'
        )

        # IPS 修正部分：仅对观测样本
        loss_observed = nn.functional.binary_cross_entropy_with_logits(
            predictions, labels, reduction='none'
        )
        loss_imputed = nn.functional.binary_cross_entropy_with_logits(
            predictions, imputed_labels, reduction='none'
        )
        residual = (loss_observed - loss_imputed) * observed_mask / clipped_ps

        # DR = DM + IPS correction
        dr_loss = loss_dm + residual
        return dr_loss.mean()

class DoublyRobustTrainer:
    """DR 训练框架"""

    def __init__(self, main_model, imputation_model, propensity_model):
        self.main_model = main_model
        self.imputation_model = imputation_model
        self.propensity_model = propensity_model
        self.dr_loss = DoublyRobustLoss()

    def train_step(self, batch, optimizer):
        user_ids, item_ids, labels, observed_mask = batch

        # Step 1: 获取填补预测（detach，不参与梯度）
        with torch.no_grad():
            imputed = torch.sigmoid(self.imputation_model(user_ids, item_ids).squeeze())
            propensity = self.propensity_model(item_ids)

        # Step 2: 主模型前向
        predictions = self.main_model(user_ids, item_ids).squeeze()

        # Step 3: DR 损失
        loss = self.dr_loss(predictions, labels.float(), propensity,
                            imputed, observed_mask.float())

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return loss.item()

class DICE(nn.Module):
    """DICE: 解耦兴趣和遵从性"""

    def __init__(self, num_users, num_items, embed_dim):
        super().__init__()
        # 兴趣 embedding
        self.user_interest = nn.Embedding(num_users, embed_dim)
        self.item_interest = nn.Embedding(num_items, embed_dim)
        # 遵从性 embedding
        self.user_conformity = nn.Embedding(num_users, embed_dim)
        self.item_popularity = nn.Embedding(num_items, embed_dim)

    def forward(self, user_ids, item_ids, mode='train'):
        # 兴趣得分
        u_int = self.user_interest(user_ids)
        i_int = self.item_interest(item_ids)
        interest_score = (u_int * i_int).sum(dim=-1)

        if mode == 'inference':
            # 推理时仅使用兴趣得分
            return interest_score

        # 遵从性得分
        u_conf = self.user_conformity(user_ids)
        i_pop = self.item_popularity(item_ids)
        conformity_score = (u_conf * i_pop).sum(dim=-1)

        return interest_score, conformity_score

    def compute_loss(self, user_ids, pos_items, neg_items):
        """
        DICE 训练：正负样本需要区分 click 来源
        """
        # 正样本
        int_pos, conf_pos = self.forward(user_ids, pos_items, mode='train')
        # 负样本
        int_neg, conf_neg = self.forward(user_ids, neg_items, mode='train')

        # 兴趣损失：正 > 负
        interest_loss = -torch.log(
            torch.sigmoid(int_pos - int_neg) + 1e-8
        ).mean()

        # 遵从性损失：正样本中的从众部分不应过大
        # 使用因果正则化
        conformity_loss = -torch.log(
            torch.sigmoid(conf_pos) + 1e-8
        ).mean()

        # 解耦正则：兴趣和遵从 embedding 应尽量正交
        discrepancy = (
            (self.user_interest.weight * self.user_conformity.weight).sum(dim=-1).pow(2).mean()
            + (self.item_interest.weight * self.item_popularity.weight).sum(dim=-1).pow(2).mean()
        )

        return interest_loss + 0.1 * conformity_loss + 0.01 * discrepancy

class CausE(nn.Module):
    """CausE: 因果 Embedding"""

    def __init__(self, num_users, num_items, embed_dim):
        super().__init__()
        self.user_emb = nn.Embedding(num_users, embed_dim)
        # 两套 item embedding
        self.item_emb_control = nn.Embedding(num_items, embed_dim)  # 常规数据
        self.item_emb_treatment = nn.Embedding(num_items, embed_dim)  # 随机数据

    def forward(self, user_ids, item_ids, is_random=False):
        u = self.user_emb(user_ids)
        if is_random:
            i = self.item_emb_treatment(item_ids)
        else:
            i = self.item_emb_control(item_ids)
        return (u * i).sum(dim=-1)

    def compute_loss(self, batch_control, batch_treatment, reg_weight=0.1):
        """
        Args:
            batch_control: (user, item, label) 常规推荐数据
            batch_treatment: (user, item, label) 随机曝光数据
        """
        u_c, i_c, y_c = batch_control
        u_t, i_t, y_t = batch_treatment

        # 常规数据损失
        pred_c = self.forward(u_c, i_c, is_random=False)
        loss_c = nn.functional.binary_cross_entropy_with_logits(pred_c, y_c.float())

        # 随机数据损失
        pred_t = self.forward(u_t, i_t, is_random=True)
        loss_t = nn.functional.binary_cross_entropy_with_logits(pred_t, y_t.float())

        # 因果正则：两套 item embedding 不应差异过大
        shared_items = torch.unique(torch.cat([i_c, i_t]))  # torch.cat沿已有维度拼接张量
        emb_diff = (
            self.item_emb_control(shared_items) - self.item_emb_treatment(shared_items)
        ).pow(2).mean()

        return loss_c + loss_t + reg_weight * emb_diff