{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generative Models with PyTorch — **Student Version** (MNIST)\n",
    "Autoencoder → VAE (conv) → CVAE (conv)\n",
    "\n",
    "**Read this first**\n",
    "- Follow the **consignes** before each code cell marked `# TODO:`.\n",
    "- Only edit where a `# TODO:` appears. Everywhere else, just run.\n",
    "- When we say \"**use my architecture**\", we mean the conv backbone defined in the instructions here.\n",
    "\n",
    "**Latent sizes**\n",
    "- First tiny loop: latent_dim=16\n",
    "- LR/L1-vs-L2 experiments: latent_dim=32\n",
    "- Rest (AE/VAE/CVAE): latent_dim=64\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0) Setup\n",
    "\n",
    "Just run these cells. If you don't have `scikit-learn`, install it with:\n",
    "\n",
    "```\n",
    "pip install scikit-learn\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, math, random, numpy as np\n",
    "from pathlib import Path\n",
    "import torch, torch.nn as nn, torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision, torchvision.transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "SEED = 123\n",
    "random.seed(SEED)\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed_all(SEED)\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "elif torch.backends.mps.is_available():\n",
    "    device = torch.device(\"mps\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "print(\"Using device:\", device)\n",
    "\n",
    "\n",
    "def denorm(x):\n",
    "    return x  # images are 0..1 in this version\n",
    "\n",
    "\n",
    "def show_grid(imgs, nrow=8, title=None):\n",
    "    import torchvision\n",
    "\n",
    "    grid = torchvision.utils.make_grid(imgs, nrow=nrow)\n",
    "    plt.figure(figsize=(min(12, nrow + 4), 4))\n",
    "    plt.imshow(grid.permute(1, 2, 0).cpu(), cmap=\"gray\")\n",
    "    plt.axis(\"off\")\n",
    "    if title:\n",
    "        plt.title(title)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) DataLoader (MNIST)\n",
    "\n",
    "**Consignes**\n",
    "1. Run the next cell to load MNIST.\n",
    "2. Observe tensor shapes:\n",
    "   - One image: `[1, 28, 28]`\n",
    "   - Batch: `[B, 1, 28, 28]`\n",
    "3. Plot a small grid to visually confirm the data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = \"./data\"\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "transform = T.Compose(\n",
    "    [\n",
    "        T.ToTensor(),  # keep in [0,1]\n",
    "    ]\n",
    ")\n",
    "\n",
    "train_ds = torchvision.datasets.MNIST(\n",
    "    DATA_DIR, train=True, download=True, transform=transform\n",
    ")\n",
    "test_ds = torchvision.datasets.MNIST(\n",
    "    DATA_DIR, train=False, download=True, transform=transform\n",
    ")\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=2, pin_memory=True\n",
    ")\n",
    "test_loader = DataLoader(\n",
    "    test_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=2, pin_memory=True\n",
    ")\n",
    "\n",
    "xb, yb = next(iter(train_loader))\n",
    "print(\"Batch images:\", xb.shape)  # [B, 1, 28, 28]\n",
    "print(\"Batch labels:\", yb.shape)  # [B]\n",
    "show_grid(xb[:32], nrow=8, title=\"MNIST mini-batch (0..1)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) One consistent training loop\n",
    "\n",
    "We will reuse the same training utilities for **all models**.  \n",
    "You will **implement the core steps in the training loop**.\n",
    "\n",
    "**Consignes (implement inside `train_one_epoch_steps`)**\n",
    "- Put the model in train mode.\n",
    "- For each batch:\n",
    "  1. Move tensors to `device`.\n",
    "  2. Compute `loss, metrics = loss_fn(model, (x, y))`.\n",
    "  3. **Zero grads**, **backward**, **optimizer step**.\n",
    "  4. Log `'loss'` and any values from `metrics`.\n",
    "- Return a list of dicts (per-step metrics).\n",
    "\n",
    "`evaluate_epoch` is provided and should not be modified.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "def train_one_epoch_steps(model, loader, optimizer, loss_fn):\n",
    "    model.train()\n",
    "    step_metrics = []\n",
    "    for x, y in loader:\n",
    "        # TODO: move x,y to device; compute loss via loss_fn;\n",
    "        # zero gradients; backward; optimizer step; then log metrics into a dict\n",
    "        # Example keys: {'loss': loss.item(), 'recon': ..., 'kl': ...}\n",
    "        pass  # TODO\n",
    "        row = {\"loss\": float(loss.item())}\n",
    "        for k, v in metrics.items():\n",
    "            row[k] = float(v) if torch.is_tensor(v) else float(v)\n",
    "        step_metrics.append(row)\n",
    "    return step_metrics  # TODO\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def evaluate_epoch(model, loader, loss_fn):\n",
    "    model.eval()\n",
    "    meter = defaultdict(float)\n",
    "    n = 0\n",
    "    for x, y in loader:\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        loss, metrics = loss_fn(model, (x, y))\n",
    "        meter[\"loss\"] += float(loss.item())\n",
    "        for k, v in metrics.items():\n",
    "            meter[k] += float(v) if torch.is_tensor(v) else float(v)\n",
    "        n += 1\n",
    "    for k in meter:\n",
    "        meter[k] /= max(n, 1)\n",
    "    return dict(meter)\n",
    "\n",
    "\n",
    "def fit(model, train_loader, val_loader, optimizer, loss_fn, epochs=5, ckpt_path=None):\n",
    "    history = {\"train_steps\": [], \"val_epoch\": []}\n",
    "    best = math.inf\n",
    "    for ep in range(1, epochs + 1):\n",
    "        steps = train_one_epoch_steps(model, train_loader, optimizer, loss_fn)\n",
    "        val = evaluate_epoch(model, val_loader, loss_fn)\n",
    "        history[\"train_steps\"].append(steps)\n",
    "        history[\"val_epoch\"].append(val)\n",
    "        if ckpt_path and val[\"loss\"] < best:\n",
    "            best = val[\"loss\"]\n",
    "            torch.save(\n",
    "                {\"model\": model.state_dict(), \"epoch\": ep, \"val_loss\": best}, ckpt_path\n",
    "            )\n",
    "        print(f\"Epoch {ep:02d} | val:\", {k: round(v, 4) for k, v in val.items()})\n",
    "    return history\n",
    "\n",
    "\n",
    "def plot_trainsteps_vs_val(\n",
    "    history, key=\"loss\", title=\"Train (per step) vs Eval (per epoch)\"\n",
    "):\n",
    "    plt.figure(figsize=(7, 4))\n",
    "    x_step, y_step = [], []\n",
    "    t = 0\n",
    "    for ep, steps in enumerate(history[\"train_steps\"], start=1):\n",
    "        for s in steps:\n",
    "            x_step.append(t)\n",
    "            y_step.append(s.get(key, np.nan))\n",
    "            t += 1\n",
    "        x_step.append(t)\n",
    "        y_step.append(np.nan)\n",
    "        t += 1  # gap\n",
    "    plt.plot(x_step, y_step, label=f\"train/{key}\")\n",
    "    y_val = [v.get(key, np.nan) for v in history[\"val_epoch\"]]\n",
    "    x_val = np.linspace(0, t, num=len(y_val), endpoint=False) + (t / len(y_val)) / 2.0\n",
    "    plt.scatter(x_val, y_val, marker=\"o\", label=f\"val/{key}\")\n",
    "    plt.xlabel(\"training steps (gaps per epoch)\")\n",
    "    plt.ylabel(key)\n",
    "    plt.title(title)\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) Autoencoders\n",
    "\n",
    "### 3.1 Linear AE (learn `nn.Module` anatomy)\n",
    "\n",
    "**Architecture**  \n",
    "Encoder: Flatten → Linear(784→256) → ReLU → Linear(256→latent)  \n",
    "Decoder: Linear(latent→256) → ReLU → Linear(256→784) → Sigmoid\n",
    "\n",
    "**Consignes**\n",
    "- In `__init__`, **implement both encoder and decoder** using `nn.Sequential`.\n",
    "- Use **latent_dim=16**.\n",
    "- Loss is **MSE**; we provide helper code.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearAE(nn.Module):\n",
    "    def __init__(self, latent_dim=16):\n",
    "        super().__init__()\n",
    "        # TODO: Student — implement encoder and decoder with nn.Sequential as described above\n",
    "        # self.enc = ...\n",
    "        # self.dec = ...\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.enc(x)\n",
    "        x_hat = self.dec(z).view(-1, 1, 28, 28)\n",
    "        return x_hat, z\n",
    "\n",
    "\n",
    "lin_ae = LinearAE(latent_dim=16).to(device)\n",
    "mse = nn.MSELoss()\n",
    "\n",
    "\n",
    "def lin_ae_loss_fn(model, batch):\n",
    "    # TODO: Student — implement the loss function for AE with MSE reconstruction loss\n",
    "    loss = ...\n",
    "    return loss, {\"recon\": loss}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Train the Linear AE**  \n",
    "Just run the following cell **after** you implemented `LinearAE`. It uses the generic trainer you completed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = optim.Adam(lin_ae.parameters(), lr=1e-3)\n",
    "hist_lin = fit(\n",
    "    lin_ae,\n",
    "    train_loader,\n",
    "    test_loader,\n",
    "    opt,\n",
    "    lin_ae_loss_fn,\n",
    "    epochs=3,\n",
    "    ckpt_path=\"best_linearAE.pt\",\n",
    ")\n",
    "plot_trainsteps_vs_val(\n",
    "    hist_lin, key=\"recon\", title=\"LinearAE: train steps vs val (recon)\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize reconstructions\n",
    "@torch.no_grad()\n",
    "def show_recon(model, loader, n=16, title=\"Top: input, Bottom: reconstruction\"):\n",
    "    model.eval()\n",
    "    x, _ = next(iter(loader))\n",
    "    x = x.to(device)[:n]\n",
    "    x_hat, _ = model(x)\n",
    "    show_grid(torch.cat([x, x_hat], dim=0), nrow=n, title=title)\n",
    "\n",
    "\n",
    "show_recon(lin_ae, test_loader, n=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Conv Autoencoder (use the prescribed conv architecture)\n",
    "\n",
    "**Encoder**  \n",
    "`Conv(1→32, 4,2,1)` → ReLU → `Conv(32→64, 4,2,1)` → ReLU → **Flatten** → `Linear(64·7·7→latent)`  \n",
    "\n",
    "**Decoder**  \n",
    "`Linear(latent→64·7·7)` → reshape `[B,64,7,7]` → `ConvTranspose(64→32, 4,2,1)` → ReLU → `ConvTranspose(32→1, 4,2,1)` → Sigmoid\n",
    "\n",
    "**Consignes**\n",
    "- Implement encoder/decoder in `__init__` using `nn.Sequential` + `nn.Linear`.\n",
    "- Use **latent_dim=64** here and later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvEncoder(nn.Module):\n",
    "    def __init__(self, latent_dim=64):\n",
    "        super().__init__()\n",
    "        # TODO: Student — implement the conv encoder as specified\n",
    "        # self.features = ...\n",
    "        # self.fc = ...\n",
    "\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.features(x).view(x.size(0), -1)\n",
    "        return self.fc(h)\n",
    "\n",
    "\n",
    "class ConvDecoder(nn.Module):\n",
    "    def __init__(self, latent_dim=64):\n",
    "        super().__init__()\n",
    "        # TODO: Student — implement the decoder: fc then deconv stack ending with Sigmoid\n",
    "        # self.fc = ...\n",
    "        # self.deconv = ...\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def forward(self, z):\n",
    "        h = self.fc(z).view(-1, 64, 7, 7)\n",
    "        return self.deconv(h)\n",
    "\n",
    "\n",
    "class ConvAE(nn.Module):\n",
    "    def __init__(self, latent_dim=64):\n",
    "        super().__init__()\n",
    "        self.enc = ConvEncoder(latent_dim)\n",
    "        self.dec = ConvDecoder(latent_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        z = self.enc(x)\n",
    "        x_hat = self.dec(z)\n",
    "        return x_hat, z\n",
    "\n",
    "\n",
    "ae = ConvAE(latent_dim=64).to(device)\n",
    "\n",
    "\n",
    "def ae_loss_fn(model, batch):\n",
    "    x, y = batch\n",
    "    x_hat, z = model(x)\n",
    "    recon = mse(x_hat, x)\n",
    "    return recon, {\"recon\": recon}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Train the Conv AE**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = optim.Adam(ae.parameters(), lr=1e-3)\n",
    "hist_ae = fit(\n",
    "    ae, train_loader, test_loader, opt, ae_loss_fn, epochs=5, ckpt_path=\"best_convae.pt\"\n",
    ")\n",
    "plot_trainsteps_vs_val(hist_ae, key=\"recon\", title=\"ConvAE: train steps vs val (recon)\")\n",
    "show_recon(ae, test_loader, n=16, title=\"ConvAE reconstructions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.3 Experiments: LR / optimizer / L1 vs L2\n",
    "\n",
    "**Consignes**\n",
    "- Copy the Conv AE but with **latent_dim=32**.\n",
    "- Implement **L1 loss** by creating `nn.L1Loss()` and using it in the loss function.\n",
    "- Try at least two learning rates (e.g., `1e-3`, `3e-4`) and compare curves.\n",
    "- Optional: try `SGD` vs `Adam/AdamW`.\n",
    "\n",
    "Summarize your observations (2–3 bullet points) in a short Markdown cell after training.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ae32 = ConvAE(latent_dim=32).to(device)\n",
    "\n",
    "# TODO: Student — make an L1 loss instance and use it inside the loss function below\n",
    "l1 = None  # TODO\n",
    "\n",
    "\n",
    "def ae32_loss_fn(model, batch):\n",
    "    x, y = batch\n",
    "    x_hat, z = model(x)\n",
    "    # TODO: Student — use L1(x_hat, x) here\n",
    "    recon = None  # TODO\n",
    "    return recon, {\"recon\": recon}\n",
    "\n",
    "\n",
    "# TODO: Student — try other LRs/optimizers\n",
    "opt32 = optim.Adam(ae32.parameters(), lr=1e-3)\n",
    "hist_ae32 = fit(ae32, train_loader, test_loader, opt32, ae32_loss_fn, epochs=3)\n",
    "plot_trainsteps_vs_val(hist_ae32, key=\"recon\", title=\"ConvAE(32): L1 vs L2, optimizers\")\n",
    "show_recon(ae32, test_loader, n=16, title=\"ConvAE(32) reconstructions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4) Latent size exploration\n",
    "\n",
    "**Consignes**\n",
    "- Train a **latent_dim=2** Conv AE and plot a scatter of the 2D codes for a test batch.\n",
    "- Color points by digit label (different color per label).\n",
    "- Train a **latent_dim=64** Conv AE and run t‑SNE on a subset of test latents.\n",
    "\n",
    "> Expect the AE 2D latent to be **not very separable** — that’s normal.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2D latent AE\n",
    "ae2d = ConvAE(latent_dim=2).to(device)\n",
    "opt2d = optim.Adam(ae2d.parameters(), lr=1e-3)\n",
    "_ = fit(ae2d, train_loader, test_loader, opt2d, ae_loss_fn, epochs=4)\n",
    "\n",
    "with torch.no_grad():\n",
    "    x, y = next(iter(test_loader))\n",
    "    Z = ae2d.enc(x.to(device)).cpu().numpy()\n",
    "    Y = y.numpy()\n",
    "\n",
    "plt.figure(figsize=(5, 4))\n",
    "plt.scatter(Z[:, 0], Z[:, 1], c=Y, cmap=\"tab10\", s=8)\n",
    "plt.colorbar()\n",
    "plt.title(\"ConvAE latent=2 (one batch)\")\n",
    "plt.xlabel(\"z1\")\n",
    "plt.ylabel(\"z2\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# t‑SNE on latent=64\n",
    "try:\n",
    "    from sklearn.manifold import TSNE\n",
    "\n",
    "    have_tsne = True\n",
    "except Exception as e:\n",
    "    print(\"Install scikit-learn for TSNE:\", e)\n",
    "    have_tsne = False\n",
    "\n",
    "ae64 = ConvAE(latent_dim=64).to(device)\n",
    "opt64 = optim.Adam(ae64.parameters(), lr=1e-3)\n",
    "_ = fit(ae64, train_loader, test_loader, opt64, ae_loss_fn, epochs=4)\n",
    "\n",
    "Zs, Ys = [], []\n",
    "with torch.no_grad():\n",
    "    for i, (x, y) in enumerate(test_loader):\n",
    "        if i >= 6:\n",
    "            break\n",
    "        Zs.append(ae64.enc(x.to(device)).cpu().numpy())\n",
    "        Ys.append(y.numpy())\n",
    "Z = np.concatenate(Zs)\n",
    "Y = np.concatenate(Ys)\n",
    "\n",
    "if have_tsne:\n",
    "    Z2 = TSNE(\n",
    "        n_components=2, init=\"random\", perplexity=30, learning_rate=\"auto\"\n",
    "    ).fit_transform(Z)\n",
    "    plt.figure(figsize=(5, 4))\n",
    "    plt.scatter(Z2[:, 0], Z2[:, 1], c=Y, cmap=\"tab10\", s=4)\n",
    "    plt.colorbar()\n",
    "    plt.title(\"t‑SNE of ConvAE latent=64 (subset)\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5) VAE (conv): Gaussian latent with MSE + KL\n",
    "\n",
    "**Objective (we minimize its value):**\n",
    "$$\n",
    "\\mathcal{L}(x) \\;=\\; \\mathrm{MSE}(x,\\hat x) \\;+\\; \\beta \\cdot \\mathrm{KL}\\big(q_\\phi(z|x)\\,\\|\\,\\mathcal{N}(0,I)\\big).\n",
    "$$\n",
    "\n",
    "**Consignes**\n",
    "- Implement in `__init__` the **same conv encoder/decoder** as Conv AE (decoder ends with **Sigmoid**).\n",
    "- Implement the **reparameterization trick**:\n",
    "  $$z = \\mu + \\sigma \\odot \\epsilon, \\quad \\epsilon \\sim \\mathcal{N}(0,I)$$\n",
    "- Implement the **KL term**:\n",
    "  $$\\mathrm{KL} = -\\tfrac{1}{2} \\; \\mathbb{E}\\left[1 + \\log\\sigma^2 - \\mu^2 - \\sigma^2\\right]$$\n",
    "- Use **MSE** for the reconstruction term; set `beta=1.0`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvVAE(nn.Module):\n",
    "    def __init__(self, latent_dim=64):\n",
    "        super().__init__()\n",
    "        # TODO: Student — implement encoder convs + fc, and decoder fc + deconvs (end with Sigmoid)\n",
    "        # self.enc_feat = ...\n",
    "        # self.enc_fc = ...\n",
    "        # self.mu = ...\n",
    "        # self.logvar = ...\n",
    "        # self.dec_fc = ...\n",
    "        # self.dec_deconv = ...\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def encode(self, x):\n",
    "        h = self.enc_feat(x).view(x.size(0), -1)\n",
    "        h = torch.relu(self.enc_fc(h))\n",
    "        return self.mu(h), self.logvar(h)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        # TODO: Student — implement z = mu + sigma * eps, with eps ~ N(0, I)\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def decode(self, z):\n",
    "        h = self.dec_fc(z).view(-1, 64, 7, 7)\n",
    "        return self.dec_deconv(h)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu, logvar = self.encode(x)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        x_hat = self.decode(z)\n",
    "        return x_hat, mu, logvar, z\n",
    "\n",
    "\n",
    "def vae_loss_fn(beta=1.0):\n",
    "    def _fn(model, batch):\n",
    "        x, y = batch\n",
    "        x_hat, mu, logvar, z = model(x)\n",
    "        recon = mse(x_hat, x)\n",
    "        # TODO: Student — implement KL = -0.5 * mean(1 + logvar - mu^2 - exp(logvar))\n",
    "        kl = None  # TODO\n",
    "        loss = recon + beta * kl\n",
    "        return loss, {\"recon\": recon, \"kl\": kl}\n",
    "\n",
    "    return _fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Train VAE** and verify that reconstructions are still good. If not, try reducing the KL weight $\\beta$ (default 1.0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = ConvVAE(latent_dim=64).to(device)\n",
    "opt_vae = optim.Adam(vae.parameters(), lr=1e-3)\n",
    "hist_vae = fit(\n",
    "    vae,\n",
    "    train_loader,\n",
    "    test_loader,\n",
    "    opt_vae,\n",
    "    vae_loss_fn(beta=1.0),\n",
    "    epochs=5,\n",
    "    ckpt_path=\"best_convvae.pt\",\n",
    ")\n",
    "plot_trainsteps_vs_val(\n",
    "    hist_vae, key=\"loss\", title=\"ConvVAE: total loss (train steps vs val)\"\n",
    ")\n",
    "plot_trainsteps_vs_val(\n",
    "    hist_vae, key=\"recon\", title=\"ConvVAE: recon (train steps vs val)\"\n",
    ")\n",
    "plot_trainsteps_vs_val(hist_vae, key=\"kl\", title=\"ConvVAE: KL (train steps vs val)\")\n",
    "show_recon(vae, test_loader, n=16, title=\"ConvVAE reconstructions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Verify that the latent space is roughly Gaussian (e.g., plot histograms of latent dimensions). If not, try increasing $\\beta$ (e.g., 0.1 → 1.0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# t-SNE on VAE latent=64 (already trained)\n",
    "\n",
    "\n",
    "Zs, Ys = [], []\n",
    "with torch.no_grad():\n",
    "    for i, (x, y) in enumerate(test_loader):\n",
    "        if i >= 10:\n",
    "            break\n",
    "        mu, logvar = vae.encode(x.to(device))\n",
    "        Zs.append(vae.reparameterize(mu, logvar).cpu().numpy())\n",
    "        Ys.append(y.numpy())\n",
    "Z = np.concatenate(Zs)\n",
    "Y = np.concatenate(Ys)\n",
    "\n",
    "if have_tsne:\n",
    "    Z2 = TSNE(\n",
    "        n_components=2, init=\"random\", perplexity=30, learning_rate=\"auto\"\n",
    "    ).fit_transform(Z)\n",
    "    plt.figure(figsize=(5, 4))\n",
    "    plt.scatter(Z2[:, 0], Z2[:, 1], c=Y, cmap=\"tab10\", s=4)\n",
    "    plt.colorbar()\n",
    "    plt.axis(\"equal\")\n",
    "    plt.title(\"t‑SNE of ConvVAE latent=64 (subset)\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def interpolate_vae(model, xA, xB, steps=9):\n",
    "    model.eval()\n",
    "    muA, logvarA = model.encode(xA)\n",
    "    zA = model.reparameterize(muA, logvarA)\n",
    "    muB, logvarB = model.encode(xB)\n",
    "    zB = model.reparameterize(muB, logvarB)\n",
    "    alphas = torch.linspace(0, 1, steps, device=device).view(-1, 1)\n",
    "    Z = (1 - alphas) * zA + alphas * zB\n",
    "    return model.decode(Z)\n",
    "\n",
    "\n",
    "xb, _ = next(iter(test_loader))\n",
    "xA, xB = xb[:1].to(device), xb[1:2].to(device)\n",
    "interp_imgs = interpolate_vae(vae, xA, xB, steps=12)\n",
    "show_grid(torch.cat([xA, interp_imgs, xB], 0), nrow=12, title=\"VAE interpolation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Sampling & interpolation\n",
    "\n",
    "**Consignes**\n",
    "- Implement `sample_vae(model, n, latent_dim)`:\n",
    "  1. Sample `z ~ N(0, I)`\n",
    "  2. Decode to images in `[0,1]`\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def sample_vae(model, n=16, latent_dim=64):\n",
    "    model.eval()\n",
    "    # TODO: Student — sample z ~ N(0, I) and decode\n",
    "    raise NotImplementedError  # TODO\n",
    "\n",
    "\n",
    "samples = sample_vae(vae, n=16, latent_dim=64)\n",
    "show_grid(samples, nrow=4, title=\"VAE prior samples\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6) Conditional VAE (CVAE)\n",
    "\n",
    "We condition on the class label $y$ by concatenating a learned one‑hot embedding to $z$ before decoding:\n",
    "$$\n",
    "\\text{decoder input} = [z;\\; e(y)].\n",
    "$$\n",
    "\n",
    "**Consignes**\n",
    "- Reuse the same encoder as VAE (same convs and FCs).\n",
    "- Add a small linear layer to embed one‑hot labels to size `emb_dim` (e.g., 16).\n",
    "- In the decoder, **concatenate** `z` and `e(y)` and then decode.\n",
    "- Loss is the same: **MSE + KL**.\n",
    "- Implement a **class-conditional sampler** `sample_cvae(model, labels, latent_dim)`:\n",
    "  - sample `z ~ N(0, I)`\n",
    "  - decode with provided labels\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_hot(y, num_classes=10):\n",
    "    oh = torch.zeros(y.size(0), num_classes, device=y.device)\n",
    "    oh.scatter_(1, y.view(-1, 1), 1.0)\n",
    "    return oh\n",
    "\n",
    "\n",
    "class ConvCVAE(nn.Module):\n",
    "    def __init__(self, latent_dim=64, num_classes=10, emb_dim=16):\n",
    "        super().__init__()\n",
    "        # TODO: Student — implement emb layer + encoder convs/fc + mu/logvar + decoder(fc+deconvs with Sigmoid)\n",
    "        # self.num_classes = num_classes\n",
    "        # self.emb = ...\n",
    "        # self.enc_feat = ...\n",
    "        # self.enc_fc = ...\n",
    "        # self.mu = ...\n",
    "        # self.logvar = ...\n",
    "        # self.dec_fc = ...\n",
    "        # self.dec_deconv = ...\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def encode(self, x):\n",
    "        h = self.enc_feat(x).view(x.size(0), -1)\n",
    "        h = torch.relu(self.enc_fc(h))\n",
    "        return self.mu(h), self.logvar(h)\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return mu + std * eps\n",
    "\n",
    "    def decode(self, z, y):\n",
    "        # TODO: Student — build one-hot, project with self.emb, concat with z, then decode\n",
    "        raise NotImplementedError  # TODO\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        mu, logvar = self.encode(x)\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "        x_hat = self.decode(z, y)\n",
    "        return x_hat, mu, logvar, z\n",
    "\n",
    "\n",
    "def cvae_loss_fn(beta=1.0):\n",
    "    def _fn(model, batch):\n",
    "        x, y = batch\n",
    "        x_hat, mu, logvar, z = model(x, y)\n",
    "        recon = mse(x_hat, x)\n",
    "        # TODO: Student — same KL as VAE\n",
    "        kl = None  # TODO\n",
    "        loss = recon + beta * kl\n",
    "        return loss, {\"recon\": recon, \"kl\": kl}\n",
    "\n",
    "    return _fn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Train CVAE**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cvae = ConvCVAE(latent_dim=64).to(device)\n",
    "opt_cvae = optim.Adam(cvae.parameters(), lr=1e-3)\n",
    "hist_cvae = fit(\n",
    "    cvae,\n",
    "    train_loader,\n",
    "    test_loader,\n",
    "    opt_cvae,\n",
    "    cvae_loss_fn(beta=1.0),\n",
    "    epochs=4,\n",
    "    ckpt_path=\"best_convcvae.pt\",\n",
    ")\n",
    "plot_trainsteps_vs_val(hist_cvae, key=\"loss\", title=\"ConvCVAE: total loss\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Class-conditional sampling**\n",
    "\n",
    "**Consignes**\n",
    "- Implement `sample_cvae(model, labels, latent_dim)`:\n",
    "  1. Labels is a CPU LongTensor (length = n).\n",
    "  2. Move labels to device, sample `z ~ N(0,I)` of shape `[n, latent_dim]`.\n",
    "  3. Decode with `(z, labels)`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def sample_cvae(model, labels, latent_dim=64):\n",
    "    model.eval()\n",
    "    # TODO: Student — implement class-conditional sampling\n",
    "    raise NotImplementedError  # TODO\n",
    "\n",
    "\n",
    "labels = torch.arange(10)\n",
    "cvae_samples = sample_cvae(cvae, labels, latent_dim=64)\n",
    "show_grid(cvae_samples, nrow=10, title=\"CVAE class-conditional prior samples (0..9)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7) Wrap‑up\n",
    "\n",
    "- You implemented: training loop core, AE/VAE/CVAE modules, VAE KL + reparam, CVAE conditioning, and samplers.\n",
    "- You ran LR/L1-vs-L2 experiments and visualized latent spaces.\n",
    "- Reflect: what does the VAE's KL do to the latent structure vs plain AE?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.x"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}