# GENAI Experiment 2: Basic GAN
# ===========================================================================

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 — Build the Generator and the Discriminator

Xg = Xt * 2 - 1                      # scale real images to [-1, 1] to match the generator's tanh
Z = 32                                 # noise size

G = nn.Sequential(nn.Linear(Z, 128), nn.LeakyReLU(0.2),
                  nn.Linear(128, 256), nn.LeakyReLU(0.2),
                  nn.Linear(256, D), nn.Tanh())
Dnet = nn.Sequential(nn.Linear(D, 256), nn.LeakyReLU(0.2),
                     nn.Linear(256, 128), nn.LeakyReLU(0.2),
                     nn.Linear(128, 1))          # logit
optG = torch.optim.Adam(G.parameters(), lr=2e-4, betas=(0.5, 0.999))
optD = torch.optim.Adam(Dnet.parameters(), lr=2e-4, betas=(0.5, 0.999))
bce = nn.BCEWithLogitsLoss()
print(G); print(Dnet)

# # Part 2 — Adversarial training
# Each batch: (1) update D on real (1) and fake (0), then (2) update G so that D says 1 for fakes.
# Snapshots of fixed noise are saved to compare epochs.

fixed_z = torch.randn(10, Z)
lossD_hist, lossG_hist, snaps = [], [], {}
EPOCHS = 60
for epoch in range(1, EPOCHS + 1):
    perm = torch.randperm(len(Xg)); dl = gl = 0; nb = 0
    for i in range(0, len(Xg), 64):
        real = Xg[perm[i:i+64]]; b = len(real)
        ones, zeros = torch.ones(b, 1), torch.zeros(b, 1)
        # (1) Discriminator
        fake = G(torch.randn(b, Z))
        lossD = bce(Dnet(real), ones) + bce(Dnet(fake.detach()), zeros)
        optD.zero_grad(); lossD.backward(); optD.step()
        # (2) Generator
        lossG = bce(Dnet(fake), ones)
        optG.zero_grad(); lossG.backward(); optG.step()
        dl += lossD.item(); gl += lossG.item(); nb += 1
    lossD_hist.append(dl / nb); lossG_hist.append(gl / nb)
    if epoch in (1, 10, 30, 60):
        with torch.no_grad(): snaps[epoch] = (G(fixed_z) + 1) / 2
    if epoch % 10 == 0 or epoch == 1:
        print(f"epoch {epoch:2d}  L_D {lossD_hist[-1]:.3f}  L_G {lossG_hist[-1]:.3f}")

# # Part 3 — Visualise samples at different epochs

for e, imgs in snaps.items():
    show(imgs, title=f"Generated images after epoch {e} (same noise vectors)")
show(Xt[:10], title="Real images for comparison")

# # Part 4 — Loss curves and simple analysis (convergence, stability, mode collapse)

plt.plot(lossD_hist, label="Discriminator loss"); plt.plot(lossG_hist, label="Generator loss")
plt.axhline(2*np.log(2), ls=":", c="C0"); plt.axhline(np.log(2), ls=":", c="C1")
plt.xlabel("epoch"); plt.ylabel("loss"); plt.legend(); plt.title("GAN losses (dotted = equilibrium)"); plt.show()

with torch.no_grad():
    s = (G(torch.randn(500, Z)) + 1) / 2
    d_real = torch.sigmoid(Dnet(Xg[:500])).mean().item(); d_fake = torch.sigmoid(Dnet(s * 2 - 1)).mean().item()
print(f"D(real) = {d_real:.3f}   D(fake) = {d_fake:.3f}   (closer to 0.5 = more balanced game)")

# Mode-collapse check without any pretrained model: give each generated image the label of its nearest real image
nearest = torch.cdist(s, Xt[:2000]).argmin(1)
counts = np.bincount(y[:2000][nearest.numpy()], minlength=len(np.unique(y)))
print("Generated samples per class (nearest real image):", counts.tolist(), " real data is roughly uniform")
diversity = torch.pdist(s).mean().item(); real_div = torch.pdist(Xt[:500]).mean().item()
print(f"Diversity (mean pairwise distance): generated {diversity:.2f} vs real {real_div:.2f}")

# Exam explanation: early in training D wins easily and the samples are noise. As G improves, both
# losses settle and oscillate around a level instead of going to zero, and the samples become clean
# shapes. Two simple mode-collapse checks: (1) label each generated image with its nearest real image
# and count the classes, which should be roughly uniform; (2) compare pixel diversity of generated and
# real samples. A basic MLP GAN often prefers the easiest shape (blobs or circles), which is *partial*
# mode collapse. More epochs, dropout in D, or a lower D learning rate reduce it. FID and IS (lab
# theory) compare features of real and generated images; here pixel diversity is used because no
# pretrained network is allowed.
