# GENAI Experiment 4: Text-to-image with a conditional VAE (DALL·E Mini idea)
# ===========================================================================

# # Part 1 — Create the image–text dataset

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)
S = 16
SHAPES, COLORS, SIZES = ["circle", "square", "plus"], ["red", "green", "blue"], ["small", "big"]
WORDS = SHAPES + COLORS + SIZES                          # vocabulary of the "text encoder"
yy, xx = np.mgrid[0:S, 0:S]

def draw(shape, color, size, rng):
    r = rng.uniform(2.5, 3.5) if size == "small" else rng.uniform(4.5, 5.5)
    cx, cy = rng.uniform(r, S - r, 2)
    if shape == "circle":   m = (xx-cx)**2 + (yy-cy)**2 <= r*r
    elif shape == "square": 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))
    img = np.zeros((S, S, 3), np.float32); img[m, COLORS.index(color)] = 1.0
    return img

def make_dataset(n=6000, seed=0):
    rng = np.random.default_rng(seed); imgs, caps = [], []
    for _ in range(n):
        sh, co, si = rng.choice(SHAPES), rng.choice(COLORS), rng.choice(SIZES)
        imgs.append(draw(sh, co, si, rng)); caps.append(f"{si} {co} {sh}")
    return np.stack(imgs), caps

X, captions = make_dataset()
print("images:", X.shape, "| example captions:", captions[:4])
fig, ax = plt.subplots(1, 8, figsize=(10, 1.6))
for i in range(8): ax[i].imshow(X[i]); ax[i].set_title(captions[i], fontsize=7); ax[i].axis("off")
plt.show()

# # Part 2 — Text encoder: tokenize the prompt, then a multi-hot condition vector

def encode_text(prompt, rng=None):
    tokens = prompt.lower().replace(",", " ").split()           # tokenization
    v = torch.zeros(len(WORDS))
    for t in tokens:
        if t in WORDS: v[WORDS.index(t)] = 1.0                 # unknown words are ignored
    if rng is not None:                                        # attribute not mentioned -> choose it at random
        for group in (SHAPES, COLORS, SIZES):
            if not any(w in tokens for w in group):
                v[WORDS.index(group[rng.integers(len(group))])] = 1.0
    return v

print("'big red circle' ->", encode_text("big red circle").tolist(), "  (order:", WORDS, ")")
Xt = torch.tensor(X).view(len(X), -1)                          # N x 768
Ct = torch.stack([encode_text(c) for c in captions])           # N x 8
D, CD, LATENT = Xt.shape[1], Ct.shape[1], 3      # small z: it should only hold position / size details
BETA = 4                                           # stronger KL (beta-VAE) so z cannot "steal" colour/shape from the text

# # Part 3 — Conditional VAE (encoder and decoder both see the text vector)

class CVAE(nn.Module):
    def __init__(self):
        super().__init__()
        self.enc = nn.Sequential(nn.Linear(D + CD, 256), nn.ReLU())
        self.mu, self.logvar = nn.Linear(256, LATENT), nn.Linear(256, LATENT)
        self.dec = nn.Sequential(nn.Linear(LATENT + CD, 256), nn.ReLU(), nn.Linear(256, D), nn.Sigmoid())
    def forward(self, x, c):
        h = self.enc(torch.cat([x, c], 1))
        mu, logvar = self.mu(h), self.logvar(h)
        z = mu + torch.exp(0.5 * logvar) * torch.randn_like(mu)
        return self.dec(torch.cat([z, c], 1)), mu, logvar

model = CVAE(); opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(30):
    perm = torch.randperm(len(Xt)); tot = 0
    for i in range(0, len(Xt), 64):
        xb, cb = Xt[perm[i:i+64]], Ct[perm[i:i+64]]
        recon, mu, logvar = model(xb, cb)
        loss = (F.binary_cross_entropy(recon, xb, reduction="sum")
                - BETA * 0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())) / len(xb)
        opt.zero_grad(); loss.backward(); opt.step(); tot += loss.item() * len(xb)
    if epoch % 5 == 0 or epoch == 29: print(f"epoch {epoch:2d}  loss {tot/len(Xt):.2f}")

# # Part 4 — Generate images from prompts

def generate(prompt, n=6, seed=0, temperature=1.0):
    g = torch.Generator().manual_seed(seed)
    z = torch.randn(n, LATENT, generator=g) * temperature
    rng = np.random.default_rng(seed)
    c = torch.stack([encode_text(prompt, rng) for _ in range(n)])   # each image fills missing words randomly
    with torch.no_grad():
        return model.dec(torch.cat([z, c], 1)).view(n, S, S, 3).numpy()

def show_row(imgs, title):
    fig, ax = plt.subplots(1, len(imgs), figsize=(len(imgs) * 1.3, 1.5))
    for i, im in enumerate(imgs): ax[i].imshow(np.clip(im, 0, 1)); ax[i].axis("off")
    plt.suptitle(title, fontsize=9); plt.show()

def relevance(imgs, prompt):
    """Simple automatic check: does the dominant colour match the colour word in the prompt?"""
    col = [w for w in prompt.split() if w in COLORS]
    if not col: return None
    return float(np.mean([COLORS[im.reshape(-1, 3).sum(0).argmax()] == col[0] for im in imgs]))

for p in ["circle", "red circle", "big red circle", "small blue plus", "big green square"]:
    imgs = generate(p); r = relevance(imgs, p)
    show_row(imgs, f"prompt: '{p}'" + (f"   colour relevance = {r:.2f}" if r is not None else "   (colour not specified -> random)"))

# # Part 5 — Effect of the random seed and of the temperature (scale of z)

for s in [0, 1, 2]:
    show_row(generate("big blue circle", seed=s), f"seed {s}: different positions / details, same content")
for t in [0.3, 1.0, 2.5]:
    show_row(generate("small red square", temperature=t), f"temperature {t}")
a, b = generate("big blue circle", seed=5), generate("big blue circle", seed=5)
print("same prompt + same seed -> identical images:", np.allclose(a, b))

# Why BETA = 4 and a 3-D z? If z is large it memorises colour and shape itself and the decoder ignores
# the text. A small z with a stronger KL term forces the decoder to take *what* to draw from the text
# vector, and z keeps only *where / how big*.
#
# Exam explanation: the prompt is tokenized into words and mapped to a condition vector, the "text
# embedding". The CVAE decoder combines it with random z. Richer prompts fix more attributes: "circle"
# gives random colours, "red circle" fixes the colour, "big red circle" fixes the size too, so
# relevance and consistency go up. The seed changes z, which changes position and details, and the
# same seed reproduces the same image. A low temperature (z near 0) gives faint "average" images,
# while a higher temperature gives more varied positions and sizes but sometimes broken shapes. This
# mirrors DALL·E Mini, where text → tokens → condition → generator, here in miniature.
