@author: alex
"""

import numpy as np

train_data = np.array([[0,0],[0,1],[1,0],[1,1]]).T
print(train_data)
labels = np.array([0,1,1,0])

def sigmoid(z,deriv = False):
    sig =  1/(1+np.exp(-z))
    if deriv == True:
        return np.multiply(sig,1-sig)
    return sig


w1 = np.random.randn(2,2)*0.01
b1 = np.zeros((2,1))

w2 = np.random.randn(1,2)*0.01
b2 = np.zeros((1,1))

lr = 5.5

for i in range(5000):
    # Shuffle data samples
    permutation = np.random.permutation([0,1,2,3])
    train_data = train_data[:,permutation]
    labels = labels[permutation]

    z1 = np.dot(w1,train_data) + b1
    a1 = sigmoid(z1)
    z2 = np.dot(w2,a1) + b2
    a2 = sigmoid(z2)


    cost = np.dot(labels,np.log(a2).T) + np.dot(1-labels,np.log(1-a2).T)
    cost = cost*(-1/4)

    # cost = ((1/4) * np.dot(labels,np.log(a2).T) + np.dot(1-labels,np.log(1-a2).T))
    cost = np.squeeze(cost)

    #dz2 = np.multiply(da2,sigmoid(z2,deriv = True))
    dz2 = a2-labels
    dw2 = (1/4) * np.dot(dz2,a1.T)
    db2 = (1/4) * np.sum(dz2,axis=1,keepdims = True)
    dz1 = np.multiply(np.dot(w2.T, dz2),sigmoid(z1,deriv = True))
    dw1 = (1/4) * np.dot(dz1,train_data.T)
    db1 = (1/4) * np.sum(dz1,axis=1,keepdims = True) #backprop

    w1 = w1 - lr*dw1
    w2 = w2 - lr*dw2
    b1 = b1 - lr*db1
    b2 = b2 - lr*db2 #update params
    print(np.abs(cost),'------',str(i))


a1 = sigmoid(np.dot(w1, train_data) + b1)
a2 = sigmoid(np.dot(w2, a1) + b2)
print("Correct labels:", labels)
print("Predicted labels:", np.round(a2[0]))