# GENAI Experiment 1: Autoencoder + Variational Autoencoder (VAE)
# ===========================================================================

import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import matplotlib.pyplot as plt
torch.manual_seed(0); np.random.seed(0)

DATA = "shapes"          # "shapes" (no internet) | "digits" (sklearn, no internet) | "mnist" (needs internet)

def make_shapes(n=3000, size=16, seed=0):
    """Draw n binary images of circles (0), squares (1) and plus signs (2)."""
    rng = np.random.default_rng(seed)
    yy, xx = np.mgrid[0:size, 0:size]
    X = np.zeros((n, size, size), np.float32); y = rng.integers(0, 3, n)
    for i in range(n):
        r = rng.uniform(3, 5); cx, cy = rng.uniform(r, size - r, 2)
        if y[i] == 0:   m = (xx - cx)**2 + (yy - cy)**2 <= r*r
        elif y[i] == 1: m = (abs(xx - cx) <= 0.8*r) & (abs(yy - cy) <= 0.8*r)
        else:           m = ((abs(xx - cx) <= 0.3*r) & (abs(yy - cy) <= r)) | ((abs(yy - cy) <= 0.3*r) & (abs(xx - cx) <= r))
        X[i][m] = 1.0
    return X, y

if DATA == "shapes":
    X, y = make_shapes()
elif DATA == "digits":
    from sklearn.datasets import load_digits
    d = load_digits(); X = (d.images / 16.0).astype(np.float32); y = d.target
else:
    from torchvision import datasets
    m = datasets.MNIST("data", train=True, download=True)
    X = (m.data[:10000].numpy() / 255.0).astype(np.float32); y = m.targets[:10000].numpy()

IMG = X.shape[1]; D = IMG * IMG
Xt = torch.tensor(X).view(-1, D)                      # flatten: N x D
print("data:", DATA, "| images:", X.shape, "| pixel range", X.min(), X.max())

def show(imgs, n=10, title=""):
    imgs = np.asarray(imgs).reshape(-1, IMG, IMG)
    fig, ax = plt.subplots(1, n, figsize=(n * 1.1, 1.3))
    for i in range(n): ax[i].imshow(imgs[i], cmap="gray", vmin=0, vmax=1); ax[i].axis("off")
    plt.suptitle(title, fontsize=9); plt.show()

show(X[:10], title="Sample training images")

# # Part 1 — Simple Autoencoder + latent-space interpolation
# Step 1. Model. Encoder D → 64 → LATENT and decoder LATENT → 64 → D. The sigmoid output keeps pixels
# in 0–1.

LATENT = 8
class AE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, 64), nn.ReLU(), nn.Linear(64, LATENT))
        self.dec = nn.Sequential(nn.Linear(LATENT, 64), nn.ReLU(), nn.Linear(64, D), nn.Sigmoid())
    def forward(self, x):
        z = self.enc(x)
        return self.dec(z), z

ae = AE()
opt = torch.optim.Adam(ae.parameters(), lr=1e-3)

# Step 2. Train with mini-batches and BCE reconstruction loss. Record the loss per epoch for the loss
# curve.

losses = []
for epoch in range(20):
    perm = torch.randperm(len(Xt)); total = 0
    for i in range(0, len(Xt), 64):
        xb = Xt[perm[i:i+64]]
        recon, _ = ae(xb)
        loss = F.binary_cross_entropy(recon, xb)
        opt.zero_grad(); loss.backward(); opt.step()
        total += loss.item() * len(xb)
    losses.append(total / len(Xt))
    if epoch % 5 == 0 or epoch == 19: print(f"epoch {epoch:2d}  loss {losses[-1]:.4f}")

plt.plot(losses, marker="o"); plt.xlabel("epoch"); plt.ylabel("BCE loss"); plt.title("AE training loss"); plt.show()

# Step 3. Original vs reconstructed images, and latent vectors

with torch.no_grad():
    recon, z = ae(Xt[:10])
show(Xt[:10], title="Original"); show(recon, title="Reconstructed")
np.set_printoptions(suppress=True, precision=2)
print("Latent vectors of the first 3 images:\n", z[:3].numpy())

# Step 4. Latent interpolation z = (1−α)·z_A + α·z_B, decoded to generate in-between images.

with torch.no_grad():
    zA, zB = ae.enc(Xt[0:1]), ae.enc(Xt[1:2])
    inter = [ae.dec((1 - a) * zA + a * zB) for a in np.linspace(0, 1, 10)]
show(torch.cat(inter), title="Interpolation from image 0 to image 1")

# # Part 2 — Variational Autoencoder
# Step 1. Encoder outputs μ and log σ²; reparameterization trick; decoder

class VAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D, 128), nn.ReLU())
        self.mu = nn.Linear(128, LATENT)
        self.logvar = nn.Linear(128, LATENT)
        self.dec = nn.Sequential(nn.Linear(LATENT, 128), nn.ReLU(), nn.Linear(128, D), nn.Sigmoid())
    def forward(self, x):
        h = self.enc(x)
        mu, logvar = self.mu(h), self.logvar(h)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)   # reparameterization: z = mu + sigma * eps
        return self.dec(z), mu, logvar

def vae_loss(recon, x, mu, logvar):
    rec = F.binary_cross_entropy(recon, x, reduction="sum")
    kl = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())
    return rec, kl

vae = VAE()
opt = torch.optim.Adam(vae.parameters(), lr=1e-3)

# Step 2. Train with reconstruction + KL loss

hist = {"rec": [], "kl": []}
for epoch in range(30):
    perm = torch.randperm(len(Xt)); r_sum = k_sum = 0
    for i in range(0, len(Xt), 64):
        xb = Xt[perm[i:i+64]]
        recon, mu, logvar = vae(xb)
        rec, kl = vae_loss(recon, xb, mu, logvar)
        loss = (rec + kl) / len(xb)
        opt.zero_grad(); loss.backward(); opt.step()
        r_sum += rec.item(); k_sum += kl.item()
    hist["rec"].append(r_sum / len(Xt)); hist["kl"].append(k_sum / len(Xt))
    if epoch % 10 == 0 or epoch == 29:
        print(f"epoch {epoch:2d}  recon {hist['rec'][-1]:.2f}  KL {hist['kl'][-1]:.2f}")

plt.plot(hist["rec"], label="reconstruction"); plt.plot(hist["kl"], label="KL")
plt.xlabel("epoch"); plt.legend(); plt.title("VAE losses"); plt.show()

# Step 3. Sample latent vectors z ~ N(0, I) and generate NEW images (images that are not in the
# dataset)

with torch.no_grad():
    z = torch.randn(20, LATENT)
    new_imgs = vae.dec(z)
show(new_imgs[:10], title="New images generated by the VAE (1)")
show(new_imgs[10:], title="New images generated by the VAE (2)")

# Exam explanation (3 lines): the AE reconstructs well, but random z gives garbage because its latent
# space has gaps. The VAE's KL term pulls every code towards N(0, I), so the latent space is
# continuous and random samples decode to valid new images. The reconstruction loss falls during
# training while the KL term rises a little: the trade-off between accuracy and a smooth latent space.
