import matplotlib.pyplot as plt
import numpy as np
# from pprint import pprint

EPSILON = 0.1
G = {
    0: {'color': 'blue'},
    1: {'color': 'red'},
}


def openData(f_path):
    f = open(f_path, 'r')
    inputData = np.matrix([
        [float(x) for x in l.split()]
        for l in f.readlines()
    ])
    N = len(inputData)
    correction, data = inputData[:, 0], inputData[:, 1:3]
    return N, correction, data


def show_data(data, correction, N, theta, title):
    for g in G:
        xs = [data[i, 0] for i in range(N) if correction[i, 0] == g]
        ys = [data[i, 1] for i in range(N) if correction[i, 0] == g]
        options = {'linestyle': '', 'marker': 'o'}
        options.update(G[g])
        plt.plot(xs, ys, **options)

    def f(x):
        return (0.5 - theta[0] - theta[1] * x) / theta[2]

    min_x, max_x = min(data[:, 0]), max(data[:, 0])
    print(min_x, max_x)
    xs = np.arange(
        min_x,
        max_x,
        (max_x - min_x) / 100,
    )
    ys = [f(x) for x in xs]
    plt.plot(xs, ys)
    plt.title(title)
    plt.show()


# theta : (1, 3)
# x: (200, 3)
def fitted_p(x, theta):
    tmp = np.exp(np.dot(x, theta.T)) + EPSILON
    return tmp / (1 + tmp)


def newton_gradient(x, y, theta_0, iter_max=60):
    k = 0
    mm = 1
    theta = theta_0
    while k < iter_max and mm > 1e-08:
        p = fitted_p(x, theta)
        D = np.diagflat(np.multiply(p, 1 - p))
        first = (x.T).dot(D).dot(x).I
        sec = x.T.dot(y - p)
        theta = theta + first.dot(sec).T
        print(f'Iter {k+1}, Theta: {theta}')
        k += 1

    return theta


iter_max = 60
if __name__ == '__main__':
    n, y, x = openData("mixture2D")
    x = np.insert(x, 0, 1, axis=1)

    theta_0 = np.random.random_sample((1, 3))
    theta = newton_gradient(x, y, theta_0, iter_max=5)

    print(theta)
    show_data(x[:, 1:], y, n, theta.tolist()[0], "foo")