import cPickle, gzip
import numpy as np

LEARNING_RATE = 0.001
EPOCHS_COUNT = 15

def output_fixer(data_set):
    out_ds = []
    for i in data_set[1]:
        arr = [0] * 10
        arr[i] = 1
        out_ds.append(arr)
    return [
        data_set[0],
        np.array(out_ds)
    ]

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

# Load the dataset
with gzip.open('mnist.pkl.gz', 'rb') as f:
    train_set, valid_set, test_set = cPickle.load(f)
    train_set = output_fixer(train_set)
    valid_set = output_fixer(valid_set)
    test_set = output_fixer(test_set)

network = {
    'weights': [],
    'biases': np.random.standard_normal(len(train_set[1][0]))
}
for i in range(len(train_set[1][0])):
    network['weights'].append(np.random.standard_normal(len(train_set[0][0])))
network['weights'] = np.array(network['weights'])


for i in range(EPOCHS_COUNT):
    # print("Starting epoch " + str(i))
    for (X, Y) in zip(train_set[0], train_set[1]):
        computed_Y = network['weights'].dot(X) + network['biases']
        diff_Y = (Y - computed_Y) * LEARNING_RATE

        network['weights'] += np.transpose(np.matrix(diff_Y)).dot(np.matrix(X))
        network['biases'] += diff_Y

        # print("NETWORK-w")
        # print(network['weights'])
        # print("NETWORK-b")
        # print(network['biases'])
        # print("INPUT")
        # print(X)
        # print("EXPECTED OUTPUT")
        # print(Y)
        # print("OUTPUT")
        # print(computed_Y)
        # print("OUTPUT DIFF")
        # print(diff_Y)

    matches = 0
    for (X, Y) in zip(valid_set[0], valid_set[1]):
        computed_Y = network['weights'].dot(X) + network['biases']

        computed_Y_val = np.where(computed_Y == np.amax(computed_Y))[0][0]
        initial_Y_val = np.where(Y == np.amax(Y))[0][0]

        if computed_Y_val == initial_Y_val:
            matches += 1
    print("Epoch " + str(i) + " done. Accuracy: " + str((float(matches) / len(valid_set[0]))))