# GENAI Experiment 9: Diffusion model (DDPM): forward/reverse process, schedules, guided generation
# ===========================================================================

# # Part 1 — Basic DDPM on 2-D data
# Dataset: two moons generated with NumPy (option: sklearn.datasets.make_moons; both work offline).

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

USE_SKLEARN = False
def two_moons(n=4000, noise=0.05):
    if USE_SKLEARN:
        from sklearn.datasets import make_moons
        X = make_moons(n, noise=noise, random_state=0)[0]
    else:
        t = np.random.rand(n) * np.pi; half = n // 2
        X = np.r_[np.c_[np.cos(t[:half]), np.sin(t[:half])], np.c_[1 - np.cos(t[half:]), 0.5 - np.sin(t[half:])]]
        X += np.random.randn(n, 2) * noise
    return (X - X.mean(0)) / X.std(0)

data = torch.tensor(two_moons(), dtype=torch.float32)
plt.figure(figsize=(4, 3)); plt.scatter(data[:, 0], data[:, 1], s=1); plt.title("real data"); plt.show()

# Step 1. Noise schedule and the forward process

def schedule(T, kind="linear"):
    if kind == "linear":
        betas = torch.linspace(1e-4, min(0.999, 0.02 * 1000 / T), T)
    else:                                                     # cosine
        s = 0.008; t = torch.linspace(0, T, T + 1) / T
        f = torch.cos((t + s) / (1 + s) * math.pi / 2) ** 2
        betas = (1 - f[1:] / f[:-1]).clamp(max=0.999)
    alphas = 1 - betas
    return betas, alphas, torch.cumprod(alphas, 0)

T = 100
betas, alphas, abar = schedule(T)
fig, ax = plt.subplots(1, 5, figsize=(15, 2.8))
for a_, t in zip(ax, [0, 10, 30, 60, 99]):
    eps = torch.randn_like(data)
    xt = abar[t].sqrt() * data + (1 - abar[t]).sqrt() * eps   # forward process in one step
    a_.scatter(xt[:, 0], xt[:, 1], s=1); a_.set_title(f"t={t+1}, signal={abar[t].sqrt():.2f}"); a_.set_xlim(-3, 3); a_.set_ylim(-3, 3)
plt.show()

# Step 2. Noise-prediction network (input: x_t and t) and training

class Net(nn.Module):
    def __init__(self, cond_classes=0):
        super().__init__()
        self.cond = cond_classes
        extra = cond_classes + 1 if cond_classes else 0                 # one-hot label + "no label" slot
        self.net = nn.Sequential(nn.Linear(2 + 1 + extra, 128), nn.SiLU(),
                                 nn.Linear(128, 128), nn.SiLU(), nn.Linear(128, 2))
    def forward(self, x, t, T, y=None):
        inp = [x, (t.float() / T).unsqueeze(1)]                        # time as a number in [0, 1]
        if self.cond: inp.append(nn.functional.one_hot(y, self.cond + 1).float())
        return self.net(torch.cat(inp, 1))

def train(T=100, kind="linear", iters=3000, data=data, labels=None, n_cls=0, p_uncond=0.2):
    betas, alphas, abar = schedule(T, kind)
    model = Net(n_cls); opt = torch.optim.Adam(model.parameters(), lr=2e-3); losses = []
    for it in range(iters):
        i = torch.randint(0, len(data), (256,)); x0 = data[i]
        t = torch.randint(0, T, (256,)); eps = torch.randn_like(x0)
        xt = abar[t].sqrt().unsqueeze(1) * x0 + (1 - abar[t]).sqrt().unsqueeze(1) * eps
        y = None
        if n_cls:
            y = labels[i].clone(); y[torch.rand(256) < p_uncond] = n_cls     # sometimes drop the label (for guidance)
        loss = ((model(xt, t, T, y) - eps) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step(); losses.append(loss.item())
    return model, losses, (betas, alphas, abar)

model, losses, sch = train(T=100)
plt.plot(np.convolve(losses, np.ones(50) / 50, "valid")); plt.title("training loss"); plt.xlabel("iteration"); plt.show()
print("final loss:", np.mean(losses[-200:]).round(4))

# Step 3. Reverse process: generate new samples from pure noise

@torch.no_grad()
def sample(model, sch, T, n=2000, y=None, w=0.0, n_cls=0):
    betas, alphas, abar = sch
    x = torch.randn(n, 2)                                           # start from pure noise
    for t in reversed(range(T)):
        tt = torch.full((n,), t)
        if n_cls:                                                   # classifier-free guidance
            e_c = model(x, tt, T, y); e_u = model(x, tt, T, torch.full((n,), n_cls))
            eps = e_u + w * (e_c - e_u)
        else:
            eps = model(x, tt, T)
        x = (x - betas[t] / (1 - abar[t]).sqrt() * eps) / alphas[t].sqrt()
        if t > 0: x = x + betas[t].sqrt() * torch.randn_like(x)
    return x

gen = sample(model, sch, 100)
def closeness(g, real=data):                                      # mean distance to the nearest real point
    return torch.cdist(g, real).min(1).values.mean().item()
fig, ax = plt.subplots(1, 2, figsize=(9, 3.5), sharex=True, sharey=True)
ax[0].scatter(data[:, 0], data[:, 1], s=1); ax[0].set_title("real")
ax[1].scatter(gen[:, 0], gen[:, 1], s=1, c="r"); ax[1].set_title(f"generated (closeness {closeness(gen):.3f})"); plt.show()
print("closeness of pure noise for comparison:", round(closeness(torch.randn(2000, 2)), 3))

# # Part 2 — Controlled study: noise schedule (linear vs cosine) × number of steps T
# Data, network and iterations are kept fixed; only the schedule and T change.

results = []
fig, ax = plt.subplots(2, 3, figsize=(12, 7), sharex=True, sharey=True)
for r, kind in enumerate(["linear", "cosine"]):
    for c, T_ in enumerate([10, 50, 200]):
        m, l, s = train(T=T_, kind=kind, iters=2000)
        g = sample(m, s, T_)
        res = (kind, T_, np.mean(l[-200:]), s[2][-1].sqrt().item(), closeness(g)); results.append(res)
        ax[r, c].scatter(data[:, 0], data[:, 1], s=1, c="lightgray"); ax[r, c].scatter(g[:, 0], g[:, 1], s=1, c="r")
        ax[r, c].set_title(f"{kind}, T={T_}, closeness {res[4]:.3f}")
plt.tight_layout(); plt.show()
print(f"{'schedule':<8}{'T':>5}{'final loss':>12}{'signal at T':>13}{'closeness':>11}")
for k, T_, l, sig, cl in results: print(f"{k:<8}{T_:>5}{l:>12.4f}{sig:>13.4f}{cl:>11.3f}")

# # Part 3 — Guided generation (the Stable Diffusion idea, offline)
# Stable Diffusion conditions the denoiser on text and uses a guidance scale w: ε = ε_uncond +
# w·(ε_cond − ε_uncond). Here the "prompt" is a class label (0, 1, 2 = three clusters). The model is
# trained with the label sometimes dropped, so it learns both conditional and unconditional
# predictions.

centers = torch.tensor([[-2.0, 0.0], [2.0, 0.0], [0.0, 2.5]])
labels3 = torch.randint(0, 3, (3000,)); data3 = centers[labels3] + 0.35 * torch.randn(3000, 2)
model3, _, sch3 = train(T=50, kind="cosine", iters=3000, data=data3, labels=labels3, n_cls=3)

fig, ax = plt.subplots(1, 4, figsize=(15, 3.4), sharex=True, sharey=True)
for a_, w in zip(ax, [0.0, 1.0, 3.0, 7.0]):
    g = sample(model3, sch3, 50, n=1500, y=torch.full((1500,), 2), w=w, n_cls=3)     # "prompt" = class 2 (top cluster)
    acc = ((g - centers[2]).norm(dim=1) < 1.0).float().mean().item()
    a_.scatter(data3[:, 0], data3[:, 1], s=1, c="lightgray"); a_.scatter(g[:, 0], g[:, 1], s=1, c="r")
    a_.set_title(f"guidance w={w}: {acc*100:.0f}% in target cluster")
plt.tight_layout(); plt.show()

# Exam explanation:
# * Part 1: the forward process destroys the moons into Gaussian noise, the network learns to predict
# that noise, and the reverse loop turns fresh noise into new moon-shaped samples.
# * Part 2: very small T (10) gives scattered samples, while larger T gives sharper ones. The cosine
# schedule keeps more signal early and works better at small T. The final loss is not comparable
# across settings.
# * Part 3: w = 0 ignores the "prompt" (samples spread over all clusters), w ≈ 1–3 puts almost all
# samples in the requested cluster, and a very large w (7) over-shoots and distorts them. This is the
# guidance-scale trade-off seen in Stable Diffusion, where very high guidance gives over-saturated,
# distorted images.
