import numpy as np
import matplotlib.pyplot as plt
from __future__ import division
import pandas as pd
from math import exp

w, h = 6, 4
data = [[0 for x in range(w)] for y in range(h)]

A = 0
B = 0
values = []
z = []
a0_list = [0.33, 0.5, 0.66, 0.75]
sigma0 = 1
sigma1 = 1
m0_list = [2, 3, 4, 5]
m1_list = [-2, -3, -4, -5]
N = 150
n0 = 0
n1 = 0
q1 = []

for j in range(4):

    a0 = a0_list[j]
    m0 = m0_list[j]
    m1 = m1_list[j]
    n0 = 0
    n1 = 0
    m0_curr = m0
    m1_curr = m1
    for t in range(200):
        n1 = 0
        n0 = 0
        q1 = []
        values = []
        for i in range(N):
            # generating values[i] and q1[i]
            alpha = np.random.uniform(0, 1, 1)[0]
            if (alpha < a0):
                values.append(np.random.normal(m0, sigma0, 1)[0])
            else:
                values.append(np.random.normal(m1, sigma1, 1)[0])
            q1.append(exp(-(values[i] - m1_curr)*(values[i] - m1_curr) / 2) / (exp(-(values[i] - m0_curr) * (values[i] - m0_curr) / 2) + exp(-(values[i] - m1_curr) * (values[i] - m1_curr) / 2)))
            n1 = n1 + q1[i]

        n0 = N - n1
        m0_curr = 0
        m1_curr = 0
        for f in range(N):
            m0_curr = m0_curr + values[f] * (1 - q1[f])
            m1_curr = m1_curr + values[f] * q1[f]

        m0_curr = m0_curr / n0
        m1_curr = m1_curr / n1

    m0_hat = m0_curr
    m1_hat = m1_curr

    data[j][0] = a0
    data[j][1] = m0
    data[j][2] = m1

    data[j][3] = round(n0 / N, 3)
    data[j][4] = round(m0_hat, 3)
    data[j][5] = round(m1_hat, 3)


df = pd.DataFrame(data = data)
df.columns = ['p_real', 'm0_real', 'm1_real', 'p_estimated', 'm0_estimated', 'm1_estimated']
df.head()