import utils.VisualizeNN as VisNN
from sklearn.neural_network import MLPClassifier
import numpy as np

if __name__ == "__main__":
    training_set_inputs = np.array([[0, 0, 1], [0, 1, 1], [1, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 1], [0, 0, 0]])
    training_set_outputs = np.array([[0, 1, 1, 1, 1, 0, 0]]).T
    X = training_set_inputs
    y = training_set_outputs

    classifier = MLPClassifier(hidden_layer_sizes=(4,), alpha=0.01, tol=0.001, random_state=1)
    classifier.fit(X, y.ravel())

    network_structure = np.hstack(([X.shape[1]], np.asarray(classifier.hidden_layer_sizes), [y.shape[1]]))
    print(network_structure)

    # Draw the Neural Network with weights
    network = VisNN.DrawNN(network_structure, classifier.coefs_)

    network.draw()

    # Draw the Neural Network without weights
    network = VisNN.DrawNN(network_structure)
    network.draw()
