import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image

device = "cuda" if torch.cuda.is_available() else "cpu"

# 1. 데이터: 28x28 손글씨 숫자. 픽셀 값을 -1 ~ 1 범위로 맞춘다.
tf = transforms.Compose([transforms.ToTensor(), transforms.Normalize(0.5, 0.5)])
data = datasets.MNIST("data", train=True, download=True, transform=tf)
loader = DataLoader(data, batch_size=128, shuffle=True, drop_last=True)


# 2. 모델: 작은 U-Net
class Block(nn.Module):
    """합성곱 두 번. 사이에 시간 정보를 더한다."""

    def __init__(self, cin, cout):
        super().__init__()
        self.conv1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.conv2 = nn.Conv2d(cout, cout, 3, padding=1)
        self.time = nn.Linear(64, cout)

    def forward(self, x, temb):
        h = F.silu(self.conv1(x))
        h = h + self.time(temb)[:, :, None, None]
        return F.silu(self.conv2(h))


class TinyUNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.tmlp = nn.Sequential(nn.Linear(1, 64), nn.SiLU(), nn.Linear(64, 64))
        self.down1 = Block(1, 32)         # 28x28
        self.down2 = Block(32, 64)        # 14x14
        self.mid = Block(64, 128)         # 7x7
        self.up2 = Block(128 + 64, 64)    # 14x14
        self.up1 = Block(64 + 32, 32)     # 28x28
        self.out = nn.Conv2d(32, 1, 1)

    def forward(self, x, t):
        temb = self.tmlp(t[:, None])
        d1 = self.down1(x, temb)
        d2 = self.down2(F.avg_pool2d(d1, 2), temb)
        m = self.mid(F.avg_pool2d(d2, 2), temb)
        u2 = self.up2(torch.cat([F.interpolate(m, scale_factor=2), d2], dim=1), temb)
        u1 = self.up1(torch.cat([F.interpolate(u2, scale_factor=2), d1], dim=1), temb)
        return self.out(u1)


model = TinyUNet().to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
print("파라미터 수:", sum(p.numel() for p in model.parameters()))

# 3. 학습: 노이즈를 섞고, 섞은 노이즈를 맞히게 한다.
for epoch in range(3):
    total = 0.0
    for x0, _ in loader:
        x0 = x0.to(device)
        t = torch.rand(x0.shape[0], device=device)      # 그림마다 다른 노이즈 비율
        eps = torch.randn_like(x0)                      # 픽셀마다 따로 뽑은 노이즈
        tt = t[:, None, None, None]
        xt = (1 - tt) * x0 + tt * eps                   # 0-1 수업의 식
        loss = F.mse_loss(model(xt, t), eps)            # 예측한 노이즈와 실제 노이즈의 차이
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item()
    print(f"에폭 {epoch + 1}: 평균 손실 {total / len(loader):.4f}")

# 4. 확인: 노이즈가 반쯤 섞인 그림에서 원본을 되찾는다.
x0 = next(iter(loader))[0][:8].to(device)
t = torch.full((8,), 0.5, device=device)
tt = t[:, None, None, None]
eps = torch.randn_like(x0)
xt = (1 - tt) * x0 + tt * eps
with torch.no_grad():
    eps_hat = model(xt, t)
x0_hat = (xt - tt * eps_hat) / (1 - tt)                 # 식을 x0 에 대해 풀었다
print("복원 오차:", F.mse_loss(x0_hat, x0).item())
save_image(torch.cat([x0, xt, x0_hat]) * 0.5 + 0.5, "result.png", nrow=8)