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


# ---------- 模型定义（修正） ----------
class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()  # 修正：去掉多余的 self
        self.encoder = nn.Sequential(
            nn.Linear(784, 256),
            nn.ReLU(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Linear(64, 20),  # 输出 mu 和 log_var（或 sigma）
        )
        self.decoder = nn.Sequential(
            nn.Linear(10, 64),
            nn.ReLU(),
            nn.Linear(64, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Sigmoid(),
        )

    def forward(self, x):
        hidden = self.encoder(x)
        mu, log_var = hidden.chunk(2, dim=1)  # 用 log_var 更稳定
        # 重参数化：sigma = exp(0.5 * log_var)
        std = torch.exp(0.5 * log_var)
        eps = torch.randn_like(std)
        z = mu + eps * std
        x_hat = self.decoder(z)

        # KL 散度（按 batch 和像素平均）
        KL = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp())
        KL = KL / (x.size(0) * 28 * 28)  # 与重构损失尺度一致

        return x_hat, KL


# ---------- 超参数设置 ----------
batch_size = 128
epochs = 20
learning_rate = 1e-3
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# ---------- 数据加载 ----------
transform = transforms.Compose(
    [
        transforms.ToTensor(),  # 将 [0,255] 转为 [0,1] 的 Tensor
        transforms.Lambda(lambda x: x.view(-1)),  # 展平为 784 维向量
    ]
)

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)

# ---------- 初始化模型、优化器 ----------
model = VAE().to(device)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 重构损失（二元交叉熵，因为像素在 [0,1]）
criterion = nn.BCELoss(reduction="sum")  # sum 后与 KL 求和，再除以像素数

# ---------- 训练循环 ----------
for epoch in range(1, epochs + 1):
    model.train()
    total_loss = 0
    total_recon = 0
    total_kl = 0

    for batch_idx, (data, _) in enumerate(train_loader):
        data = data.to(device)
        optimizer.zero_grad()

        x_hat, kl = model(data)
        recon_loss = criterion(x_hat, data)  # 按 batch 求和（每个像素的 BCE 之和）

        # 总损失 = 重构损失 + KL 散度（都已除以像素数，但 recon 未除，所以需统一）
        # 此处将 recon 也除以像素数，使两项量级匹配
        recon_loss = recon_loss / (data.size(0) * 28 * 28)
        loss = recon_loss + kl

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        total_recon += recon_loss.item()
        total_kl += kl.item()

    avg_loss = total_loss / len(train_loader)
    avg_recon = total_recon / len(train_loader)
    avg_kl = total_kl / len(train_loader)

    print(
        f"Epoch {epoch:2d} | Avg Loss: {avg_loss:.4f} | Recon: {avg_recon:.4f} | KL: {avg_kl:.4f}"
    )

    # 每 5 个 epoch 生成一些样本看看效果（可选）
    if epoch % 5 == 0:
        model.eval()
        with torch.no_grad():
            # 从标准正态分布采样 16 个 latent code
            sample_z = torch.randn(16, 10).to(device)
            generated = model.decoder(sample_z).cpu().numpy()
            # 显示
            fig, axes = plt.subplots(4, 4, figsize=(6, 6))
            for i, ax in enumerate(axes.flat):
                ax.imshow(generated[i].reshape(28, 28), cmap="gray")
                ax.axis("off")
            plt.suptitle(f"Epoch {epoch} Generated Samples")
            plt.show()
            plt.close()

# ---------- 测试集评估（可选） ----------
model.eval()
test_loss = 0
with torch.no_grad():
    for data, _ in test_loader:
        data = data.to(device)
        x_hat, kl = model(data)
        recon = criterion(x_hat, data) / (data.size(0) * 28 * 28)
        test_loss += (recon + kl).item()
print(f"Test Average Loss: {test_loss / len(test_loader):.4f}")