docs = [[0,0,1], [2,2,1], [0,1], [2]]
theta = [1/3, 2/3]
beta = [[1/4,1/4,2/4], [2/4,1/4,1/4]]

M=3
K=2
D=4

c = [[0 for x in range(D)] for y in range(M)]
r = [[0 for x in range(D)] for y in range(K)]

# compute counts
for d in range(D):
    for i in range(len(docs[d])):
        c[docs[d][i]][d] += 1

print("COUNTS:", c)

for it in range(20):
    
    print("Iteration", it)
    
    # compute responsibilities
    for d in range(D):
        total = 0
        for k in range(K):
            r[k][d] = theta[k]
            for m in range(M):
                r[k][d] *= beta[k][m]**c[m][d]
            total += r[k][d]
        for k in range(K):
            r[k][d] /= total

    print("RESPONSIBILITIES:",r)

    # update theta
    total = 0
    for k in range(K):
        theta[k] = 0
        for d in range(D):
            theta[k] += r[k][d]
            total += r[k][d]
    for k in range(K):
        theta[k] /= total

    # update beta
    for k in range(K):
        total = 0
        for m in range(M):
            beta[k][m] = 0 
        for m in range(M):
            for d in range(D):
                beta[k][m] += r[k][d]*c[m][d]
                total += r[k][d]*c[m][d]
        for m in range(M):
            beta[k][m] /= total

    print("THETA:", theta)
    print("BETA:",beta)