# Source: content/notes/deep-learning/generative-models.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import math
import torch
from torch import nn
from torch.distributions import Normal, kl_divergence

torch.manual_seed(31)
torch.set_num_threads(1)
centers = torch.tensor([[-1., -1.], [-1., 1.], [1., -1.], [1., 1.]])
x = centers[torch.randint(4, (512,))] + 0.12 * torch.randn(512, 2)
encoder = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 4))
decoder = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))
opt = torch.optim.Adam(list(encoder.parameters()) + list(decoder.parameters()), lr=0.01)
sigma_x = 0.25
fixed_noise = torch.randn(len(x), 2)

def objective(noise):
    mu, logvar = encoder(x).chunk(2, dim=-1)
    logvar = logvar.clamp(-8, 4)
    std = (0.5 * logvar).exp()
    z = mu + std * noise
    reconstruction = decoder(z)
    nll = (0.5 * ((x - reconstruction) / sigma_x).square()
           + math.log(sigma_x * math.sqrt(2 * math.pi))).sum(-1)
    kl = 0.5 * (mu.square() + logvar.exp() - 1 - logvar).sum(-1)
    return (nll + kl).mean(), kl, mu, std

with torch.no_grad():
    initial = objective(fixed_noise)[0].item()
for _ in range(180):
    opt.zero_grad(set_to_none=True)
    loss, _, _, _ = objective(torch.randn(len(x), 2))
    loss.backward()
    opt.step()
encoder.eval()
decoder.eval()
with torch.no_grad():
    final, kl, mu, std = objective(fixed_noise)
    library_kl = kl_divergence(Normal(mu, std), Normal(torch.zeros_like(mu),
                                                       torch.ones_like(std))).sum(-1)
    torch.testing.assert_close(kl, library_kl, atol=1e-5, rtol=1e-5)
    prior_means = decoder(torch.randn(256, 2))
    samples = prior_means + sigma_x * torch.randn_like(prior_means)
    nearest = torch.cdist(samples, centers).argmin(-1)
    counts = torch.bincount(nearest, minlength=4)
assert final.item() < initial
assert samples.shape == (256, 2) and torch.isfinite(samples).all()
print({"initial_negative_elbo": initial, "final_negative_elbo": final.item(),
       "mean_kl": kl.mean().item(), "nearest_mode_counts": counts.tolist()})