import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# ==========================================
# 1. 扩散调度器：定义加噪过程
# ==========================================
class DiffusionScheduler:
    def __init__(self, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.timesteps = timesteps
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def add_noise(self, x_0, t):
        """给干净图像加噪"""
        sqrt_alpha_bar = self.alpha_bars[t].sqrt().view(-1,1,1,1)
        sqrt_one_minus_alpha_bar = (1 - self.alpha_bars[t]).sqrt().view(-1,1,1,1)
        epsilon = torch.randn_like(x_0)
        x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * epsilon
        return x_t, epsilon

# ==========================================
# 2. 简化的 U-Net 去噪网络
# ==========================================
class SimpleUNet(nn.Module):
    def __init__(self, in_channels=1, time_emb_dim=32):
        super().__init__()
        # 时间嵌入（告诉网络当前是第几步）
        self.time_mlp = nn.Sequential(
            nn.Linear(1, time_emb_dim),
            nn.ReLU(),
            nn.Linear(time_emb_dim, time_emb_dim)
        )

        # 编码器
        self.conv1 = nn.Conv2d(in_channels, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, 3, padding=1)

        # 时间调节
        self.time_proj1 = nn.Linear(time_emb_dim, 64)
        self.time_proj2 = nn.Linear(time_emb_dim, 128)

        # 解码器
        self.deconv1 = nn.ConvTranspose2d(128, 64, 3, padding=1)
        self.deconv2 = nn.ConvTranspose2d(64, in_channels, 3, padding=1)

    def forward(self, x, t):
        # 时间嵌入
        t_emb = self.time_mlp(t.float().unsqueeze(-1))

        # 编码
        h1 = F.relu(self.conv1(x) + self.time_proj1(t_emb).unsqueeze(-1).unsqueeze(-1))
        h2 = F.relu(self.conv2(h1) + self.time_proj2(t_emb).unsqueeze(-1).unsqueeze(-1))

        # 解码
        h3 = F.relu(self.deconv1(h2))
        out = self.deconv2(h3)

        return out  # 预测的噪声

# ==========================================
# 3. 训练扩散模型
# ==========================================
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleUNet().to(device)
scheduler = DiffusionScheduler(timesteps=1000)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for epoch in range(100):
    for x_0 in dataloader:  # x_0: 真实图像
        x_0 = x_0.to(device)

        # 随机采样时间步
        t = torch.randint(0, 1000, (x_0.shape[0],)).to(device)

        # 加噪
        x_t, true_noise = scheduler.add_noise(x_0, t)

        # 预测噪声
        pred_noise = model(x_t, t)

        # 噪声预测损失
        loss = F.mse_loss(pred_noise, true_noise)

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

    print(f"Epoch {epoch+1}: loss = {loss.item():.4f}")

print("训练完成！")

# ==========================================
# 4. 采样生成新图像
# ==========================================
@torch.no_grad()
def sample(model, scheduler, img_size=(1, 28, 28)):
    """从纯噪声逐步去噪生成图像"""
    model.eval()
    x = torch.randn((1, *img_size)).to(device)  # 纯噪声

    for t in reversed(range(scheduler.timesteps)):
        t_batch = torch.full((1,), t, device=device)
        pred_noise = model(x, t_batch)

        # DDIM 采样公式
        alpha = scheduler.alphas[t]
        alpha_bar = scheduler.alpha_bars[t]
        beta = scheduler.betas[t]

        if t > 0:
            noise = torch.randn_like(x)
        else:
            noise = 0

        x = (1 / alpha.sqrt()) * (x - (beta / (1 - alpha_bar).sqrt()) * pred_noise) + beta.sqrt() * noise

    return x.clamp(0, 1)

generated_img = sample(model, scheduler)