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"
torch.manual_seed(0)

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)


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)
        self.down2 = Block(32, 64)
        self.mid = Block(64, 128)
        self.up2 = Block(128 + 64, 64)
        self.up1 = Block(64 + 32, 32)
        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)

# 학습: 0-6 과 한 줄만 다르다. 맞힐 대상이 노이즈(eps)가 아니라 속도(eps - x0)다.
for epoch in range(10):
    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
        loss = F.mse_loss(model(xt, t), eps - x0)
        opt.zero_grad()
        loss.backward()
        opt.step()
        total += loss.item()
    print(f"에폭 {epoch + 1}: 평균 손실 {total / len(loader):.4f}")


# 생성: 순수 노이즈(t = 1)에서 출발해 속도를 따라 t = 0 까지 걸어온다.
@torch.no_grad()
def sample(n, steps):
    x = torch.randn(n, 1, 28, 28, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((n,), 1.0 - i * dt, device=device)
        x = x - dt * model(x, t)
    return x


for steps in (1, 4, 20, 50):
    torch.manual_seed(1)
    save_image(sample(16, steps) * 0.5 + 0.5, f"gen_{steps}.png", nrow=16)