import math
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torch.nn import functional as F
import numpy as np
import matplotlib.pyplot as plt



class SinusoidalEmbedding(nn.Module):
    """DDPM 标准的正弦位置编码，把时间步 t 编码成向量"""
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb)
        emb = t[:, None] * emb[None, :]
        return torch.cat([emb.sin(), emb.cos()], dim=-1)


class ResidualBlock(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1, time_dim=None):
        super().__init__()
        # 第一个卷积：可能改变通道数和步长
        self.conv1 = nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_ch)

        # 时间嵌入投影：把 t 的嵌入映射到 out_ch，加到特征图上
        self.time_proj = None
        if time_dim is not None:
            self.time_proj = nn.Sequential(
                nn.SiLU(),
                nn.Linear(time_dim, out_ch),
            )

        # shortcut：如果维度变化，用1x1卷积（也可加BN）
        self.shortcut = nn.Sequential()
        if in_ch != out_ch :
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=1, bias=False),
                nn.BatchNorm2d(out_ch)
            )

    def forward(self, x, t_emb=None):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        # 添加时间嵌入（FiLM 式相加注入）
        if t_emb is not None and self.time_proj is not None:
            out = out + self.time_proj(t_emb)[:, :, None, None]

        # 添加 shortcut
        out += self.shortcut(residual)

        out = self.relu(out)   # 标准做法：最后加一次激活
        return out
class DownSample(nn.Module):
    def __init__(self,ch):
        super().__init__()
        self.conv1=nn.Conv2d(ch,ch,2,2)
    def forward(self,x):
        return self.conv1(x)

class UpSample(nn.Module):
    def __init__(self, ch):
        super().__init__()
        # 使用 ch//2 确保整数
        self.conv1 = nn.ConvTranspose2d(ch, ch//2, 4, 2, 1)
    def forward(self, x):
        return self.conv1(x)

class UNet(nn.Module):
    def __init__(self, time_dim=64):
        super().__init__()
        self.time_emb = SinusoidalEmbedding(time_dim)
        self.time_mlp = nn.Sequential(
            nn.Linear(time_dim, time_dim * 4),
            nn.SiLU(),
            nn.Linear(time_dim * 4, time_dim * 4),
        )
        t_dim = time_dim * 4

        self.conv1 = nn.Conv2d(1, 32, 3, 1, 1)
        self.res1 = ResidualBlock(32, 64, time_dim=t_dim) 
        self.down1 = DownSample(64)
        self.res2 = ResidualBlock(64, 128, time_dim=t_dim)
        self.down2 = DownSample(128)
        self.res3 = ResidualBlock(128, 256, time_dim=t_dim)  
        self.down3 = DownSample(256)
        self.res4 = ResidualBlock(256, 512, time_dim=t_dim)

        # --- 上采样部分（重新设计通道匹配）---
        self.up1 = UpSample(512)          # 512 -> 256
        self.res5 = ResidualBlock(512, 256, time_dim=t_dim)  # 拼接后 512 -> 256

        self.up2 = UpSample(256)          # 256 -> 128
        self.res6 = ResidualBlock(256, 128, time_dim=t_dim)  # 拼接后 256 -> 128

        self.up3 = UpSample(128)          # 128 -> 64
        self.res7 = ResidualBlock(128, 64, time_dim=t_dim)   # 拼接后 128 -> 64

        self.res8 = ResidualBlock(64, 32, time_dim=t_dim)
        self.conv2 = nn.Conv2d(32, 1, 3, 1, 1)

    def forward(self, x, t):
        t_emb = self.time_mlp(self.time_emb(t))  # [b, time_dim*4]
        # ----- 下采样（保存跳跃连接）-----
        x1 = self.res1(self.conv1(x), t_emb)          # [b, 64, 28, 28]
        x2 = self.res2(self.down1(x1), t_emb)         # [b, 128, 14, 14]
        x3 = self.res3(self.down2(x2), t_emb)         # [b, 256, 7, 7]
        x4 = self.res4(self.down3(x3), t_emb)         # [b, 512, 3, 3]

        # ----- 上采样 1（3×3 → 7×7）-----
        x4_up = self.up1(x4)                   # [b, 256, 6, 6]
        x4_up = F.interpolate(x4_up, size=7, mode='bilinear')  # [b, 256, 7, 7]
        x4_cat = torch.cat([x4_up, x3], dim=1) # [b, 512, 7, 7]
        x3_new = self.res5(x4_cat, t_emb)             # [b, 256, 7, 7]

        # ----- 上采样 2（7×7 → 14×14）-----
        x3_up = self.up2(x3_new)               # [b, 128, 14, 14]  尺寸恰好为14
        x3_cat = torch.cat([x3_up, x2], dim=1) # [b, 256, 14, 14]
        x2_new = self.res6(x3_cat, t_emb)             # [b, 128, 14, 14]

        # ----- 上采样 3（14×14 → 28×28）-----
        x2_up = self.up3(x2_new)               # [b, 64, 28, 28]
        x2_cat = torch.cat([x2_up, x1], dim=1) # [b, 128, 28, 28]
        x1_new = self.res7(x2_cat, t_emb)             # [b, 64, 28, 28]

        # ----- 最终输出-----
        x1_new = self.res8(x1_new, t_emb)             # [b, 32, 28, 28]
        out = self.conv2(x1_new)               # [b, 1, 28, 28]
        return out
if __name__=='__main__':
    # ---------- 数据加载 ----------
    batch_size = 16
    transform = transforms.Compose(
        [
            transforms.ToTensor(),  # 将 [0,255] 转为 [0,1] 的 Tensor
            transforms.Normalize((0.5,), (0.5,)),  #归一化到[-1,1]
        ]
    )

    train_dataset = datasets.MNIST(
        root="./data", train=True, download=True, transform=transform
    )
    test_dataset = datasets.MNIST(
        root="./data", train=False, download=True, transform=transform
    )

    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    # ---------- DDPM 超参数 ----------
    timesteps = 600
    betas = torch.linspace(1e-4, 0.02, timesteps)   # 线性采样 beta
    alphas = 1.0 - betas
    alpha_bar = torch.cumprod(alphas, dim=0)        # \bar{alpha}_t

    def q_sample(x_0, t):
        """前向扩散：q(x_t | x_0) = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * eps"""
        alpha_bar_t = alpha_bar[t][:, None, None, None]
        eps = torch.randn_like(x_0)
        x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1 - alpha_bar_t) * eps
        return x_t, eps

    def sample(model, n=16, device="cpu"):
        """简易采样：从纯噪声逐步去噪"""
        model.eval()
        x = torch.randn(n, 1, 28, 28, device=device)
        with torch.no_grad():
            for t in reversed(range(timesteps)):
                t_batch = torch.full((n,), t, device=device, dtype=torch.long)
                eps_pred = model(x, t_batch)
                alpha_t = alphas[t].to(device)
                alpha_bar_t = alpha_bar[t].to(device)
                x = (x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_t)
                if t > 0:
                    x += torch.sqrt(betas[t].to(device)) * torch.randn_like(x)
        model.train()
        return x

    # ---------- 训练 ----------
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = UNet().to(device)
    optimizer = optim.Adam(model.parameters(), lr=1e-3)
    epochs = 20

    for epoch in range(epochs):
        total_loss = 0.0
        for x_0, _ in train_loader:
            x_0 = x_0.to(device)
            t = torch.randint(0, timesteps, (x_0.size(0),), device=device)  # 随机采样时间
            x_t, eps = q_sample(x_0, t)

            eps_pred = model(x_t, t)
            loss = F.mse_loss(eps_pred, eps)  # 预测噪声，MSE 损失

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(train_loader):.4f}")

    # ---------- 采样可视化 ----------
    samples = sample(model, n=16, device=device).cpu().clamp(-1, 1)
    samples = (samples + 1) / 2  # [-1,1] -> [0,1]

    fig, axes = plt.subplots(4, 4, figsize=(8, 8))
    for i, ax in enumerate(axes.flat):
        ax.imshow(samples[i, 0], cmap="gray")
        ax.axis("off")
    plt.tight_layout()
    plt.savefig("samples.png")
    plt.show()

        