import numpy as np

from mnist import MNIST
from sklearn.tree import DecisionTreeClassifier  # decision tree
from sklearn.neural_network import MLPClassifier # neural networks
from sklearn.metrics import accuracy_score

def main():
    # This method reads in the training and testing data from MNIST
    X_train, y_train, X_test, y_test = load_dataset()
    print('Evaluating decision tree')
    evaluate_model(DecisionTreeClassifier(), X_train, y_train, X_test, y_test)
    print('Evaluating neural network')
    evaluate_model(MLPClassifier(), X_train, y_train, X_test, y_test)


# The stuff in here is more important to your machine learning task
def evaluate_model(model, X_train, y_train, X_test, y_test):
    """
    Evaluates the given model by training on X_train and y_train and
    testing it with X_test and y_test. Uses name to print better name
    """
    model.fit(X_train, y_train)

    train_predictions = model.predict(X_train)
    test_predictions = model.predict(X_test)

    print('   Train accuracy: ', accuracy_score(y_train, train_predictions))
    print('   Test accuracy : ', accuracy_score(y_test, test_predictions))


def load_dataset():
    """
    Loads the MNIST dataset from the data directory

    Returns:
    tuple of np.arrays (training data, training labels, test data, test labels)
    """
    mndata = MNIST('./data/')
    X_train, labels_train = map(np.array, mndata.load_training())
    X_test, labels_test = map(np.array, mndata.load_testing())
    X_train = X_train / 255.0
    X_test = X_test / 255.0
    return X_train, labels_train, X_test, labels_test


# Python magic to get main() to run
if __name__ == '__main__':
    main()