"""Readable Week 07 VAE core: binary 28x28 images, two continuous latent coordinates.

Requires PyTorch. Run this file for a small shape/gradient smoke check (no download).
The separate figure generator trains the lecture's fixed-budget MNIST pilot.
"""
import torch
from torch import nn
from torch.nn import functional as F


class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder = nn.Sequential(nn.Linear(784, 128), nn.ReLU())
        self.mu_head = nn.Linear(128, 2)
        self.logvar_head = nn.Linear(128, 2)
        self.decoder = nn.Sequential(
            nn.Linear(2, 128), nn.ReLU(), nn.Linear(128, 784)
        )

    def encode(self, x):
        h = self.encoder(x)
        return self.mu_head(h), self.logvar_head(h)

    def forward(self, x, epsilon=None):
        mu, logvar = self.encode(x)
        if epsilon is None:
            epsilon = torch.randn_like(mu)
        u = mu + torch.exp(0.5 * logvar) * epsilon
        return self.decoder(u), mu, logvar


def loss_terms(pixel_logits, x, mu, logvar):
    """Return one reconstruction NLL and one prior KL per image, both in nats."""
    reconstruction = F.binary_cross_entropy_with_logits(
        pixel_logits, x, reduction='none'
    ).sum(dim=1)
    kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar).sum(dim=1)
    return reconstruction, kl


if __name__ == '__main__':
    torch.manual_seed(17)
    model = VAE()
    x = torch.randint(0, 2, (4, 784)).float()
    pixel_logits, mu, logvar = model(x)
    reconstruction, kl = loss_terms(pixel_logits, x, mu, logvar)
    loss = (reconstruction + kl).mean()  # beta = 1
    loss.backward()
    assert pixel_logits.shape == x.shape and mu.shape == logvar.shape == (4, 2)
    assert reconstruction.shape == kl.shape == (4,)
    assert all(p.grad is not None and torch.isfinite(p.grad).all()
               for p in model.parameters())
    print('Forward shapes and encoder/decoder gradients: PASS')
    print('Parameter count:', sum(p.numel() for p in model.parameters()))