#encoding:utf-8

import random
import math
from matplotlib import pyplot
import sys
import numpy

pyplot.rcParams['font.family'] = 'Times New Roman'


def f(x, m, s):
    return math.exp(-0.5 * (x - m) ** 2 / s) / math.sqrt(2 * math.pi * s)


def gaussian(m, s):
    def f(x):
        return numpy.exp(-0.5 * (x - m) ** 2 / s) / numpy.sqrt(2 * math.pi * s)
    return f


def calc_likelihood(x, m, s, p):
    sum = 0.0
    for n in range(len(x)):
        temp = 0.0
        for k in range(len(p)):
            temp += p[k] * f(x[n], m[k], s[k])
        sum += math.log(temp)
    return sum


# pi = [1 - random.random() for _ in range(k)]
pi = [0.3, 0.7]

k = len(pi)
n = 500

s = 0
for elem in pi:
    s += elem
for i in range(k):
    pi[i] /= s

#mu = [random.uniform(-k, k) for _ in range(k)]
#sigma = [random.uniform(0.1, k) for _ in range(k)]

mu = [-5, 5]
sigma = [1, 1]

x = [0 for _ in range(n)]

mini = 100000
maxi = -100000

for i in range(n):
    r = random.uniform(0, 1)
    cum = 0
    for j in range(k):
        if cum <= r and r <= cum + pi[j]:
            x[i] = random.gauss(mu[j], sigma[j])
            mini = min(x[i], mini)
            maxi = max(x[i], maxi)
        cum += pi[j]

print(pi)
print(mu)
print(sigma)

pi_est = [1.0 / k for _ in range(k)]
mu_est = [random.uniform(mini, maxi) for _ in range(k)]
sigma_est = [random.uniform(0.1, maxi - mini) for _ in range(k)]

l_prev = 0
step = 0
while True:
    step += 1
    gamma = [[0 for _ in range(n)] for _ in range(k)]

    for t in range(n):
        tmp = 0
        for i in range(k):
            gamma[i][t] = pi_est[i] * f(x[t], mu_est[i], sigma_est[i])
            tmp += gamma[i][t]
        for i in range(k):
            gamma[i][t] /= tmp

    for i in range(k):
        tmp = 0
        for t in range(n):
            tmp += gamma[i][t]
        pi_est[i] = tmp / n

        tmp1 = 0
        tmp2 = 0
        tmp3 = 0
        for t in range(n):
            tmp1 += gamma[i][t] * x[t] * x[t]
            tmp2 += gamma[i][t] * x[t]
            tmp3 += gamma[i][t]
        mu_est[i] = tmp2 / tmp3
        sigma_est[i] = tmp1 / tmp3 - mu_est[i] * mu_est[i]

    l = calc_likelihood(x, mu_est, sigma_est, pi_est)

    print(l)

    if (abs(l - l_prev) < 1e-8):
        break
    l_prev = l

print("converged at step %d" % step)

print("Real Value:")
print(pi)
print(mu)
print(sigma)

print("Estimated Value:")
print(pi_est)
print(mu_est)
print(sigma_est)

n, bins, _ = pyplot.hist(x, 100, normed = 1, alpha = 0.3)
seq = numpy.arange(-15, 15, 0.02)
for i in range(len(pi_est)):
    pyplot.plot(seq, pi_est[i] * gaussian(mu_est[i], sigma_est[i])(seq), linewidth = 2.0)
pyplot.show()