# Code from Chapter 7 of Machine Learning: An Algorithmic Perspective (2nd Edition)
# by Stephen Marsland (http://stephenmonika.net)

# You are free to use, change, or redistribute the code in any way you wish for
# non-commercial purposes, but please maintain the name of the original author.
# This code comes with no warranty of any kind.

# Stephen Marsland, 2008, 2014

import pylab as pl
import numpy as np


def GMM():
    """ Fits two Gaussians to data using the EM algorithm """
    N = 100
    pl.ion()

    y = 1. * np.zeros(N)
    # Set up data
    out1 = np.random.normal(6, 1, N)
    out2 = np.random.normal(1, 1, N)
    choice = np.random.rand(N)

    w = [choice >= 0.5]
    y[w] = out1[w]
    w = [choice < 0.5]
    y[w] = out2[w]

    pl.clf()
    pl.hist(y, fc='0.5')

    # Now do some learning

    # Initialisation
    mu1 = y[np.random.randint(0, N - 1, 1)]
    mu2 = y[np.random.randint(0, N - 1, 1)]
    s1 = np.sum((y - np.mean(y)) ** 2) / N
    s2 = s1
    pi = 0.5

    # EM loop
    count = 0
    gamma = 1. * np.zeros(N)
    nits = 20

    ll = 1. * np.zeros(nits)

    while count < nits:
        count = count + 1

        # E-step
        for i in range(N):
            gamma[i] = pi * np.exp(-(y[i] - mu1) ** 2 / (2 * s1)) / (
                pi * np.exp(-(y[i] - mu1) ** 2 / (2 * s1)) + (1 - pi) * np.exp(-(y[i] - mu2) ** 2 / 2 * s2))

        # M-step
        mu1 = np.sum((1 - gamma) * y) / np.sum(1 - gamma)
        mu2 = np.sum(gamma * y) / np.sum(gamma)
        s1 = np.sum((1 - gamma) * (y - mu1) ** 2) / np.sum(1 - gamma)
        s2 = np.sum(gamma * (y - mu2) ** 2) / np.sum(gamma)
        pi = np.sum(gamma) / N

        ll[count - 1] = np.sum(
            np.log(pi * np.exp(-(y[i] - mu1) ** 2 / (2 * s1)) + (1 - pi) * np.exp(-(y[i] - mu2) ** 2 / (2 * s2))))

    x = np.arange(-2, 8.5, 0.1)
    y = 35 * pi * np.exp(-(x - mu1) ** 2 / (2 * s1)) + 35 * (1 - pi) * np.exp(-(x - mu2) ** 2 / (2 * s2))

    pl.plot(x, y, 'k', linewidth=4)
    pl.figure(), pl.plot(ll, 'ko-')
    pl.show()


GMM()
