VAE

# VAE (Variational Autoencoder) PyTorch 구현
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

class VAE(nn.Module):
    """Variational Autoencoder"""

    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super(VAE, self).__init__()

        self.latent_dim = latent_dim

        # Encoder: 입력 -> 잠재 분포의 평균과 분산
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU()
        )

        # 잠재 분포 파라미터
        self.fc_mu = nn.Linear(hidden_dim // 2, latent_dim)      # 평균
        self.fc_logvar = nn.Linear(hidden_dim // 2, latent_dim)  # 로그 분산

        # Decoder: 잠재 벡터 -> 재구성된 입력
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Linear(hidden_dim // 2, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()  # 0-1 범위로 출력
        )

    def encode(self, x):
        """입력을 잠재 분포의 파라미터로 인코딩"""
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        """
        Reparameterization Trick
        z = mu + std * epsilon 형태로 샘플링하여 역전파 가능하게 함
        """
        std = torch.exp(0.5 * logvar)
        epsilon = torch.randn_like(std)  # 표준 정규분포에서 샘플링
        z = mu + std * epsilon
        return z

    def decode(self, z):
        """잠재 벡터를 원본 공간으로 디코딩"""
        return self.decoder(z)

    def forward(self, x):
        """순전파: 인코딩 -> 샘플링 -> 디코딩"""
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        reconstructed = self.decode(z)
        return reconstructed, mu, logvar

    def generate(self, num_samples, device='cpu'):
        """새로운 샘플 생성"""
        # 표준 정규분포에서 샘플링
        z = torch.randn(num_samples, self.latent_dim).to(device)
        samples = self.decode(z)
        return samples


def vae_loss(reconstructed, original, mu, logvar, beta=1.0):
    """
    VAE 손실 함수

    Args:
        reconstructed: 재구성된 출력
        original: 원본 입력
        mu: 잠재 분포의 평균
        logvar: 잠재 분포의 로그 분산
        beta: KL 발산 가중치 (Beta-VAE용)

    Returns:
        total_loss, recon_loss, kl_loss
    """
    # 1. 재구성 손실 (Binary Cross Entropy)
    recon_loss = F.binary_cross_entropy(reconstructed, original, reduction='sum')

    # 2. KL Divergence: 잠재 분포와 표준 정규분포 간의 거리
    # KL(q(z|x) || p(z)) = -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
    kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())

    total_loss = recon_loss + beta * kl_loss
    return total_loss, recon_loss, kl_loss


def train_vae(model, train_loader, optimizer, device, epoch):
    """VAE 학습 함수"""
    model.train()
    total_loss = 0
    total_recon = 0
    total_kl = 0

    for batch_idx, (data, _) in enumerate(train_loader):
        data = data.view(-1, 784).to(device)  # Flatten

        optimizer.zero_grad()
        reconstructed, mu, logvar = model(data)

        loss, recon_loss, kl_loss = vae_loss(reconstructed, data, mu, logvar)
        loss.backward()
        optimizer.step()

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

    avg_loss = total_loss / len(train_loader.dataset)
    print(f'Epoch {epoch}: Loss={avg_loss:.4f}, '
          f'Recon={total_recon/len(train_loader.dataset):.4f}, '
          f'KL={total_kl/len(train_loader.dataset):.4f}')

    return avg_loss


def visualize_latent_space(model, test_loader, device):
    """잠재 공간 시각화 (2D 잠재 공간인 경우)"""
    model.eval()
    z_list = []
    label_list = []

    with torch.no_grad():
        for data, labels in test_loader:
            data = data.view(-1, 784).to(device)
            mu, _ = model.encode(data)
            z_list.append(mu.cpu())
            label_list.append(labels)

    z = torch.cat(z_list, dim=0).numpy()
    labels = torch.cat(label_list, dim=0).numpy()

    if z.shape[1] == 2:
        plt.figure(figsize=(10, 8))
        scatter = plt.scatter(z[:, 0], z[:, 1], c=labels, cmap='tab10', alpha=0.6)
        plt.colorbar(scatter)
        plt.xlabel('Latent Dimension 1')
        plt.ylabel('Latent Dimension 2')
        plt.title('VAE Latent Space Visualization')
        plt.savefig('vae_latent_space.png')
        print("Latent space visualization saved to 'vae_latent_space.png'")


# 사용 예제
if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # MNIST 데이터셋 로드
    transform = transforms.ToTensor()
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)

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

    # VAE 모델 초기화
    vae = VAE(input_dim=784, hidden_dim=400, latent_dim=20).to(device)
    optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)

    # 학습
    print("Training VAE on MNIST...")
    for epoch in range(1, 11):
        train_vae(vae, train_loader, optimizer, device, epoch)

    # 새로운 샘플 생성
    print("\nGenerating new samples...")
    samples = vae.generate(16, device)
    print(f"Generated samples shape: {samples.shape}")