import math
from math import exp,log,sqrt
PI = math.pi

EPSILON = 0.000001                      # Avoid division by zero
DELTA = 0.00001

tMax = 10000


#data = [2.8962, 0.4077, 1.1863, 2.4843]
data = [float(x) for x in file("data.txt").readlines()]
iMax = 2
jMax = len(data)

p = [[0.5], [0.5]]
mu = [[-1.0], [1.0]]
var = [[1.0], [1.0]]
#pij = [[[None], [None], [None], [None]], [[None], [None], [None], [None]]]
pij = [[[None] for j in range(jMax)] for i in range(iMax)]

def gaussianComponentDensity(i, x, t):
    expt = (-(x - mu[i][t])**2) / (2 * var[i][t])
    return 1/(sqrt(2 * PI * var[i][t])) * exp(expt)

def gaussianMixtureDensity(x, t):
    return sum([p[i][t] * gaussianComponentDensity(i, x, t)
                for i in range(iMax)])

def eStep(t):
    # Calc pij
    for i in range(iMax):
        for j in range(jMax):
            pij[i][j].append(gaussianComponentDensity(i, data[j], t) *
                             p[i][t] /
                             gaussianMixtureDensity(data[j], t))

def mStep(t):
    for i in range(iMax):
        sumpij = sum([pij[i][j][t+1] for j in range(jMax)])
        
        # Compute mu_i
        mu[i].append(sum([pij[i][j][t+1] * data[j]
                          for j
                          in range(jMax)]) / sumpij)

        # Compute var_i
        var[i].append(sum([pij[i][j][t+1] *
                           (data[j] - mu[i][t+1])**2
                           for j in range(jMax)]) / sumpij)
        if var[i][t+1] == 0:
            var[i][t+1] = EPSILON

        # Compute p_i
        p[i].append(1.0/jMax * sumpij)

def calcLogLikelihood(t):
    ll = 0.0
    for j in range(jMax):
        llj = log(gaussianMixtureDensity(data[j], t))
        ll += llj
    return ll

def printVars(t):
    print "E-STEP", t
    for i in range(iMax):
        print " p_%d,j =" % i, [pij[i][j][t] for j in range(jMax)]

    print "M-STEP", t
    for i in range(iMax):
        print " p_%d =" % i, p[i][t]
        print " mu_%d =" % i, mu[i][t]
        print " var_%d =" % i, var[i][t]
    print "Log-likelihood:", calcLogLikelihood(t)

def printVarsLatex(t, printPIJs=True):
    print "%d &" % t
    if printPIJs:
        for j in range(jMax):
            if pij[0][j][t] == None:
                print " &"
            else:
                print "%5.4f &" % pij[0][j][t]

    for i in range(iMax):
        print " %5.4f &" % p[i][t]
        print " %5.4f &" % mu[i][t]
        print " %5.4f &" % var[i][t]
    print "%5.4f \\\\" % calcLogLikelihood(t)

def quiescent(t, delta):
    for i in range(iMax):
        if abs(p[i][t+1] - p[i][t]) > delta:
            return False
        if abs(mu[i][t+1] - mu[i][t]) > delta:
            return False
        if abs(var[i][t+1] - var[i][t]) > delta:
            return False
    return True

# for t in range(tMax):
#     printVarsLatex(t)
#     eStep(t)
#     mStep(t)
# printVarsLatex(tMax)

for t in range(tMax):
    eStep(t)
    mStep(t)
    if quiescent(t, DELTA):
        printVarsLatex(t, False)
        printVarsLatex(t+1, False)
        break