import numpy as np
import pandas as pd

data_location = 'D:/Data/'
train_file = data_location + 'mnist_train.csv'
test_file = data_location + 'mnist_test.csv'

def get_number_mnist(to_image=False, one_hot=True):
    # Read the csv's
    test_values_csv = pd.read_csv(test_file)
    train_values_csv = pd.read_csv(train_file)

    # prepare the data
    test_amount = len(test_values_csv.index)
    test_values = test_values_csv.values[:, 1:] / 255
    test_labels = np.reshape(test_values_csv.values[:, 0:1], (test_amount,))

    train_amount = len(train_values_csv.index)
    train_values = train_values_csv.values[:, 1:] / 255
    train_labels = np.reshape(train_values_csv.values[:, 0:1], (train_amount,))

    # to 28x28
    if(to_image):
        train_values = train_values.reshape(train_values.shape[0], 28, 28, 1)
        test_values = test_values.reshape(test_values.shape[0], 28, 28, 1)


    if one_hot:
        return train_values, get_one_hot(train_labels, 10), test_values, get_one_hot(test_labels, 10)
    else:
        return train_values, train_labels, test_values, test_labels

def get_small_test_batch(amount, to_image=False, one_hot=True):
    # Read the csv's
    test_values_csv = pd.read_csv(test_file, nrows=amount)
    train_values_csv = pd.read_csv(train_file, nrows=amount)

    # prepare the data
    test_amount = len(test_values_csv.index)
    test_values = test_values_csv.values[:, 1:] / 255
    test_labels = np.reshape(test_values_csv.values[:, 0:1], (test_amount,))

    train_amount = len(train_values_csv.index)
    train_values = train_values_csv.values[:, 1:] / 255
    train_labels = np.reshape(train_values_csv.values[:, 0:1], (train_amount,))

    # to 28x28
    if (to_image):
        train_values = train_values.reshape(train_values.shape[0], 28, 28, 1)
        test_values = test_values.reshape(test_values.shape[0], 28, 28, 1)

    if one_hot:
        return train_values, get_one_hot(train_labels, 10), test_values, get_one_hot(test_labels, 10)
    else:
        return train_values, train_labels, test_values, test_labels

def get_one_hot(targets, nb_classes):
    res = np.eye(nb_classes)[np.array(targets).reshape(-1)]
    return res.reshape(list(targets.shape)+[nb_classes])