import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

class VAE(nn.Module):
    def __init__(self, latent_dim=2, input_dim=2):
        super(VAE, self).__init__()
        self.latent_dim = latent_dim

        # Encoder: From input space to latent space (mean and log variance)
        self.fc1 = nn.Linear(input_dim, 16)
        self.fc2_mu = nn.Linear(16, latent_dim)
        self.fc2_logvar = nn.Linear(16, latent_dim)

        # Decoder: From latent space to input space (mean of Gaussian)
        self.fc3 = nn.Linear(latent_dim, 16)
        self.fc4 = nn.Linear(16, input_dim)

    def encode(self, x):
        h = torch.relu(self.fc1(x))
        mu = self.fc2_mu(h)
        logvar = self.fc2_logvar(h)
        return mu, logvar

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        return mu + eps * std  # Latent variable z

    def decode(self, z):
        h = torch.relu(self.fc3(z))
        return self.fc4(h)  # Output is mean of Gaussian (the same dim as the input)

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        return self.decode(z), mu, logvar

# Loss function: Reconstruction loss + KL divergence
def vae_loss(recon_x, x, mu, logvar):
    recon_loss = nn.functional.mse_loss(recon_x, x, reduction='sum') # MSE loss for Gaussian decoder
    # KL Divergence (between approximate posterior and unit Gaussian prior)
    kl_div = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return recon_loss + 0.5 * kl_div

# Hyperparameters
latent_dim = 2
input_dim = 2
batch_size = 256
epochs = 1000
lr = 0.001

model = VAE(latent_dim, input_dim)
optimizer = optim.Adam(model.parameters(), lr=lr)

def create_data(n_samples):
    return torch.randn(n_samples, input_dim)  # 2D Gaussian data

# Training loop
for epoch in range(epochs):
    model.train()
    real_data = create_data(batch_size)
    
    optimizer.zero_grad()

    # Forward pass
    recon_data_mean, mu, logvar = model(real_data)

    # Compute loss
    loss = vae_loss(recon_data_mean, real_data, mu, logvar)

    # Backpropagation
    loss.backward()
    optimizer.step()

    if epoch % 100 == 0:
        print(f"Epoch {epoch} | Loss: {loss.item():.4f}")

print("Training complete!")

Epoch 0 | Loss: 516.2407
Epoch 100 | Loss: 472.4721
Epoch 200 | Loss: 421.2310
Epoch 300 | Loss: 334.9844
Epoch 400 | Loss: 332.3295
Epoch 500 | Loss: 298.6973
Epoch 600 | Loss: 309.5663
Epoch 700 | Loss: 321.7454
Epoch 800 | Loss: 321.4826
Epoch 900 | Loss: 307.2536
Training complete!

# Visualize the generated latent space and reconstructions
with torch.no_grad():
    model.eval()
    real_data = create_data(1000)
    recon_data_mean, _, _ = model(real_data)
    
    # Create a figure with two subplots side by side
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
    
    # Left plot: Real data vs mean reconstruction from sampled latent variables
    ax1.scatter(real_data[:, 0], real_data[:, 1], label="Real Data", alpha=0.6)
    # Sample from standard normal distribution in latent space
    z = torch.randn(1000, latent_dim)
    # Pass through decoder to get reconstructions
    sampled_recon = model.decode(z)
    ax1.scatter(sampled_recon[:, 0], sampled_recon[:, 1], alpha=0.6, color='orange', label='Reconstructed Mean (new samples from latent)')
    ax1.legend()
    ax1.set_title("Real vs. Reconstructed Mean (new samples from latent)")
    
    # Right plot: Real data vs generated data from reconstructed mean
    ax2.scatter(real_data[:, 0], real_data[:, 1], label="Real Data", alpha=0.6)
    # Generate new data by sampling around reconstructed mean
    generated = torch.normal(mean=recon_data_mean, std=torch.ones_like(recon_data_mean))
    ax2.scatter(generated[:, 0], generated[:, 1], alpha=0.6, color='green', label='Generated Data (sampled from reconstructed mean)')
    ax2.legend()
    ax2.set_title("Real vs. Generated Data")
    
    plt.tight_layout()
    plt.show()

--- VAE Architecture¶

--- Training Objective of VAE: From Maximizing Likelihood¶

Intractable $p_\theta(\mathbf{x}) = \int p_\theta(\mathbf{x} \mid \mathbf{z}) \, p(\mathbf{z}) \, d\mathbf{z}$¶

Lower Bound Surrogate of $p_\theta(\mathbf{x})$ By Variational Inference¶

VAE Loss Function: From ELBO to a Practical Implementable Loss¶

--- VAE Implementation¶

--- History of VAE¶

--- VAE Applications¶