# -*- coding: utf-8 -*-
"""GMM_EM.ipynb

Automatically generated by Colaboratory.

# Gaussian mixture model with expectation maximization algorithm

GMM with EM.

This notebook implements the following:

1) Function that avoids computing inverse of matrix when computing y=(A_inverse)x by solving system of linear equations.

2) Log sum trick to avoid underflow when multiplying small numbers.

3) pdf of the Multivariate normal distribution

4) E-step function of the EM algorithm

5) M-step function of the EM algorithm

6) Variational lower bound function

7) GMM function

8) Training function for GMM

9) Scatter plot of clusters (Plot at the bottom of this notebook shows 8 clusters from a dataset of 100 points)

# Imports
"""

# Commented out IPython magic to ensure Python compatibility.
import sys
import numpy as np
from numpy.linalg import det, solve
import matplotlib
import matplotlib.pyplot as plt
# %matplotlib inline

"""# Variables info

N: number of data (rows)

d: dimension of data

X: (N x d), data points

C: int, number of clusters

Computed from E-step:

gamma: (N x C), distribution q(T), probabilities of clusters for objects

Initial values are also subsequently computed & updated from M-step:

pi: (C), mixture component weights, initial weights of T (latent variable), sum to 1, t=1, 2 or 3.
mu: (C x d), mixture component means

sigma: (C x d x d), # mixture component covariance matrices

# Generate random data.
"""

N = 100
d = 2

X = np.random.rand(N,d)

print(X[:5])
print(X.shape)

fig, ax = plt.subplots(1,1, figsize=(15,10))
ax.set_title('Data')
ax.scatter(X[:, 0], X[:, 1], c='black', s=100)
#plt.axis('equal')
plt.show()

"""# Generate initial values"""

epsilon = 1e-10 # Use in stopping criterion as well as in preventing numerical errors.

C = 7

def rand_input(C, d):
  # https://stackoverflow.com/questions/18659858/generating-a-list-of-random-numbers-summing-to-1
  pi0 = np.random.dirichlet(np.ones(C),size=1)[0] # Generating a list of random numbers, summing to 1
  mu0 = np.random.rand(C,d)
  sigma0 = np.random.rand(C,d,d)
  return pi0, mu0, sigma0

pi0, mu0, sigma0 = rand_input(C, d)

print(pi0)
print(pi0.shape)
print(mu0)
print(mu0.shape)
print(sigma0)
print(sigma0.shape)

"""# Avoid computing inverse of matrix when computing y=(A_inverse)x."""

# Function which avoids computing inverse of matrix when computing y=(A_inverse)x by solving linear equations.
# Use in E-step.

def _A_inverse_times_X(A, X):
  # A is nxn
  # X is rxn

  Y = []
  for row_data in X:
    Y_new = np.linalg.solve(A, row_data)
    Y.append(Y_new)

  Y = np.asarray(Y)
  assert Y.shape == X.shape, "Output shape must be equal to shape of X."

  return Y

"""# Multivariate normal (Gaussian) distribution

$MVN = \frac{1}{\sqrt{(2\pi)^n|\boldsymbol\Sigma_c|}}
\exp\left(-\frac{1}{2}({x}-{\mu_c})^T{\boldsymbol\Sigma_c}^{-1}({x}-{\mu_c})\right)$

Computes pdf of Multivariate normal (Gaussian) distribution.
"""

# Alternatively, one could also use multivariate_normal.pdf from scipy.stats
# instead of this function

def __mvg(cov, X, mean):
  diff = X - mean
  Y = _A_inverse_times_X(cov, diff)
  pow_term = -0.5 * np.matmul(Y, diff.T)
  e_term = np.exp(pow_term)

  const_term = (2*np.pi)**(X.shape[1])
  det_term = np.linalg.det(cov)
  deno_term = np.sqrt(np.multiply(const_term, det_term))

  P = np.divide(e_term, deno_term)

  return P.diagonal() # returns the pdf, shape=(num_X,)

# Returns pdf for multiple components.

def _mvg(cov, X, mean):
  P = []
  for i, r in enumerate(mean):
    P.append(__mvg(cov[i], X, mean[i]))

  return P # shape=(C, num_X)

"""# Log sum trick"""

# log sum trick to prevent underflow in E-step.
# https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/
# https://web.archive.org/web/20150502150148/http://machineintelligence.tumblr.com/post/4998477107/the-log-sum-exp-trick
# https://www.quora.com/Why-is-e-log_-e-x-equal-to-x

def exp_normalize(x):
  b = x.max()
  y = np.exp(x - b)

  return y / y.sum()

"""# E-step

Multiply the initial weight of the class with the multivariate guassian pdf of each data from the class.
"""

def E_step(X, pi, mu, sigma):
    N = X.shape[0] # number of objects
    C = pi.shape[0] # number of clusters
    d = mu.shape[1] # dimension of each object
    gamma = np.zeros((N, C)) # distribution q(T)
    gamma = np.mat(np.zeros((N, C)))

    prob = _mvg(sigma, X, mu) # pdf of data X in each class
    prob = np.mat(prob)

    for c in range(C):
        # Instead of multiplying probabilities directly which could result in underflow,
        # we'll work in log scale.
        # pi[c] = P(T=c), prob[c, :] = P(X|T=c)
        #gamma[:, c] = np.multiply(pi[c], prob[c, :].T)
        gamma[:, c] = np.log(pi[c] + epsilon) + np.log(prob[c, :].T + epsilon)

    for i in range(N):
        # Instead of summing the denominator, we'll use the log sum trick coded in exp_normalize function.
        gamma[i, :] = exp_normalize(gamma[i, :])

    return gamma # Q(T) = P(T|X,theta), weights of each model (class) for each data in X.

"""# M-Step

Compute the following:

[Equations from wiki](https://en.wikipedia.org/wiki/Expectation%E2%80%93maximization_algorithm#E_step)

![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/0e0327c8676ae66ec651b422a19f5ea532913c7a)

![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/45f3e73f50d396aadc98182709eee0c0d513aa6b)

![alt text](https://wikimedia.org/api/rest_v1/media/math/render/svg/a92651be432155520db19dc0b4da807039d96eb0)
"""

def M_step(X, gamma):
    N = X.shape[0] # number of objects
    C = gamma.shape[1] # number of clusters
    d = X.shape[1] # dimension of each object

    mu = np.zeros((C, d))
    sigma = []
    pi = np.zeros(C)

    # for each model in C
    for c in range(C):
        # sum of all Q(t) of model c
        sum_Q_t = np.sum(gamma[:, c])

        # mean of model c
        mu[c, :] = np.sum(np.multiply(X, gamma[:, c]), axis=0) / sum_Q_t

        # cov of model c
        diff = X - mu[c]
        sigma.append(diff.T @ np.multiply(diff, gamma[:, c]) / sum_Q_t)

        # weight of model c
        pi[c] = sum_Q_t / N

    return pi, mu, np.asarray(sigma)

"""# Variational lower bound

Computes the scalar output of the following:

$$\sum_{i=1}^{N} \sum_{c=1}^{C} q(t_i =c) (\log \pi_c + \log(MVN)) - \sum_{i=1}^{N} \sum_{c=1}^{K} q(t_i =c) \log q(t_i =c)$$
"""

def compute_vlb(X, pi, mu, sigma, gamma):
    """
    Each input is numpy array:
    X: (N x d), data points
    gamma: (N x C), distribution q(T)
    pi: (C)
    mu: (C x d)
    sigma: (C x d x d)

    Returns value of variational lower bound
    """
    N = X.shape[0] # number of objects
    C = gamma.shape[1] # number of clusters
    d = X.shape[1] # dimension of each object

    VLB = 0.0
    for c in range(C):
      mu_c = np.expand_dims(mu[c,:], axis=0)
      sigma_c = np.expand_dims(sigma[c,:], axis=0)
      gamma_c = gamma[:,c]

      mvg = np.asarray(_mvg(sigma_c, X, mu_c)) # 1xc
      sum = np.log(pi[c] + epsilon) + np.log(mvg + epsilon) # 1xc, + 1e-30 to prevent log(0)
      prod = np.multiply(gamma_c, sum.T) # transpose sum for element wise multiplication
      prod2 = np.multiply(gamma_c, np.log(gamma_c + epsilon)) # element wise multiplication, + 1e-30 to prevent log(0)
      VLB += (prod - prod2)

    VLB = np.sum(VLB, axis=0) # sum all values for all rows

    return VLB

"""# GMM

Find the best parameters by optimizing with the following criterion:

Stopping threshold: ($|\frac{\mathcal{L}_i-\mathcal{L}_{i-1}}{\mathcal{L}_{i-1}}| \le \text{threshold}$)
"""

def GMM(X, C, d, threshold=epsilon, max_iter=1000, trial=500):
    N = X.shape[0] # number of objects
    d = X.shape[1] # dimension of each object
    best_VLB = None
    best_pi = None
    best_mu = None
    best_sigma = None

    for rs in range(trial):
        try:
            pi, mu, sigma = rand_input(C, d) # Try random initial values
            curr_LVB, prev_LVB = 0.0, 0.0
            iter = 0
            while iter < max_iter:
                #print('iter, rs', iter, rs)

                prev_LVB = curr_LVB

                gamma = E_step(X, pi, mu, sigma)
                pi, mu, sigma = M_step(X, gamma)
                curr_LVB = compute_vlb(X, pi, mu, sigma, gamma)

                #print('prev_LVB', prev_LVB)
                #print('curr_LVB', curr_LVB)

                # LVB is the variation lower bound function. It must NOT be decreasing.
                # We are trying to maximize LVB so that the gap between LVB & GMM is minimized.
                if prev_LVB != 0.0 and curr_LVB < prev_LVB:
                    print('VLB ERROR EXIT!: curr_LVB < prev_LVB')
                    sys.exit(1)

                # If numerical error in LVB, goto next trial.
                if np.isnan(curr_LVB) == True:
                    break

                if prev_LVB != 0.0 and abs((curr_LVB - prev_LVB) / (prev_LVB)) <= threshold:
                    if best_VLB == None or curr_LVB > np.float32(best_VLB):
                        best_VLB = curr_LVB
                        best_pi = pi
                        best_mu = mu
                        best_sigma = sigma
                    break # end while loop, goto for loop

                iter += 1

        except np.linalg.LinAlgError:
            print("Singular matrix not allowed.")
            pass

    return best_VLB, best_pi, best_mu, best_sigma

"""# Train"""

# Train
# If numerical errors occured, run a couple of more times.

best_VLB, best_pi, best_mu, best_sigma = GMM(X, C, d)
print('best_VLB', best_VLB)
print('best_pi', best_pi)
print('best_mu', best_mu)
print('best_sigma', best_sigma)

# Use the best values to do 1 more E-step to get gamma.
gamma = E_step(X, best_pi, best_mu, best_sigma)
labels = np.ravel(gamma.argmax(axis=1))

"""# Scatter plot"""

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

'''
# Generate colors for each class.
# only works for max of 4 classes.
def gen_col(C):
  colors =[]
  for c in range(C):
    colors.append(np.random.randint(0, 255, C) / 255)
  print(colors)
  return colors

colors = gen_col(C)

plt.scatter(X[:, 0], X[:, 1], c=labels, cmap=matplotlib.colors.ListedColormap(colors), s=30)
plt.axis('equal')
plt.show()
'''

# https://stackoverflow.com/questions/12487060/matplotlib-color-according-to-class-labels

N = C # Number of labels

# setup the plot
fig, ax = plt.subplots(1,1, figsize=(15,10))
# define the data
x = X[:, 0]
y = X[:, 1]
tag = labels

# define the colormap
cmap = plt.cm.jet
# extract all colors from the .jet map
cmaplist = [cmap(i) for i in range(cmap.N)]
# create the new map
cmap = cmap.from_list('Custom cmap', cmaplist, cmap.N)

# define the bins and normalize
bounds = np.linspace(0,N,N+1)
norm = mpl.colors.BoundaryNorm(bounds, cmap.N)

# make the scatter
scat = ax.scatter(x, y, c=tag, s=np.random.randint(100,500,N), cmap=cmap, norm=norm)

# create the colorbar
cb = plt.colorbar(scat, spacing='proportional',ticks=bounds)
cb.set_label('Custom cbar')
ax.set_title('Discrete color mappings')
plt.show()
