{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Generative Adversarial Networks (GANs)\n",
    "\n",
    "In this notebook, we will introduce Generative Adversarial Networks (GANs) and implement a simple GAN using PyTorch. We will first discuss the basic concepts of GANs and training models on 2D data. Then, we will implement a GAN to generate handwritten digits from the MNIST dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.datasets as datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let us define the device to work on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"GPU available using Cuda: {torch.cuda.is_available()}\")\n",
    "print(\n",
    "    f\"GPU available using MPS: {torch.backends.mps.is_available() and torch.backends.mps.is_built()}\"\n",
    ")\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    gpu_available = True\n",
    "elif torch.backends.mps.is_available() and torch.backends.mps.is_built():\n",
    "    device = torch.device(\"mps\")\n",
    "    gpu_available = True\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    gpu_available = False\n",
    "\n",
    "\n",
    "print(f\"Device: {device}\")\n",
    "print(f\"Number of GPUs: {torch.cuda.device_count()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1 The GAN Model\n",
    "\n",
    "Generative Adversarials Networks (GANs) are a class of generative models that are used to generate new data samples. GANs consist of two neural networks, a generator and a discriminator, that are trained simultaneously. The generator generates new data samples, while the discriminator tries to distinguish between real and generated samples. The generator and discriminator are trained in an adversarial manner, where the generator tries to generate samples that are indistinguishable from real samples, and the discriminator tries to distinguish between real and generated samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 The Generator\n",
    "\n",
    "One easy way to thing about the generator is as a function that takes random noise as input and generates new data samples. The generator is a neural network that takes random noise as input and generates new data samples. The generator is trained to generate samples that are indistinguishable from real samples. \n",
    "\n",
    "Typically the generator is a feedforward neural network that takes random noise as input in $R^d$ and generates samples in $R^D$. The distribution of the output is the generated distribution. Typically, the input noise is sampled from a normal distribution and for instance, if the generator function is linear:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def input_noise(n=1000, noise_dim=1):\n",
    "    return torch.randn(n, noise_dim)\n",
    "\n",
    "\n",
    "def generator_function(x):\n",
    "    return 2 * x - 3.5\n",
    "\n",
    "\n",
    "input_data = input_noise()\n",
    "output_data = generator_function(input_data)\n",
    "plt.figure(figsize=(10, 5))\n",
    "plt.subplot(1, 2, 1)\n",
    "plt.hist(input_data.numpy(), bins=100, range=(-10, 10))\n",
    "plt.title(\"Input data\")\n",
    "plt.subplot(1, 2, 2)\n",
    "plt.hist(output_data.numpy(), bins=100, range=(-10, 10))\n",
    "plt.title(\"Output data\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generator can be implemented with a neural network to represent more complex functions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 The Discriminator\n",
    "\n",
    "Let us assume that we have a target distribution $p_{\\text{data}}(x)$ and a generator distribution $p_{\\text{model}}(x)$. The discriminator is trained to distinguish between samples from the target distribution and the generator distribution. The discriminator is trained to differentiate between samples from the target distribution and the generator distribution. In other words, the discriminator outputs the probability that a sample is from the target distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def target_density(x):\n",
    "    return 0.5 * torch.exp(\n",
    "        torch.distributions.Normal(-4, 1).log_prob(x)\n",
    "    ) + 0.5 * torch.exp(torch.distributions.Normal(4, 1).log_prob(x))\n",
    "\n",
    "\n",
    "def model_density(x):\n",
    "    return torch.exp(torch.distributions.Normal(-3.5, 2).log_prob(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = torch.linspace(-10, 10, 1000)\n",
    "plt.plot(x, target_density(x), label=\"Target\")\n",
    "plt.plot(x, model_density(x), label=\"Model\")\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimal discriminator is given by:\n",
    "$$D^*(x) = \\frac{p_{\\text{data}}(x)}{p_{\\text{data}}(x) + p_{\\text{model}}(x)}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Dopt = target_density(x) / (target_density(x) + model_density(x) + 1e-6)\n",
    "plt.figure()\n",
    "plt.plot(x, Dopt)\n",
    "plt.title(\"Optimal Density Ratio\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Generator is trained to generate samples that are indistinguishable from samples from the target distribution. In other words, the generator tries to generate samples with the highest probability of being from the target distribution.\n",
    "We will now implement a simple GAN on 2D data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. GAN on two dimensional data Gaussian Mixture\n",
    "\n",
    "In this section, we will implement a simple GAN on 2D data. We will generate samples from a Gaussian Mixture distribution and train a GAN to generate samples from the same distribution using PyTorch.\n",
    "We will implement a linear generator and discriminator for this task, then move on to more complex models. \n",
    "\n",
    "\n",
    "## 1.1 Training loop on simple data with linear models\n",
    "\n",
    "First let us visualize the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simple_2D_data(n=1000):\n",
    "    means = torch.Tensor([[-4, 0], [0, 4], [-2, 2]])\n",
    "    sigmas = torch.Tensor([0.6] * 3)\n",
    "    n_per_cluster = n // 3\n",
    "    data = []\n",
    "    for i in range(3):\n",
    "        data.append(\n",
    "            torch.distributions.Normal(means[i], sigmas[i]).sample((n_per_cluster,))\n",
    "        )\n",
    "    data = torch.cat(data) / 6\n",
    "    return data\n",
    "\n",
    "\n",
    "dataset = simple_2D_data(n=5000)\n",
    "plt.scatter(dataset[:, 0], dataset[:, 1], c=\"black\", s=1)\n",
    "plt.xlim(-1.1, 1.1)\n",
    "plt.ylim(-1.1, 1.1)\n",
    "plt.title(\"Target Data\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Complete the code below to implement a linear generator and discriminator:\n",
    "- The generator should be a linear layer that takes an input noise in dimension `noise_dim` and outputs a sample in dimension 2. You can add output activation if needed.\n",
    "- The discriminator should be a linear layer that takes an input sample in dimension 2 and outputs a probability that the sample is from the target distribution. The output should be a scalar between 0 and 1. You can add output activation if needed.\n",
    "\n",
    "```python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearGenerator(nn.Module):\n",
    "    def __init__(self, noise_dim=1):\n",
    "        super(LinearGenerator, self).__init__()\n",
    "        # Complete the code here\n",
    "        self.activation = nn.Tanh()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Complete the code here\n",
    "        return self.actication(x)\n",
    "\n",
    "class LinearGenerator(nn.Module):\n",
    "    def __init__(self, noise_dim=1):\n",
    "        super(LinearGenerator, self).__init__()\n",
    " \n",
    "    def forward(self, x):\n",
    "        # Complete the code here\n",
    "        return F.tanh(x)\n",
    "    \n",
    "\n",
    "class LinearGenerator(nn.Module):\n",
    "    def __init__(self, noise_dim=1):\n",
    "        super(LinearGenerator, self).__init__()\n",
    "        self.blocks = nn.Sequential(\n",
    "            nn.Linear(noise_dim, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 2),\n",
    "        )\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Complete the code here\n",
    "        return self.blocks(x)\n",
    "    \n",
    "\n",
    "class LinearDiscriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LinearDiscriminator, self).__init__()\n",
    "        # Complete the code here\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Complete the code here "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us the both models:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_generator(generator, noise_dim, n=1000):\n",
    "    noise = input_noise(n, noise_dim=noise_dim).to(device)\n",
    "    generated_data = generator(noise).detach().cpu()\n",
    "    plt.figure()\n",
    "    plt.scatter(generated_data[:, 0], generated_data[:, 1], c=\"blue\", s=1)\n",
    "    plt.xlim(-1.1, 1.1)\n",
    "    plt.ylim(-1.1, 1.1)\n",
    "    plt.title(\"Generated Data\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_discriminator(discriminator, n=1000):\n",
    "    x = torch.linspace(-1, 1, n)\n",
    "    y = torch.linspace(-1, 1, n)\n",
    "    xx, yy = torch.meshgrid(x, y)\n",
    "    grid = (\n",
    "        torch.from_numpy(np.stack([xx, yy], axis=-1).reshape(-1, 2)).float().to(device)\n",
    "    )\n",
    "    with torch.no_grad():\n",
    "        predictions = discriminator(grid).cpu().view(n, n)\n",
    "    plt.imshow(predictions.T, extent=(-1, 1, -1, 1), origin=\"lower\")\n",
    "    plt.colorbar()\n",
    "    plt.title(\"Discriminator Output\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_loss(loss_D, loss_G):\n",
    "    plt.figure()\n",
    "    plt.plot(loss_D, label=\"Discriminator Loss\")\n",
    "    plt.plot(loss_G, label=\"Generator Loss\")\n",
    "    plt.legend()\n",
    "    plt.title(\"Losses\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "generator = LinearGenerator(noise_dim).to(device)\n",
    "discriminator = LinearDiscriminator().to(device)\n",
    "\n",
    "visualize_generator(generator, noise_dim)\n",
    "visualize_discriminator(discriminator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the discriminator is trained to differentiate between samples from the target distribution and the generator distribution, we can train it as a binary classifier. Let us say that the point from the target distribution is labeled as 1 and the point from the generator distribution is labeled as 0.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complete the code below to implement the training loss for the discriminator using the binary cross entropy loss ```torch.nn.BCELoss``` [(Doc)](https://pytorch.org/docs/stable/generated/torch.nn.BCELoss.html). The discriminator should output the probability that the sample is from the target distribution. The target should be 1 for samples from the target distribution and 0 for samples from the generator distribution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCELoss()\n",
    "\n",
    "\n",
    "def discriminator_loss(output_real, output_fake):\n",
    "    real_labels = torch.ones_like(output_real)\n",
    "    fake_labels = torch.zeros_like(output_fake)\n",
    "\n",
    "    real_loss = # Complete here\n",
    "    fake_loss = # Complete here\n",
    "\n",
    "    return real_loss + fake_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The generator loss is the opposite of the discriminator loss. We can write it as the Binary Cross Entropy loss with the labels flipped.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complete the code below to implement the training loss for the generator. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator_loss(output_fake):\n",
    "    fake_labels = # Complete here\n",
    "    fake_loss = # Complete here\n",
    "    return fake_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font>  Run the code below to make sure that the training losses are implemented correctly.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "noise = input_noise(1000, noise_dim=noise_dim).to(device)\n",
    "data = simple_2D_data(n=5000).to(device)\n",
    "\n",
    "generator = LinearGenerator(noise_dim).to(device)\n",
    "discriminator = LinearDiscriminator().to(device)\n",
    "\n",
    "print(\n",
    "    \"The discriminator loss is: \",\n",
    "    discriminator_loss(discriminator(data), discriminator(generator(noise))),\n",
    ")\n",
    "print(\"The generator loss is: \", generator_loss(discriminator(generator(noise))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training loop for the GAN consists in alternating a gradient step on the discriminator and the generator:\n",
    "1. Sample a batch a data from the target distribution.\n",
    "2. Sample a batch of noise from a normal distribution with the same batch size.\n",
    "3. Generate samples from the generator.\n",
    "4. Compute the discriminator loss based on real and generated samples.\n",
    "5. Apply a gradient step on the discriminator.\n",
    "6. Compute the generator loss based on generated samples.\n",
    "7. Apply a gradient step on the generator.\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complete the code below to implement the training step for the GAN. The training step should apply a gradient step on the discriminator and the generator. The training step should return the discriminator and generator losses.  The training step should be able to apply a gradient step on the discriminator `num_steps` times before applying a gradient step on the generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def training_step(\n",
    "    generator, discriminator, optimizer_D, optimizer_G, noise, data, num_steps=1\n",
    "):\n",
    "    for _ in range(num_steps):\n",
    "        discriminator.train()\n",
    "        generator.eval()\n",
    "        # Complete here to train the discriminator\n",
    "\n",
    "    discriminator.eval()\n",
    "    generator.train()\n",
    "    # Complete here to train the generator\n",
    "\n",
    "    return loss_D.item(), loss_G.item()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to make sure that the training step is implemented correctly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "batch_size = 100\n",
    "\n",
    "generator = LinearGenerator(noise_dim).to(device)\n",
    "discriminator = LinearDiscriminator().to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), lr=0.002)\n",
    "optimizer_G = optim.Adam(generator.parameters(), lr=0.002)\n",
    "\n",
    "noise = input_noise(n=batch_size, noise_dim=noise_dim).to(device)\n",
    "data = simple_2D_data(n=batch_size).to(device)\n",
    "\n",
    "print(\"Before training\")\n",
    "print(\n",
    "    \"The discriminator loss is: \",\n",
    "    discriminator_loss(discriminator(data), discriminator(generator(noise))),\n",
    ")\n",
    "print(\"The generator loss is: \", generator_loss(discriminator(generator(noise))))\n",
    "\n",
    "for _ in range(100):\n",
    "    loss_D, loss_G = training_step(\n",
    "        generator, discriminator, optimizer_D, optimizer_G, noise, data\n",
    "    )\n",
    "\n",
    "print(\"\\nAfter a few training steps\")\n",
    "print(\n",
    "    \"The discriminator loss is: \",\n",
    "    discriminator_loss(discriminator(data), discriminator(generator(noise))),\n",
    ")\n",
    "print(\"The generator loss is: \", generator_loss(discriminator(generator(noise))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font>  Complete the code below to implement the training loop for the GAN. At each epoch, we should sample 5000 data points from the target distribution and split them into batches of size `batch_size`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_GAN(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    sample_function,\n",
    "    n_epochs=20,\n",
    "    batch_size=100,\n",
    "    num_steps=1,\n",
    "):\n",
    "    loss_D_values = []\n",
    "    loss_G_values = []\n",
    "    noise = # Complete here\n",
    "    data = # Complete here\n",
    "    n_batches = len(data) // batch_size\n",
    "    for epoch in range(n_epochs):\n",
    "        data = data[torch.randperm(len(data))]\n",
    "        noise = noise[torch.randperm(len(noise))]\n",
    "        avg_loss_D = 0\n",
    "        avg_loss_G = 0\n",
    "\n",
    "        for i in range(n_batches):\n",
    "            noise_i = # Complete here\n",
    "            data_i = # Complete here\n",
    "            loss_D, loss_G = # Complete here\n",
    "\n",
    "\n",
    "            loss_D_values.append(loss_D)\n",
    "            loss_G_values.append(loss_G)\n",
    "            avg_loss_D += loss_D\n",
    "            avg_loss_G += loss_G\n",
    "\n",
    "        print(\n",
    "            f\"Epoch {epoch+1}/{n_epochs} Loss D: {avg_loss_D/n_batches:.4f} Loss G: {avg_loss_G/n_batches:.4f}\"\n",
    "        )\n",
    "\n",
    "    return loss_D_values, loss_G_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to train the GAN on 2D data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = LinearGenerator(noise_dim).to(device)\n",
    "discriminator = LinearDiscriminator().to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), betas=(0.5, 0.9), lr=5e-4)\n",
    "optimizer_G = optim.Adam(generator.parameters(), betas=(0.5, 0.9), lr=5e-4)\n",
    "\n",
    "loss_D_values, loss_G_values = train_GAN(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    simple_2D_data,\n",
    "    n_epochs=20,\n",
    "    batch_size=100,\n",
    "    num_steps=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_loss(loss_D_values, loss_G_values)\n",
    "visualize_discriminator(discriminator)\n",
    "visualize_generator(generator, noise_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us try with a more complex model:\n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complete the code below to implement a more complex generator and discriminator:\n",
    "- Generator: Implement a neural network with 4 hidden layers of size 256, 512, 1024 with Leaky ReLU activation functions. The output layer should output a sample in dimension 2.\n",
    "- Discriminator: Implement a neural network with 7 hidden layers of size 1024, 512, 256 with Leaky ReLU activation functions. The output layer should output a scalar between 0 and 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DenseGenerator(nn.Module):\n",
    "    def __init__(self, noise_dim=1, out_dim=2):\n",
    "        super(DenseGenerator, self).__init__()\n",
    "        # Complete the code here\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Complete the code here\n",
    "\n",
    "\n",
    "class DenseDiscriminator(nn.Module):\n",
    "    def __init__(self, in_dim=2):\n",
    "        super(DenseDiscriminator, self).__init__()\n",
    "        # Complete the code here\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Complete the code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to train the GAN on 2D data with a more complex model. Observe the generated samples and the discriminator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "batch_size = 100\n",
    "\n",
    "generator = DenseGenerator(noise_dim).to(device)\n",
    "discriminator = DenseDiscriminator().to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)\n",
    "optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)\n",
    "\n",
    "loss_D_values, loss_G_values = train_GAN(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    simple_2D_data,\n",
    "    n_epochs=20,\n",
    "    batch_size=100,\n",
    "    num_steps=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_loss(loss_D_values, loss_G_values)\n",
    "visualize_discriminator(discriminator)\n",
    "visualize_generator(generator, noise_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now try with 2D data slightly more complex:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def complex_2D_data(n=1000):\n",
    "    means = torch.Tensor([[-2, 0], [2, 0], [0, 2], [0, -2]])\n",
    "    sigmas = torch.Tensor([0.5] * 4)\n",
    "    n_per_cluster = n // 4\n",
    "    data = []\n",
    "    for i in range(4):\n",
    "        data.append(\n",
    "            torch.distributions.Normal(means[i], sigmas[i]).sample((n_per_cluster,))\n",
    "        )\n",
    "    data = torch.cat(data) / 6\n",
    "    return data\n",
    "\n",
    "\n",
    "dataset = complex_2D_data(n=5000)\n",
    "plt.scatter(dataset[:, 0], dataset[:, 1], c=\"black\", s=1)\n",
    "plt.xlim(-1.1, 1.1)\n",
    "plt.ylim(-1.1, 1.1)\n",
    "plt.title(\"Target Data\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to train the GAN on 2D data with a more complex model. Observe the generated samples and the discriminator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "batch_size = 100\n",
    "\n",
    "generator = DenseGenerator(noise_dim).to(device)\n",
    "discriminator = DenseDiscriminator().to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)\n",
    "optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)\n",
    "\n",
    "loss_D_values, loss_G_values = train_GAN(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    complex_2D_data,\n",
    "    n_epochs=20,\n",
    "    batch_size=100,\n",
    "    num_steps=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_loss(loss_D_values, loss_G_values)\n",
    "visualize_discriminator(discriminator)\n",
    "visualize_generator(generator, noise_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GAN model is know for being hard to train for two main reasons:\n",
    "- Mode collapse: The generator collapses to a single point and generates the same sample.\n",
    "- Model instability: The generator and discriminator are trained simultaneously and the training can be unstable.\n",
    "\n",
    "We can observe the model collapse easilly with the 2D data on this dataset:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def complex_large_2D_data(n=1000):\n",
    "    means = torch.Tensor([[-4, 0], [4, 0], [0, 4], [0, -4]])\n",
    "    sigmas = torch.Tensor([0.6] * 4)\n",
    "    n_per_cluster = n // 4\n",
    "    data = []\n",
    "    for i in range(4):\n",
    "        data.append(\n",
    "            torch.distributions.Normal(means[i], sigmas[i]).sample((n_per_cluster,))\n",
    "        )\n",
    "    data = torch.cat(data) / 6\n",
    "    return data\n",
    "\n",
    "\n",
    "dataset = complex_large_2D_data(n=5000)\n",
    "plt.scatter(dataset[:, 0], dataset[:, 1], c=\"black\", s=1)\n",
    "plt.xlim(-1.1, 1.1)\n",
    "plt.ylim(-1.1, 1.1)\n",
    "plt.title(\"Target Data\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to train the GAN on 2D data with a more complex model. Conclude on mode collapse."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 10\n",
    "batch_size = 100\n",
    "\n",
    "generator = DenseGenerator(noise_dim).to(device)\n",
    "discriminator = DenseDiscriminator().to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)\n",
    "optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)\n",
    "\n",
    "loss_D_values, loss_G_values = train_GAN(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    complex_large_2D_data,\n",
    "    n_epochs=20,\n",
    "    batch_size=100,\n",
    "    num_steps=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_loss(loss_D_values, loss_G_values)\n",
    "visualize_discriminator(discriminator)\n",
    "visualize_generator(generator, noise_dim)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. GAN on MNIST\n",
    "\n",
    "We will now implement a GAN to generate handwritten digits from the MNIST dataset. The MNIST dataset consists of 28x28 grayscale images of handwritten digits from 0 to 9. We will use the same architectyre as before. For compuational reasons, we will use a rescaled version of the MNIST dataset with images of size 'size'x'size'."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "size = 14"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The code below loads the MNIST dataset and rescales the images to the desired size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Transform to resize, convert to tensor, and flatten to a vector of size (size)\n",
    "transform = transforms.Compose(\n",
    "    [\n",
    "        transforms.Resize((size, size)),  # Resize to sizexsize pixels\n",
    "        transforms.ToTensor(),  # Convert to a tensor with shape (1, size, size)\n",
    "        transforms.Lambda(\n",
    "            lambda x: x.view(size * size)\n",
    "        ),  # Flatten to (size*size,) and normalize\n",
    "        transforms.Lambda(\n",
    "            lambda x: 2 * x - 1\n",
    "        ),  # Normalize to [-1, 1] to avoid overflow\n",
    "    ]\n",
    ")\n",
    "\n",
    "\n",
    "# Custom dataset to filter only labels 0 and 1\n",
    "class FilteredMNIST(datasets.MNIST):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        # Filter to keep only images with labels 0 and 1\n",
    "        indices = [\n",
    "            i for i, label in enumerate(self.targets) if label in [0, 1, 2, 3, 4]\n",
    "        ]\n",
    "        self.data = self.data[indices]\n",
    "        self.targets = self.targets[indices]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        # Return only the image, not the label\n",
    "        image, _ = super().__getitem__(index)\n",
    "        return image\n",
    "\n",
    "\n",
    "# Create datasets and loaders\n",
    "dataset_train = FilteredMNIST(\n",
    "    root=\"./data\", train=True, download=True, transform=transform\n",
    ")\n",
    "dataset_test = FilteredMNIST(\n",
    "    root=\"./data\", train=False, download=True, transform=transform\n",
    ")\n",
    "\n",
    "loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=128, shuffle=True)\n",
    "loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=128, shuffle=True)\n",
    "\n",
    "first_batch = next(iter(loader_test))\n",
    "\n",
    "plt.figure(1, figsize=(10, 10))\n",
    "plt.suptitle(\"Test Set\", y=0.57)\n",
    "for i in range(10):\n",
    "    plt.subplot(1, 10, i + 1)\n",
    "    plt.imshow(first_batch[i].view(size, size).numpy(), cmap=\"gray\")\n",
    "    plt.axis(\"off\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The function below will plot the generated samples from the generator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_images(generator, noise_dim):\n",
    "    generator.eval()\n",
    "    noise = input_noise(10, noise_dim=noise_dim).to(device)\n",
    "    generated_images = generator(noise).detach().cpu()\n",
    "    generated_images = (generated_images + 1) / 2\n",
    "    plt.figure(figsize=(10, 1))\n",
    "    for i in range(10):\n",
    "        plt.subplot(1, 10, i + 1)\n",
    "        plt.imshow(generated_images[i].view(size, size).numpy(), cmap=\"gray\")\n",
    "        plt.axis(\"off\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As usual for image data, we will split the data into batches of size `batch_size`. The training loop must be adapted to train the GAN on image data. \n",
    "\n",
    "<font color='blue'>**TODO:**</font> Complete the code below to implement the training loop for the GAN on MNIST. At each epoch, we iterate over the batches of the dataset and train the GAN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_GAN_MNIST(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    n_epochs=20,\n",
    "    num_steps=1,\n",
    "):\n",
    "    loss_D_values = []\n",
    "    loss_G_values = []\n",
    "\n",
    "    for epoch in range(n_epochs):\n",
    "        avg_loss_D = 0\n",
    "        avg_loss_G = 0\n",
    "\n",
    "        # Complete here to train the GAN\n",
    "\n",
    "        print(\n",
    "            f\"Epoch {epoch+1}/{n_epochs} Loss D: {avg_loss_D/len(loader_train):.4f} Loss G: {avg_loss_G/len(loader_train):.4f}\"\n",
    "        )\n",
    "        if epoch % 5 == 0:\n",
    "            plot_images(generator, noise_dim)\n",
    "\n",
    "    return loss_D_values, loss_G_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<font color='blue'>**TODO:**</font> Run the code below to train the GAN on MNIST."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noise_dim = 100\n",
    "generator = DenseGenerator(noise_dim, out_dim=size * size).to(device)\n",
    "discriminator = DenseDiscriminator(in_dim=size * size).to(device)\n",
    "\n",
    "optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002)\n",
    "optimizer_G = optim.Adam(generator.parameters(), lr=0.0002)\n",
    "\n",
    "loss_D_values, loss_G_values = train_GAN_MNIST(\n",
    "    generator,\n",
    "    discriminator,\n",
    "    optimizer_D,\n",
    "    optimizer_G,\n",
    "    noise_dim,\n",
    "    n_epochs=100,\n",
    "    num_steps=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualize_loss(loss_D_values, loss_G_values)\n",
    "plot_images(generator, noise_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}