# 预计算 noise schedule（在训练开始前做一次）
T = 1000
betas = torch.linspace(1e-4, 0.02, T)       # β_1, β_2, ..., β_T，线性增长
alphas = 1 - betas                            # α_t = 1 - β_t
alpha_bar = torch.cumprod(alphas, dim=0)      # ᾱ_t = α_1 * α_2 * ... * α_t

# 给定 x0 和时间步 t，一步得到 x_t（即上面的采样公式）
def q_sample(x0, t):
    # x0: [B, C, H, W]，t: [B]
    eps = torch.randn_like(x0)                           # ε ~ N(0,I)
    sqrt_ab    = alpha_bar[t][:, None, None, None].sqrt()  # √ᾱ_t
    sqrt_1_ab  = (1 - alpha_bar[t])[:, None, None, None].sqrt()  # √(1-ᾱ_t)
    x_t = sqrt_ab * x0 + sqrt_1_ab * eps                # 公式直接翻译
    return x_t, eps

def timestep_embedding(t, dim):
    # t: [B]，整数时间步
    # 输出: [B, dim]，每个时间步对应一个 dim 维向量
    half = dim // 2
    # 生成 half 个频率（从高频到低频）
    freqs = torch.exp(-math.log(10000) * torch.arange(half) / half)
    args = t[:, None].float() * freqs[None]   # [B, half]，时间步 × 频率
    # 用 cos 和 sin 各表示一半维度，拼在一起
    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)  # [B, dim]
    return emb

# 在 UNet 的每个 ResBlock 里，把时间 embedding 加到特征图上
class ResBlock(nn.Module):
    def __init__(self, channels, t_dim):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
        # 把 t_dim 维的时间向量投影到 channels 维，然后加到特征图
        self.time_proj = nn.Linear(t_dim, channels)

    def forward(self, x, t_emb):
        # x:     [B, C, H, W]  当前特征图
        # t_emb: [B, t_dim]    时间步 embedding
        scale = self.time_proj(t_emb)[:, :, None, None]  # [B, C, 1, 1]
        return self.conv(x) + scale   # 广播加到每个空间位置

T = 1000
betas = torch.linspace(1e-4, 0.02, T)   # β_t，线性 schedule
alphas = 1 - betas                        # α_t = 1 - β_t
alpha_bar = torch.cumprod(alphas, dim=0)  # ᾱ_t = ∏α_s

def train_step(x0, unet):
    # x0: [B, C, H, W]，像素值归一化到 [-1, 1]
    B = x0.shape[0]

    # 步骤1：t ~ Uniform{1,...,T}
    t = torch.randint(0, T, (B,))               # [B]，每个样本随机一个时间步

    # 步骤2：ε ~ N(0,I)，然后算 x_t（前向公式的直接实现）
    eps = torch.randn_like(x0)                  # [B, C, H, W]
    sqrt_ab   = alpha_bar[t][:, None, None, None].sqrt()      # √ᾱ_t，形状 [B,1,1,1]
    sqrt_1_ab = (1 - alpha_bar[t])[:, None, None, None].sqrt()
    x_t = sqrt_ab * x0 + sqrt_1_ab * eps        # [B, C, H, W]，加噪后的图像

    # 步骤3：网络预测噪声 ε_θ(x_t, t)
    eps_pred = unet(x_t, t)                     # [B, C, H, W]

    # 步骤4：MSE 损失，即 ||ε - ε_θ||²
    loss = F.mse_loss(eps_pred, eps)
    return loss

@torch.no_grad()
def ddpm_sample(unet, shape):
    # shape: (B, C, H, W)
    x = torch.randn(shape)   # x_T ~ N(0, I)，从纯噪声出发

    for t in reversed(range(T)):    # t = T-1, T-2, ..., 1, 0
        t_batch = torch.full((shape[0],), t)

        # 用网络预测噪声 ε̂ = ε_θ(x_t, t)
        eps_pred = unet(x, t_batch)         # [B, C, H, W]

        # 计算均值 μ_t（对应上面公式）
        a  = alphas[t]        # α_t
        ab = alpha_bar[t]     # ᾱ_t
        mu = (x - (1-a) / (1-ab).sqrt() * eps_pred) / a.sqrt()

        # x_{t-1} = μ_t + √β_t * z（最后一步 t=0 不加噪声）
        if t > 0:
            z = torch.randn_like(x)
            x = mu + betas[t].sqrt() * z
        else:
            x = mu    # t=0 时直接返回均值，不加额外噪声

    return x   # [B, C, H, W]，去噪后的图像

def forward_diffuse_text(x0, t, alpha_bar):
    # x0: [B, L, d]，归一化的 token embeddings
    # t:  [B]
    eps = torch.randn_like(x0)              # [B, L, d]，ε ~ N(0,I)
    ab  = alpha_bar[t][:, None, None]       # [B, 1, 1]，广播到 [B, L, d]
    # 公式直接翻译：x_t = √ᾱ_t * x0 + √(1-ᾱ_t) * ε
    x_t = ab.sqrt() * x0 + (1 - ab).sqrt() * eps
    return x_t, eps   # x_t: [B, L, d]，形状始终不变

class TextDiffusionTransformer(nn.Module):
    def __init__(self, d_model, nhead, num_layers, vocab_size, max_len):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)   # token id → embedding
        self.pos_embed = nn.Embedding(max_len, d_model)  # 位置 embedding
        self.time_embed = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )
        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, batch_first=True)
        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
        # 输出仍是 embedding 空间（d_model 维），不是 vocab logits
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x_t, t):
        # x_t: [B, L, d]  加噪后的 embedding（连续值，不是 token id）
        # t:   [B]         时间步
        B, L, d = x_t.shape

        # 时间步编码：标量 t → d 维向量，广播加到序列每个位置
        t_emb = timestep_embedding(t, d)        # [B, d]
        t_emb = self.time_embed(t_emb)          # [B, d]，经过 MLP 变换
        t_emb = t_emb[:, None, :]               # [B, 1, d] → 广播到 [B, L, d]

        # 位置编码
        pos = torch.arange(L, device=x_t.device)
        p_emb = self.pos_embed(pos)[None]       # [1, L, d]

        # 融合：加噪 embedding + 时间信息 + 位置信息
        h = x_t + t_emb + p_emb                # [B, L, d]

        # 双向 self-attention：每个位置都能看到整个序列
        h = self.transformer(h)                 # [B, L, d]

        # 输出：预测噪声 ε̂（也在 embedding 空间，和 x_t 同维度）
        return self.out_proj(h)                 # [B, L, d]

def train_step_text_continuous(token_ids, model, alpha_bar):
    # token_ids: [B, L]，整数 token id
    B, L = token_ids.shape

    # 步骤1：token id → embedding，并归一化（让量级和噪声匹配）
    x0 = model.embed(token_ids)              # [B, L, d]
    x0 = F.normalize(x0, dim=-1)            # 归一化到单位球面

    # 步骤2：随机时间步 + 加噪
    t = torch.randint(0, T, (B,))
    x_t, eps = forward_diffuse_text(x0, t, alpha_bar)   # [B, L, d]

    # 步骤3：预测噪声
    eps_pred = model(x_t, t)                # [B, L, d]

    # 步骤4：MSE 损失（和图像版完全相同的公式）
    loss = F.mse_loss(eps_pred, eps)
    return loss

@torch.no_grad()
def text_diffusion_sample(model, B, L, d):
    x = torch.randn(B, L, d)    # 从噪声出发，x_T ~ N(0, I)

    for t in reversed(range(T)):
        t_batch = torch.full((B,), t)
        eps_pred = model(x, t_batch)         # [B, L, d]

        # DDPM 去噪公式（与图像版完全相同）
        a  = alphas[t]
        ab = alpha_bar[t]
        mu = (x - (1-a)/(1-ab).sqrt() * eps_pred) / a.sqrt()
        if t > 0:
            x = mu + betas[t].sqrt() * torch.randn_like(x)
        else:
            x = mu
    # x: [B, L, d]，最终去噪后的连续 embedding，但还不是 token id
    ...

# all_embeddings: [V, d]，词表中所有 token 的 embedding 矩阵（即 Embedding 层的权重）
def decode_to_tokens(x, all_embeddings):
    # x:              [B, L, d]，去噪后的连续 embedding
    # all_embeddings: [V, d]

    # 分别归一化
    x_norm   = F.normalize(x, dim=-1)                  # [B, L, d]
    emb_norm = F.normalize(all_embeddings, dim=-1)     # [V, d]

    # 矩阵乘法：[B, L, d] × [d, V] = [B, L, V]
    # logits[b, i, v] = 第 b 个样本、第 i 个位置与第 v 个 token 的余弦相似度
    logits = x_norm @ emb_norm.T                       # [B, L, V]

    # 对每个位置，取相似度最高的 token id
    token_ids = logits.argmax(dim=-1)                  # [B, L]
    return token_ids

def mask_forward(token_ids, t, alpha_bar, mask_id):
    # token_ids: [B, L]，原始 token id（整数）
    # 返回：x_t: [B, L]，部分位置被替换为 mask_id
    ab = alpha_bar[t][:, None]                          # [B, 1]，广播到 [B, L]
    # 每个位置以概率 ᾱ_t 保留（1 = 保留，0 = mask）
    keep = torch.bernoulli(ab.expand_as(token_ids.float())).bool()  # [B, L]
    x_t = token_ids.clone()
    x_t[~keep] = mask_id                                # 未保留的位置 → [MASK]
    return x_t                                          # [B, L]，dtype=long

class MaskedDiffusionModel(nn.Module):
    def __init__(self, vocab_size, d_model, nhead, num_layers, max_len):
        super().__init__()
        self.embed     = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Embedding(max_len, d_model)
        self.time_mlp  = nn.Sequential(
            nn.Linear(d_model, d_model * 4), nn.SiLU(),
            nn.Linear(d_model * 4, d_model),
        )
        layer = nn.TransformerEncoderLayer(d_model, nhead, batch_first=True)
        self.transformer = nn.TransformerEncoder(layer, num_layers)
        # 关键区别：输出是词表大小的 logits，不是 embedding
        self.out_proj = nn.Linear(d_model, vocab_size)

    def forward(self, x_t, t):
        # x_t: [B, L]  含 MASK 的 token id 序列（整数）
        # t:   [B]
        B, L = x_t.shape

        h = self.embed(x_t)                              # [B, L, d]，id 转 embedding
        h = h + self.pos_embed(torch.arange(L))[None]   # [B, L, d]，加位置编码

        t_emb = self.time_mlp(timestep_embedding(t, d_model))  # [B, d]
        h = h + t_emb[:, None, :]                        # [B, L, d]，加时间信息

        h = self.transformer(h)                          # [B, L, d]，双向 attention
        logits = self.out_proj(h)                        # [B, L, V]，每位置的 token 分布
        return logits

def train_step_masked(token_ids, model, alpha_bar, mask_id):
    # token_ids: [B, L]
    B, L = token_ids.shape

    # 步骤1：随机时间步 + 按比例 mask
    t = torch.randint(0, T, (B,))
    x_t = mask_forward(token_ids, t, alpha_bar, mask_id)    # [B, L]

    # 步骤2：网络预测每个位置的 token 概率
    logits = model(x_t, t)                                  # [B, L, V]

    # 步骤3：只在被 mask 的位置计算损失（预测原始 token）
    mask = (x_t == mask_id)                                 # [B, L]，bool
    loss = F.cross_entropy(
        logits[mask],       # [N_masked, V]，被 mask 位置的预测
        token_ids[mask],    # [N_masked]，被 mask 位置的真实 token
    )
    return loss

@torch.no_grad()
def masked_diffusion_sample(model, B, L, mask_id, vocab_size):
    x = torch.full((B, L), mask_id)     # 全 MASK 序列，[B, L]

    for t in reversed(range(T)):
        t_batch = torch.full((B,), t)
        logits = model(x, t_batch)       # [B, L, V]

        # 对所有位置采样（但只对 MASK 位置生效）
        probs   = logits.softmax(dim=-1)                        # [B, L, V]
        sampled = torch.multinomial(probs.view(-1, vocab_size), 1).view(B, L)

        still_masked = (x == mask_id)   # [B, L]，当前还是 MASK 的位置

        if t > 0:
            # 这一步要揭开的比例：ᾱ_{t-1} / ᾱ_t（每步只揭开一点点）
            # 剩余不揭开的比例：1 - ᾱ_{t-1} / ᾱ_t
            remask_prob = 1 - alpha_bar[t-1] / alpha_bar[t]
            # 以概率 (1 - remask_prob) 揭开 MASK，以概率 remask_prob 继续保持 MASK
            reveal = torch.bernoulli(
                torch.full_like(sampled.float(), 1 - remask_prob)
            ).bool()
            x = torch.where(still_masked & reveal, sampled, x)
        else:
            # 最后一步：把所有剩余 MASK 全部填上
            x = torch.where(still_masked, sampled, x)

    return x   # [B, L]，完整的 token 序列

class ConditionalDiffusionLayer(nn.Module):
    def forward(self, x_t, condition, t_emb):
        # x_t:       [B, L, d]    噪声序列（query 来源）
        # condition: [B, Lc, d]   文本 encoder 输出（Lc 可以不等于 L）
        # t_emb:     [B, d]       时间步 embedding

        h = x_t + t_emb[:, None, :]      # [B, L, d]，注入时间信息

        # 1. Self-attention：序列内部的全局交互
        #    Q = K = V 均来自 h，输出形状仍为 [B, L, d]
        h = self.self_attn(h, h, h)       # [B, L, d]

        # 2. Cross-attention：从条件中提取信息
        #    Q 来自 h，K/V 来自 condition
        #    Q: [B, L, d_head]，K/V: [B, Lc, d_head]
        #    attn(Q, K^T): [B, L, Lc]（每个噪声位置对所有条件位置的注意力权重）
        #    乘以 V: [B, L, d_head] → 融入条件信息
        h = self.cross_attn(h, condition, condition)  # [B, L, d]

        h = self.ffn(h)                   # [B, L, d]
        return h