# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.15.2
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# + id="iRjddwOiWzr5" tags=["keep"]
# %matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

import math

# + [markdown] id="wjyOyZDn6rk2"
# ## Exercise 11.1: PCA on MNIST
#
# In the lectures the principal component analysis (PCA) was introduced as a
# method for dimensionality reduction and feature extraction, i.e., to condense
# data by mapping it to a lower dimensional space of the most important features.
#
# Let
# \begin{equation*}
#   \mathbf{X} = \begin{bmatrix} \mathbf{x}_1^\intercal \\ \vdots \\
#     \mathbf{x}_N^\intercal \end{bmatrix} \in \mathbb{R}^{N \times D}
# \end{equation*}
# be a matrix of $N$ data samples $\mathbf{x}_n \in \mathbb{R}^D$, which are
# centered around zero.
# We consider a PCA with $M < D$ components.
#
# To project the data points $\mathbf{x}_n$ to the $M$-dimensional space that is
# defined by the $M$ principal components of $\mathbf{X}$, the so-called principal
# subspace of $\mathbf{X}$, we can use the singular value decomposition of
# $\mathbf{X}$. Let $\mathbf{X} = \mathbf{U} \mathbf{\Sigma} \mathbf{V}^\intercal$ be the
# singular value decomposition of the data matrix $\mathbf{X}$ with the singular
# values sorted in descending order.
#
# Then the projection $\mathbf{z}_n$ of
# data point $\mathbf{x}_n$ to the principal subspace of $\mathbf{X}$ is given by
# \begin{equation}
#   \mathbf{z}_n^\intercal = \mathbf{x}_n^\intercal \begin{bmatrix} \mathbf{v}_1 & \cdots & \mathbf{v}_M \end{bmatrix},
# \end{equation}
# where $\mathbf{v}_i$ is the $i$th column of matrix $\mathbf{V}$. The vector
# $\mathbf{z}_n$ can be seen as an encoding of the data point
# $\mathbf{x}_n$ in a lower dimensional space that is constructed by the directions
# for which the data shows the largest variations.
#
# - **NOTE:** The
#   singular value decomposition of a matrix $\mathbf{X} \in \mathbb{R}^{N \times D}$
#   is defined as a factorization of the form $\mathbf{X} = \mathbf{U} \mathbf{\Sigma}
#   \mathbf{V}^\intercal$ where $\mathbf{U} \in \mathbb{R}^{N \times N}$ and
#   $\mathbf{V} \in \mathbb{R}^{D \times D}$ are orthogonal matrices and
#   $\mathbf{\Sigma} \in \mathbb{R}^{N \times D}$ is a rectangular diagonal matrix with
#   non-negative numbers on the diagonal. The diagonal entries of $\mathbf{\Sigma}$ are
#   the so-called singular values of $\mathbf{X}$. A common convention is to sort
#   the singular values in descending order, in which case the diagonal matrix $\mathbf{\Sigma}$ is uniquely determined by $\mathbf{X}$.
#
# In this exercise, we perform and analyse PCA of the MNIST image data set with $k = 2$ principal components.

# + id="3_ghoR-Bxh7c" tags=["keep"] colab={"base_uri": "https://localhost:8080/", "height": 471, "referenced_widgets": ["4ccd141c98f54d90a9cd11aa91d3d8db", "a34f4a7ed0aa4959865fb6c18986c8da", "f75356b111904220955975e86d84b41b", "db7c9423a3d1406a9b8966aa2832292e", "f5f181b5d01847b7b5d8b994029b0e1d", "31e0ba7ceae1456ab227c3f6be6639aa", "f3c4195724224d828230727761f49ac2", "c931646be58f4fa784c3b709594405f9", "0a6dc847c2564c92b1b0321f242e7996", "4bd7c3ead5684cdabbb051b0ef13b04b", "6ecf1e979b2c4ecbad9ef1fd59927cee", "ce58554fdafe43ed8bb84e1b25ede9ec", "8254ddf4d1f54f0b89c9d760ad226b2d", "905b3b6803af496faf4034426e6e4dac", "7975a6c7dd0f47d69f23b472d7787cc7", "d24f97fbd12245618b5e047c56ab7bce", "a8b6d601a7404561bd3557963bb824ee", "05253a71049c4ac28763a6e35d345c7f", "3b4bb56c17a546e7a5e777221bbe39b8", "b929180018304b2cac36ce8c82bc3da4", "68a352513337409d8e1960bd0ffb46dd", "86cedf4c83a845bb83b966d3c5e3a535", "3a407ae95dc14cc4835db525a6ab5c0f", "026533823a6b48d49a2f99e32d872168", "8afdf2837c5a48cdaafa95f7edb9b009", "b6cdfbf584cb4ffb87f7bf6ee310ced1", "e1d12678f9f04619b3916f07f88bddad", "1ec6de04ff464db2bb5b80be79160771", "96cf6751628f402db35adf66eecd5a5e", "8b9a19d8bff044d4b6bcad9bd68bae3a", "225e2d9f8d904b9b89f85c1637b15cba", "8400dcf30433453896870fccc519ebdd", "334cc9537c694d64ad7110c9f9bf2981", "cbaddae4996d4534950fb0b13073dbb8", "4efa3f42f07944eb8977c914a78f4fa7", "9227c8c3297648c5b45d02e9d541a6c0", "3ac44b9d564b47e0bb9c2ae747737632", "da65b97e1ec848f3893d46375fb82a3b", "883e3560a05a4fd18dcfa38608747615", "cba7e0a667d14b299c83004b58be05c0", "dc039fe63c9a404c9b168c6f9a5acd96", "b48b4ab936194ec6b51b0be1671a3fcc", "891a8cf764d549c7bbab1782f1a3412b", "cffd8ffd876948c6ac21974d9de826ee"]} executionInfo={"status": "ok", "timestamp": 1633497614385, "user_tz": -120, "elapsed": 11601, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="efb2306c-6bbe-4062-99b9-3957517a14bf"
# We use torch and torchvision for automatically downloading the MNIST dataset
import torch
import torchvision
import torchvision.transforms as transforms

# download MNIST datasets and flatten images to vectors of length 784
transform = transforms.Compose([transforms.ToTensor(),
                                transforms.Lambda(lambda x: x.view(-1))])
trainset = torchvision.datasets.MNIST(root='./data', train=True,
                                      download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False,
                                     download=True, transform=transform)

# extract a complete PyTorch dataset
def extract(dataset):
    datasize = len(dataset)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=datasize, shuffle=False)
    return next(iter(dataloader))

# extract all training and test images and labels into PyTorch tensors
train_images, train_labels = extract(trainset)
test_images, test_labels = extract(testset)

# It is possible to do the exercise both using torch or numpy.
# In the second case, you can just convert the matrices to numpy using:
# train_images = train_images.numpy()
# train_labels = train_labels.numpy()
# test_images = test_images.numpy()
# test_labels = test_labels.numpy()

# + [markdown] id="SMqxXreK5ghh"
#
# Let $\mathbf{X} \in [0,1]^{N \times 784}$ be a matrix of the MNIST training data set with $N = 60000$, where each row represents a training image. Moreover, let $\mathbf{X} - \overline{\mathbf{x}}^\intercal = \mathbf{U} \mathbf{\Sigma} \mathbf{V}^\intercal$ be the singular value decomposition of the MNIST training data after them mean has been removed, where $\overline{\mathbf{x}} = \frac{1}{N} \sum_{i=1}^N \mathbf{x}_i$.
#
# We compute the singular value decomposition $\mathbf{U} \boldsymbol{\Sigma} \mathbf{V}^\intercal$ with the function [`torch.svd`](https://pytorch.org/docs/stable/torch.html#torch.svd).

# + id="wkfNjs3XbMyC" tags=["keep"]
# center training images
train_mean = train_images.mean(axis=0)
train_images_centered = train_images - train_mean

U, S, V = torch.svd(train_images_centered)

# + [markdown] id="g9Jt2du7O4z1"
# ### Item (a)
#
# Let `train_encoding` $\in \mathbb{R}^{60000 \times 2}$ and `test_encoding` $\in \mathbb{R}^{10000 \times 2}$ be the two-dimensional principal subspace of the MNIST training images. Compute the two encodings of the images in the MNIST training and the test data set.  Do it with the help of $\mathbf{U}$, $\boldsymbol{\Sigma}$, and $\mathbf{V}$.
#
# *Hints*: Rememeber that the presence of the center $\overline{\textbf{x}}$ needs to be accounted for.
#

# + id="XnM5TzfNxh7f" colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1633499108677, "user_tz": -120, "elapsed": 424, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="30182feb-937d-445c-81a2-28060d35e19c"
train_encoding = torch.matmul(train_images_centered, V[:,0:2])
print('Train_encoding: ', train_encoding.shape)

test_encoding = torch.matmul(test_images - train_mean, V[:,0:2])
print('Test_encoding: ', test_encoding.shape)


# + [markdown] id="QqA3hX3ybMyZ"
# Generate 2D scatter plots showing of the datapoints you computed in the last item. The function `plot_encoding` bellow can be used to generate a plot of the latent space.

# + id="XnA0MGjBbMyc" tags=["keep"]
def plot_encoding(train_data, test_data):
    # create two plots side by side
    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(15, 6))

    # plot data
    titles = ["training data", "test data"]
    for (data, title, ax) in zip([train_data, test_data], titles, axes.flat):
        encodings, labels = data
        scatter = ax.scatter(encodings[:, 0], encodings[:, 1],
                             c=labels, cmap=plt.cm.tab10, vmin=-0.5,
                             vmax=9.5, alpha=0.7)
        ax.set_xlabel("$z_1$")
        ax.set_ylabel("$z_2$")
        ax.set_title(title)

    # add colorbar
    cb = fig.colorbar(scatter, ticks=np.arange(0, 10), ax=axes.ravel().tolist())
    cb.ax.set_title("digit")
    
    return fig


# + colab={"base_uri": "https://localhost:8080/", "height": 406} id="CEC8YzTmxh7h" tags=["keep"] executionInfo={"status": "ok", "timestamp": 1633499115673, "user_tz": -120, "elapsed": 3376, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="c429bb32-a86a-421f-cd5b-b14919478a6a"
plot_encoding((train_encoding, train_labels), (test_encoding, test_labels))
plt.show()

# + [markdown] id="QjsGYI6pbMyt"
# ### Item (b)
#
# With the help of the matrices $\mathbf{U}$, $\boldsymbol{\Sigma}$, and $\mathbf{V}$ computed in the last item. Map the `test_encoding` from the enconding space back to the space of image obtaining: `test_reconstruction` $\in \mathbb{R}^{10000 \times 784}$ which will consist of the images in the MNIST test data set obtained considering only the two principal components.

# + id="eNv8RcWWxh7k"
test_reconstruction = torch.matmul(test_encoding, np.transpose(V[:,0:2])) + train_mean


# + [markdown] id="QjtEPRQnbMyv"
# Plot some test images and their reconstructed counterparts. You can use the function `plot_reconstruction` bellow do the plotting. Which digits can be reconstructed and decoded quite well, and which ones seem to be more challenging?

# + id="yaF1Pr4Sxh7k"
# plot at most 8 random reconstructions
def plot_reconstruction(originals, reconstructions, labels, nrow=2):
    # indices of displayed samples
    n = originals.size(0)
    indices = range(n) if n <= 8 else np.random.choice(n, size=8, replace=False)

    # compute number of rows and columns of grid
    num_samples = len(indices)
    num_cols = min(nrow, num_samples)
    num_rows = int(math.ceil(float(num_samples) / nrow))

    fig = plt.figure(figsize=(5 * num_cols, 2.5 * num_rows))
    for (i, idx) in enumerate(indices):
        # extract original, reconstruction, and label
        original = originals[idx]
        reconstruction = reconstructions[idx]
        label = labels[idx]

        # configure subplot
        plt.subplot(num_rows, num_cols, i+1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        plt.title(f"Label: {label.item()}", fontweight='bold')

        # plot original and reconstructed image
        # design a grid, similar to torchvision.utils.make_grid
        grid = torch.ones(1, 1, 32, 62)
        grid.narrow(2, 2, 28).narrow(3, 2, 28).copy_(
            original.view(-1, 1, 28, 28))
        grid.narrow(2, 2, 28).narrow(3, 32, 28).copy_(
            reconstruction.view(-1, 1, 28, 28))
        plt.imshow(grid.squeeze().cpu().numpy(),
                   vmin=0.0, vmax=1.0, cmap='gray_r')

    return fig


# + id="pVhQvVPmbMyy" tags=["keep"] colab={"base_uri": "https://localhost:8080/", "height": 594} executionInfo={"status": "ok", "timestamp": 1633499492731, "user_tz": -120, "elapsed": 288, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="a637e4b8-1bff-43be-a9b4-c4ebc5e59ac8"
plot_reconstruction(test_images, test_reconstruction, test_labels)
plt.show()

# + [markdown] id="bNR9-T54xh7l"
# ### Item (c)
#
# The comparison of the original images and their reconstructions provides us with some intuition for how much information is lost by the compression of the images to the two-dimensional latent space. As a less subjective measure we calculate the average squared reconstruction error. An advantage of an objective measure such as the average squared reconstruction error is that it enables us to compare the PCA with other models for dimensionality reduction.
# \begin{equation*}
# \mathrm{sqerr} := \frac{1}{10000} \sum_{i=1}^{10000} \|\mathbf{x}_i - \tilde{\mathbf{x}}_i\|^2_2
# \end{equation*}
# of the images $\mathbf{x}_i \in {[0,1]}^{784}$ and their reconstruction $\tilde{\mathbf{x}}_i \in \mathbb{R}^{784}$ ($i = 1,\ldots, 10000$) in the MNIST test data set. 
#
# What average squared reconstruction error do you get with PCA?

# + id="Ho3lBFIyaKg6" tags=["keep"] colab={"base_uri": "https://localhost:8080/"} executionInfo={"status": "ok", "timestamp": 1633500299300, "user_tz": -120, "elapsed": 308, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="036db55b-b5fb-4a94-db20-e25dee04c9de"
sqerr = torch.sum((test_images - test_reconstruction)**2)/len(test_images)
print(f'Average squared reconstruction error: {sqerr}')

# + [markdown] id="rqr1928yo9M0"
# ### Item (d)
#
# For each digit from 1 to 9, compute the mean of the data training points in the encoding space.

# + id="z-lL95Elxh7i" tags=["keep"]
train_mean_encodings = torch.zeros(10,2)
repetitions = torch.zeros(10)
for (encoding, label) in zip(train_encoding, train_labels):
  repetitions[label] += 1
  train_mean_encodings[label,:] += (encoding-storage[label,])/(repetitions[label])


# + [markdown] id="WZawT13Jxh7j"
# Generate the images corresponding to each of the means obtained in the encode space. Use the SVD matrix to convert the encoding space back to the image space. Use the the function `plot_images` bellow to generate a plot of that image.

# + id="MfZfs9rZxh7j" tags=["keep"]
def plot_images(images, labels=None, nrow=5):
    # compute number of columns of grid
    num_samples = images.size(0)
    num_cols = min(nrow, num_samples)
    num_rows = int(math.ceil(float(num_samples) / nrow))

    fig = plt.figure(figsize=(2 * num_cols, 2 * num_rows))
    for i in range(num_samples):
        # extract image and labels if provided
        image = images[i]
        label = None if labels is None else labels[i]

        # configure subplot
        plt.subplot(num_rows, num_cols, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.grid(False)
        if label is not None:
            plt.title(f"Label: {label.item()}", fontweight='bold')

        # plot image
        plt.imshow(image.view(28, 28).cpu().numpy(),
                   vmin=0.0, vmax=1.0, cmap='gray_r')
        
    return fig


# + id="wXhNi4p9xh7j" tags=["keep"] colab={"base_uri": "https://localhost:8080/", "height": 265} executionInfo={"status": "ok", "timestamp": 1633502791157, "user_tz": -120, "elapsed": 966, "user": {"displayName": "Marcello Vendruscolo", "photoUrl": "https://lh3.googleusercontent.com/a/default-user=s64", "userId": "10393122258563683124"}} outputId="28789621-78a1-4282-db0e-97166f7ab9f9"
# compute mean images
train_mean_images = train_mean + train_mean_encodings.mm(V[:, :2].t())

plot_images(train_mean_images, torch.arange(10))
plt.show()

# + [markdown] id="kXOdBYQm1R29"
# ## Exercise 11.2 - Derivations for probabilistic PCA
#
#
# In constrast to (regular) PCA, the so-called probabilistic PCA (PPCA) allows a probabilistic
# interpretation of the principal components. The probabilistic formulation
# of PCA also allows us to extend the method and alter its underlying assumptions
# quite easily, as we will learn during the course of this laboratory.
#
#
# As before, let $\mathbf{x} \in \mathbb{R}^D$ represent a data sample that we
# want to decode from a lower dimensional representation
# $\mathbf{z} \in \mathbb{R}^M$ with $M < D$. The PPCA model assumes that
# $\mathbf{z}$ is standard normally distributed and $\mathbf{x}$
# can be decoded by a noisy linear transformation of $\mathbf{z}$.
# Mathematically, the model is given by
# \begin{align*}
#   p(\mathbf{x} \,|\, \mathbf{z}) &= \mathcal{N}\left(\mathbf{x}; \mathbf{W}\mathbf{z} + \boldsymbol{\mu}, \sigma^2 \mathbf{I}_D\right), \\
#   p(\mathbf{z}) &= \mathcal{N}(\mathbf{z}; \boldsymbol{0}, \mathbf{I}_M),
# \end{align*}
# with parameters $\mathbf{W} \in \mathbb{R}^{D \times M}$,
# $\boldsymbol{\mu} \in \mathbb{R}^D$, and $\sigma^2 > 0$.
# [Michael E. Tipping and Christopher M. Bishop show in "Probabilistic Principal Component Analysis"](https://www.jstor.org/stable/2680726)
# that for $\sigma^2 \to 0$
# the model recovers the standard PCA (but the components of $\mathbf{z}$ might
# be permuted).
#
# We assume that the data $\mathbf{x}_1, \ldots, \mathbf{x}_N$ is identically
# and independently distributed according to the PPCA model. In a maximum
# likelihood setting, one determines the parameters $\mathbf{W}$, $\boldsymbol{\mu}$,
# and $\sigma^2$ that maximize the likelihood
# \begin{equation*}
#   p(\mathbf{x}_1, \ldots, \mathbf{x}_N ; \mathbf{W}, \boldsymbol{\mu}, \sigma^2)
#   = \prod_{n=1}^N p(\mathbf{x}_n; \mathbf{W}, \boldsymbol{\mu}, \sigma^2),
# \end{equation*}
# or equivalently the log-likelihood
# \begin{equation*}
#   \log p(\mathbf{x}_1, \ldots, \mathbf{x}_N; \mathbf{W}, \boldsymbol{\mu}, \sigma^2)
#   = \sum_{n=1}^N \log p(\mathbf{x}_n; \mathbf{W}, \boldsymbol{\mu}, \sigma^2).
# \end{equation*}

# + [markdown] id="xQffYFcTy3jI"
# ### Item (a) 
#
# *(Pen and paper exercise)*
#
# Show that for the model of the probabilistic PCA
#   \begin{equation*}
#     p(\mathbf{x}) = \mathcal{N}(\mathbf{x}; \boldsymbol{\mu}, \mathbf{C}),
#   \end{equation*}
#   where $\mathbf{C} = \mathbf{W}\mathbf{W}^\intercal + \sigma^2 \mathbf{I}_D$.
#
#

# + [markdown] id="3UCseyss1Vsu"
# ### Item (b)
# *(Pen and paper exercise)*
#
# Show that the distribution of the latent variable $\mathbf{z}$ conditioned on
#   $\mathbf{x}$ is Gaussian as well and given by
#   \begin{equation*}
#     p(\mathbf{z} \,|\, \mathbf{x}) = \mathcal{N}\left(\mathbf{z}; \mathbf{M}^{-1} \mathbf{W}^\intercal (\mathbf{x} - \boldsymbol{\mu}), \sigma^2 \mathbf{M}^{-1} \right),
#   \end{equation*}
#   where $\mathbf{M} = \mathbf{W}^\intercal \mathbf{W} + \sigma^2 \mathbf{I}_M$.
#
#
#

# + [markdown] id="GR9ls8C-s_dd"
# ## Exercise 11.3 - Gaussian Mixture Models
#
#
# Consider the dataset generated bellow. It was generated using a gaussian mixture. That is the dataset is generated as follows:
#
# - We sample a variable `pp` from a bernouly distribution with probability `pi = 0.7`.
# - When the result is equal to one (`pp == 1`) we sample from the normal:
# $$\mathcal{N}\left(\mu_1 = \begin{bmatrix} 1 \\ 2\end{bmatrix}, \Sigma_1 = \begin{bmatrix} 2 & 0 \\ 0 & 1 \end{bmatrix}\right)$$
#
# - When the result is equal to one (`pp == 0`) we sample from the normal:
# $$\mathcal{N}\left(\mu_2 = \begin{bmatrix} 5 \\ 7\end{bmatrix}, \Sigma_2 = \begin{bmatrix} 2 & 1 \\ 1 & 2 \end{bmatrix}\right)$$
#
#
# In this exercise, you will use an expectation maximization algorithm to estimate the parameters of gaussian mixture model to fit this dataset. Using what you implemented, you will apply this model to the two dimensional encondings from the exercise 11.1.

# + id="YHHYEKlcs_TX"
from scipy.stats import multivariate_normal, uniform, bernoulli

def plot_2dnormal(ax, vnormal, color='red'):
  # Plot center
  plt.plot(vnormal.mean[0], vnormal.mean[1], '*', ms=5, color=color)
  # Get countour grid
  x, y = np.mgrid[-3:3:.01, -3:3:.01]
  # rescale x, y
  x = vnormal.cov[0,0] * x +  vnormal.mean[0]
  y = vnormal.cov[1,1] * y +  vnormal.mean[1]
  # Plot countour
  pos = np.dstack((x, y)) 
  ax.contour(x, y, vnormal.pdf(pos), colors=color, levels=3)

# First normal distribution
mu1 = [1, 2]
cov1 = [[2, 0], [0, 1]]
n1 = multivariate_normal(mean=mu1, cov= cov1)
# Parameters of the second normal distribution
mu2 = [5, 7]
cov2 = [[2, 1], [1, 2]]
n2 = multivariate_normal(mean=mu2, cov=cov2)
# Probability of belonging to each
pi = 0.7
pp = bernoulli(p=pi)
odds = [pi, 1 - pi]
# Create a list with all the normals
vnormals = [n1, n2]

# Sample N=300 samples from each
N = 300
v1 = n1.rvs(N)
v2 = n2.rvs(N)
p = pp.rvs(N)

gmm_data = np.where(p[:, None] == 1, v1, v2)

fig, ax = plt.subplots()
ax.scatter(gmm_data[:,0], gmm_data[:,1])
plot_2dnormal(ax, n1, color='red')
plot_2dnormal(ax, n2, color='blue')
plt.show()


# + [markdown] id="fluE4c4ZCRx1"
# ### Item (a)
#
# Implement a function `e_step`  that receives as input:
# -  the `gmm_data`: which is an `300 x 2` np array generated above;
# -  the `odds`: which is a list containing as elements the `pi` and `1 - pi`; and
# - the `vnormals`: a list  `[n1, n2]` of variable generated using: (`multivariate_normal`)[https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.multivariate_normal.html]. See `vnormals` above.
#
# and outputs a vector of dimension `300 x 2 ` containing the probability of `pp = 1` for each of the points. This is equivalent, in each row:
# $$\left(\frac{\pi_1\mathcal{N}\left(x; \mu_1 , \Sigma_1\right)}{\sum_i \pi_i\mathcal{N}\left(x; \mu_i , \Sigma_i\right)},  \frac{\pi_2\mathcal{N}\left(x; \mu_2 , \Sigma_2\right)}{\sum_i \pi_i\mathcal{N}\left(x; \mu_i , \Sigma_i\right)}\right) $$
#
# This computation corresponds to expectation step in the expectation maximization algorithm.

# + id="pamMsGZ3Cdhx"
def e_step(gmm_data, odds, vnormals):
    return ## TODO


# + [markdown] id="Ngfkjp6cBGzM"
# Bellow the function uses the values produced by such function to computed to decide whether a point from the data is associated with each center.

# + id="TlmAnsuLCfBR"
# appy the function to the dataset
ws = e_step(gmm_data, odds, vnormals)

# Use it to decide to which normal it belongs to
fig, ax = plt.subplots()
plot_2dnormal(ax, n1, color='black')
plot_2dnormal(ax, n2, color='black')
belong_to = ws.argmax(axis=0)
ax.scatter(gmm_data[belong_to==0,0], gmm_data[belong_to==0,1], color='red')
ax.scatter(gmm_data[belong_to==1,0], gmm_data[belong_to==1,1], color='blue')
plt.show()


# + [markdown] id="5a8a_4ny7IwG"
# ### Item (b)
#
# Implement a function `m_step`  that receives as input:
# -  the `gmm_data`: which is an `300 x 2` np array generated above;
# -  the `ws`: which is an `300 x 2` np array giving in the first row the probability of a point belonging to the first normal in the mixture. And, in the second row, the probability of a point belonging to the first normal in the mixture. I.e. similar to the output of the function `e_step` implemented above.
#
# It outputs the parameters of the mixture:
# -  the `odds`: which is a list containing as elements the `pi` and `1 - pi`; and
# - the `vnormals`: a list  `[n1, n2]` of variable generated using: (`multivariate_normal`)[https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.multivariate_normal.html]. See `vnormals` above.
#
# This can be done using Eq. 10.16 from http://smlbook.org/.

# + id="aaqRa9mDFpxP"
def m_step(gmm_data, ws):
    return ### TODO


# + [markdown] id="rvqfsNgWCt7N"
# You can use the values computed in the last few items as a sanity check:
#
#

# + id="JbLA6_i_GDQP"
odds, vnormals = m_step(gmm_data, ws)

print('odds:')
print(odds)
print('')
for i, rv in enumerate(vnormals):
  print('Normal {}'.format(i+1))
  print('- mean:')
  print(rv.mean)

  print('- cov:')
  print(rv.cov )
  print('')

# + [markdown] id="dq-x0rL8GPE5"
# If everything went right, you should obtain values close to the ones used to generate the dataset:
#
# - `odds` $\approx$ [0.7 0.3]
# - The first normal: $$\mu_1 = \begin{bmatrix} 1 \\ 2\end{bmatrix}, \Sigma_1 = \begin{bmatrix} 2 & 0 \\ 0 & 1 \end{bmatrix}$$
# -  The first normal: $$\mu_2 = \begin{bmatrix} 5 \\ 7\end{bmatrix}, \Sigma_2 = \begin{bmatrix} 2 & 1 \\ 1 & 2 \end{bmatrix}$$
#

# + [markdown] id="5c31XGxiLb1D"
# ### Item (c)
# Use the code bellow to plot 3 iterations of the Expectation-Maximization algorithm in the dataset generated above.

# + id="moydTocLvWbQ"
# Initialization
mu1 =  [0, 0]
cov1 = [[1, 0], [0, 1]]
n1_estimated = multivariate_normal(mean=[0, 0], cov= cov1)
# Parameters of the second normal distribution
mu2 = [10, 10]
cov2 = [[1, 0], [0, 1]]
n2_estimated  = multivariate_normal(mean=mu2, cov=cov2)
vnormals_estimated = [n1_estimated, n2_estimated]
# Probability of belonging to each
odds_estimated = [0.5, 0.5]


for i in range(3):
  fig, ax = plt.subplots()
  ws = e_step(gmm_data, odds_estimated, vnormals_estimated)
  belong_to = ws.argmax(axis=0)
  ax.scatter(gmm_data[:, 0], gmm_data[:, 1], c=belong_to)
  for j in range(ws.shape[0]):
    plot_2dnormal(ax, vnormals_estimated[j], color='black')

  odds_estimated, vnormals_estimated = m_step(gmm_data, ws)
  plt.show()

# + [markdown] id="gIk96nSiPy1E"
# ### Item (d)
#
# Adapt the code above and apply the method to MNIST 2 dimensional encoding you generated in Exercise 11.1. You can choose how many centers you want to use and what initialization parameters.

# + id="mWSxgXA0P2wX"

