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

class EmbeddingLayer(nn.Module):  # 继承nn.Module定义网络层
    """通用稀疏特征 Embedding 层"""
    def __init__(self, feature_dims: dict, embed_dim: int = 32):
        super().__init__()  # super()调用父类方法
        self.embeddings = nn.ModuleDict({
            name: nn.Embedding(dim, embed_dim)
            for name, dim in feature_dims.items()
        })
        self.output_dim = len(feature_dims) * embed_dim

    def forward(self, features: dict) -> torch.Tensor:
        embs = [self.embeddings[n](features[n]) for n in self.embeddings]
        return torch.cat(embs, dim=-1)  # torch.cat沿已有维度拼接张量

class Tower(nn.Module):
    """单塔 MLP"""
    def __init__(self, input_dim, hidden_dims, output_dim, dropout=0.1):
        super().__init__()
        layers = []
        prev = input_dim
        for h in hidden_dims:
            layers += [nn.Linear(prev, h), nn.BatchNorm1d(h), nn.ReLU(), nn.Dropout(dropout)]
            prev = h
        layers.append(nn.Linear(prev, output_dim))
        self.mlp = nn.Sequential(*layers)

    def forward(self, x):
        return F.normalize(self.mlp(x), p=2, dim=-1)  # F.xxx PyTorch函数式API

class TwoTowerPreRanking(nn.Module):
    """双塔粗排模型"""
    def __init__(self, user_feat_dims, item_feat_dims, embed_dim=16,
                 user_hidden=[128, 64], item_hidden=[128, 64],
                 tower_out=32, temperature=0.05):
        super().__init__()
        self.user_emb = EmbeddingLayer(user_feat_dims, embed_dim)
        self.item_emb = EmbeddingLayer(item_feat_dims, embed_dim)
        self.user_tower = Tower(self.user_emb.output_dim, user_hidden, tower_out)
        self.item_tower = Tower(self.item_emb.output_dim, item_hidden, tower_out)
        self.temperature = temperature

    def get_user_embedding(self, user_features):
        return self.user_tower(self.user_emb(user_features))

    def get_item_embedding(self, item_features):
        return self.item_tower(self.item_emb(item_features))

    def forward(self, user_features, item_features):
        u = self.get_user_embedding(user_features)   # [B, d]
        v = self.get_item_embedding(item_features)   # [B, d]
        return torch.sum(u * v, dim=-1) / self.temperature

    def batch_score(self, user_emb, item_embs):
        """在线批量打分：user_emb [1,d], item_embs [N,d] → [N]"""
        return torch.matmul(item_embs, user_emb.squeeze(0)) / self.temperature  # squeeze压缩维度

class PreRankingTrainer:
    """In-batch Negative 对比学习训练"""
    def __init__(self, model, lr=1e-3):
        self.model = model
        self.optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    def train_step(self, user_features, item_features):
        self.model.train()  # train()训练模式
        self.optimizer.zero_grad()  # 清零梯度
        u = self.model.get_user_embedding(user_features)  # [B, d]
        v = self.model.get_item_embedding(item_features)  # [B, d]
        # B×B 相似度矩阵，对角线为正样本对
        sim = torch.matmul(u, v.T) / self.model.temperature
        target = torch.arange(u.size(0), device=u.device)
        loss = F.cross_entropy(sim, target)  # F.cross_entropy PyTorch函数式交叉熵损失
        loss.backward()  # 反向传播计算梯度
        self.optimizer.step()  # 更新参数
        acc = (sim.argmax(dim=1) == target).float().mean().item()  # 将单元素张量转为Python数值
        return {"loss": loss.item(), "accuracy": acc}

class DistillationLoss(nn.Module):
    """综合蒸馏损失：soft + hard + feature"""
    def __init__(self, alpha=0.3, beta=0.3, gamma=0.4, temperature=3.0):
        super().__init__()
        self.alpha, self.beta, self.gamma = alpha, beta, gamma
        self.T = temperature

    def forward(self, s_logit, t_logit, label, s_feat=None, t_feat=None, proj=None):
        # Hard Label Loss
        l_hard = F.binary_cross_entropy_with_logits(s_logit, label)
        # Soft Label Loss (KL散度)
        s_soft = F.log_softmax(torch.stack([s_logit, -s_logit], -1) / self.T, -1)  # torch.stack沿新维度拼接张量
        t_soft = F.softmax(torch.stack([t_logit, -t_logit], -1) / self.T, -1)
        l_soft = F.kl_div(s_soft, t_soft, reduction="batchmean") * (self.T ** 2)
        # Feature Distillation
        l_feat = torch.tensor(0.0, device=s_logit.device)
        if s_feat is not None and t_feat is not None and proj is not None:
            l_feat = F.mse_loss(proj(s_feat), t_feat.detach())  # 分离计算图，不参与梯度计算
        return self.alpha * l_hard + self.beta * l_soft + self.gamma * l_feat

class SEBlock(nn.Module):
    """Squeeze-and-Excitation：动态特征字段选择"""
    def __init__(self, num_fields, reduction=4):
        super().__init__()
        self.squeeze = nn.Linear(num_fields, num_fields // reduction)
        self.excitation = nn.Linear(num_fields // reduction, num_fields)

    def forward(self, field_embs):
        """field_embs: [B, num_fields, embed_dim]"""
        w = field_embs.mean(dim=-1)                    # [B, F]
        w = torch.sigmoid(self.excitation(torch.relu(self.squeeze(w))))
        return field_embs * w.unsqueeze(-1)            # 加权

class LightCrossNetwork(nn.Module):
    """轻量交叉网络：向量级交叉 O(d)，对比精排的矩阵交叉 O(d²)"""
    def __init__(self, input_dim, num_layers=2):
        super().__init__()
        self.weights = nn.ParameterList([nn.Parameter(torch.randn(input_dim)*0.01) for _ in range(num_layers)])
        self.biases = nn.ParameterList([nn.Parameter(torch.zeros(input_dim)) for _ in range(num_layers)])

    def forward(self, x0):
        xl = x0
        for w, b in zip(self.weights, self.biases):  # zip按位置配对
            xl = x0 * (xl * w) + b + xl
        return xl

class COLDModel(nn.Module):
    """COLD粗排：SE特征选择 + 轻量交叉 + MLP"""
    def __init__(self, feature_dims, embed_dim=16, se_reduction=4,
                 cross_layers=2, mlp_dims=None):
        super().__init__()
        if mlp_dims is None: mlp_dims = [128, 64]
        self.embeddings = nn.ModuleDict({
            n: nn.Embedding(d, embed_dim) for n, d in feature_dims.items()
        })
        num_fields = len(feature_dims)
        flat_dim = num_fields * embed_dim
        self.se = SEBlock(num_fields, se_reduction)
        self.cross = LightCrossNetwork(flat_dim, cross_layers)
        layers = []
        prev = flat_dim
        for d in mlp_dims:
            layers += [nn.Linear(prev, d), nn.BatchNorm1d(d), nn.ReLU(), nn.Dropout(0.1)]
            prev = d
        layers.append(nn.Linear(prev, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, features):
        embs = torch.stack([self.embeddings[n](features[n]) for n in self.embeddings], dim=1)
        selected = self.se(embs)
        flat = selected.view(selected.size(0), -1)  # 重塑张量形状
        return self.mlp(self.cross(flat)).squeeze(-1)

import torch.quantization as quant

def quantize_preranking(model: nn.Module) -> nn.Module:
    """动态量化：不需要校准数据，适合 MLP 为主的粗排模型"""
    return torch.quantization.quantize_dynamic(
        model, {nn.Linear}, dtype=torch.qint8
    )

# 量化前后对比
# 推理速度: FP32 ~5ms → INT8 ~1.5ms (3.3x 加速)
# 模型大小: FP32 12MB → INT8 3MB (4x 压缩)
# AUC损失:  <0.1% (几乎无损)

class FeatureImportanceAnalyzer:
    """基于 SE Block 权重的特征重要性分析"""
    @staticmethod  # @staticmethod不需要实例即可调用
    def analyze(model: COLDModel, dataloader) -> dict:
        importances = {}
        model.eval()
        with torch.no_grad():  # 禁用梯度计算，节省内存
            for batch in dataloader:
                embs = torch.stack([
                    model.embeddings[n](batch[n]) for n in model.embeddings
                ], dim=1)
                w = embs.mean(dim=-1)
                w = torch.sigmoid(model.se.excitation(torch.relu(model.se.squeeze(w))))
                for i, name in enumerate(model.embeddings.keys()):  # enumerate同时获取索引和元素
                    importances.setdefault(name, 0.0)
                    importances[name] += w[:, i].mean().item()
        total = sum(importances.values())
        return {k: v/total for k, v in importances.items()}

    # 裁剪策略：保留 top-70% 重要性的特征字段
    # 通常能裁掉 30% 特征，延迟减少 20-30%，AUC损失 <0.5%

import math

class GeneralSearchUnit(nn.Module):
    """GSU：从超长序列检索 Top-K 相关行为"""
    def __init__(self, embed_dim, mode="soft"):
        super().__init__()
        self.mode = mode
        if mode == "soft":
            self.q_proj = nn.Linear(embed_dim, embed_dim // 2)
            self.k_proj = nn.Linear(embed_dim, embed_dim // 2)

    def forward(self, behavior_embs, target_emb, top_k,
                behavior_cates=None, target_cate=None):
        """behavior_embs:[B,L,D], target_emb:[B,D]"""
        B, L, D = behavior_embs.shape
        if self.mode == "hard":
            # 类目匹配
            mask = (behavior_cates == target_cate.unsqueeze(1))
            scores = mask.float() + torch.arange(L, device=mask.device).unsqueeze(0) * 1e-6
            scores[~mask] = -1e9
        else:
            # 向量检索
            q = self.q_proj(target_emb)         # [B, D//2]
            k = self.k_proj(behavior_embs)      # [B, L, D//2]
            scores = torch.bmm(k, q.unsqueeze(-1)).squeeze(-1) / math.sqrt(q.size(-1))
        K = min(top_k, L)
        topk_scores, topk_idx = scores.topk(K, dim=1)
        # gather要求索引与输入同维度：unsqueeze(-1)加嵌入维，expand复制D次，从3D embedding中按topk索引取出对应向量
        topk_embs = torch.gather(behavior_embs, 1, topk_idx.unsqueeze(-1).expand(-1,-1,D))
        return topk_embs, topk_scores

class ExactSearchUnit(nn.Module):
    """ESU：对子序列做 Multi-Head Target Attention"""
    def __init__(self, embed_dim, num_heads=4, max_time=365):
        super().__init__()
        self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.time_emb = nn.Embedding(max_time + 1, embed_dim)
        self.time_gate = nn.Sequential(nn.Linear(embed_dim, embed_dim), nn.Sigmoid())
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, sub_seq, target_emb, time_deltas=None):
        """sub_seq:[B,K,D], target_emb:[B,D]"""
        if time_deltas is not None:
            t = self.time_emb(time_deltas.clamp(0, self.time_emb.num_embeddings-1))
            sub_seq = sub_seq * self.time_gate(t)
        q = target_emb.unsqueeze(1)             # [B,1,D]
        out, _ = self.mha(q, sub_seq, sub_seq)  # [B,1,D]
        return self.norm(out.squeeze(1) + target_emb)

class SIM(nn.Module):
    """SIM: GSU 检索 + ESU 注意力"""
    def __init__(self, num_items, embed_dim=64, top_k=100, num_heads=4,
                 gsu_mode="soft", mlp_dims=None):
        super().__init__()
        if mlp_dims is None: mlp_dims = [256, 128, 64]
        self.top_k = top_k
        self.item_emb = nn.Embedding(num_items, embed_dim)
        self.gsu = GeneralSearchUnit(embed_dim, gsu_mode)
        self.esu = ExactSearchUnit(embed_dim, num_heads)
        layers = []
        prev = embed_dim * 3
        for d in mlp_dims:
            layers += [nn.Linear(prev, d), nn.ReLU(), nn.Dropout(0.1)]
            prev = d
        layers.append(nn.Linear(prev, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, behavior_ids, target_id, behavior_cates=None,
                target_cate=None, time_deltas=None):
        beh = self.item_emb(behavior_ids)     # [B,L,D]
        tgt = self.item_emb(target_id)        # [B,D]
        sub, _ = self.gsu(beh, tgt, self.top_k, behavior_cates, target_cate)
        interest = self.esu(sub, tgt, time_deltas)
        feat = torch.cat([interest, tgt, interest * tgt], dim=-1)
        return self.mlp(feat).squeeze(-1)

class SimHashLayer(nn.Module):
    """SimHash：高维embedding → 低维二进制哈希码"""
    def __init__(self, input_dim, hash_bits=32):
        super().__init__()
        self.proj = nn.Linear(input_dim, hash_bits, bias=False)
        nn.init.orthogonal_(self.proj.weight)

    def forward(self, x):
        p = self.proj(x)
        if self.training:
            # STE: 前向 sign, 反向透传梯度
            code = (p > 0).float()
            return code - p.detach() + p
        return (p > 0).float()

    @staticmethod
    def hamming(c1, c2):
        """c1:[B,L,m], c2:[B,m] → 汉明距离 [B,L]"""
        return (c1 ^ c2.unsqueeze(1)).sum(dim=-1)

class ETAModel(nn.Module):
    """ETA: SimHash 硬检索 + Target Attention"""
    def __init__(self, num_items, embed_dim=64, hash_bits=32,
                 top_k=100, num_heads=4, mlp_dims=None):
        super().__init__()
        if mlp_dims is None: mlp_dims = [256, 128, 64]
        self.top_k, self.embed_dim = top_k, embed_dim
        self.item_emb = nn.Embedding(num_items, embed_dim)
        self.sim_hash = SimHashLayer(embed_dim, hash_bits)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.norm = nn.LayerNorm(embed_dim)
        layers = []
        prev = embed_dim * 3
        for d in mlp_dims:
            layers += [nn.Linear(prev, d), nn.ReLU(), nn.Dropout(0.1)]
            prev = d
        layers.append(nn.Linear(prev, 1))
        self.mlp = nn.Sequential(*layers)

    def forward(self, behavior_ids, target_id):
        beh = self.item_emb(behavior_ids)       # [B,L,D]
        tgt = self.item_emb(target_id)          # [B,D]
        # SimHash 检索
        bh = self.sim_hash(beh)                 # [B,L,m]
        th = self.sim_hash(tgt)                 # [B,m]
        dist = SimHashLayer.hamming(bh, th)     # [B,L]
        K = min(self.top_k, beh.size(1))
        _, idx = dist.topk(K, dim=1, largest=False)  # 距离最小=最相关
        topk = torch.gather(beh, 1, idx.unsqueeze(-1).expand(-1,-1,self.embed_dim))
        # Target Attention
        q = tgt.unsqueeze(1)
        out, _ = self.attn(q, topk, topk)
        interest = self.norm(out.squeeze(1) + tgt)
        feat = torch.cat([interest, tgt, interest * tgt], dim=-1)
        return self.mlp(feat).squeeze(-1)

class MultiRoundHashSampling(nn.Module):
    """SDIM 多轮哈希采样"""
    def __init__(self, embed_dim, num_rounds=4, num_buckets=64):
        super().__init__()
        self.rounds = num_rounds
        self.projs = nn.ModuleList([
            nn.Linear(embed_dim, num_buckets, bias=False) for _ in range(num_rounds)
        ])

    def forward(self, behavior_embs, target_emb):
        """behavior_embs:[B,L,D], target_emb:[B,D] → [B,D]"""
        B, L, D = behavior_embs.shape
        interests = []
        for proj in self.projs:
            beh_bucket = proj(behavior_embs).argmax(-1)   # [B,L]
            tgt_bucket = proj(target_emb).argmax(-1)      # [B]
            # unsqueeze(1)将目标桶号(B,)扩展为(B,1)与行为桶号(B,L)广播比较，再unsqueeze(-1)变(B,L,1)作为embedding的mask权重
            mask = (beh_bucket == tgt_bucket.unsqueeze(1)).float().unsqueeze(-1)
            bucket_sum = (behavior_embs * mask).sum(dim=1)
            interests.append(bucket_sum / mask.sum(dim=1).clamp(min=1))
        return torch.stack(interests).mean(dim=0)  # [B,D]

class LifecycleRouter(nn.Module):
    """根据用户生命周期特征路由到不同序列策略"""
    def __init__(self, embed_dim):
        super().__init__()
        self.router = nn.Sequential(
            nn.Linear(3, 16), nn.ReLU(), nn.Linear(16, 3), nn.Softmax(dim=-1)
        )  # 输入: [log(序列长度), 最近7天活跃, log(注册天数)]

    def forward(self, seq_len, recent_act, reg_days):
        feats = torch.stack([seq_len.float().log1p(), recent_act.float(),
                             reg_days.float().log1p()], dim=-1)
        return self.router(feats)  # [B, 3] 三种策略权重

class TimeDecayAttention(nn.Module):
    """时间衰减注意力：越久远的行为权重越低"""
    def __init__(self, embed_dim):
        super().__init__()
        self.attn_proj = nn.Linear(embed_dim * 3, 1)

    def forward(self, behavior_embs, target_emb, time_gaps):
        """
        behavior_embs: [B,L,D], target_emb: [B,D], time_gaps: [B,L] 距今天数
        """
        decay = torch.exp(-0.01 * time_gaps).unsqueeze(-1)  # 指数衰减 [B,L,1]
        tgt = target_emb.unsqueeze(1).expand_as(behavior_embs)
        attn_in = torch.cat([behavior_embs, tgt, behavior_embs * tgt], dim=-1)
        score = self.attn_proj(attn_in).squeeze(-1) * decay.squeeze(-1)
        weight = F.softmax(score, dim=-1).unsqueeze(-1)
        return (behavior_embs * weight).sum(dim=1)  # [B,D]