import pickle
import gzip
import numpy
import io
from sagemaker.amazon.common import write_numpy_to_dense_tensor

print("Extracting MNIST data set")
# Load the dataset
with gzip.open('/opt/ml/processing/input/mnist.pkl.gz', 'rb') as f:
    train_set, valid_set, test_set = pickle.load(f, encoding='latin1')

# process the data
# Convert the training data into the format required by the SageMaker KMeans algorithm
print("Writing training data")
with open('/opt/ml/processing/output_train/train_data', 'wb') as train_file:
    write_numpy_to_dense_tensor(train_file, train_set[0], train_set[1])

print("Writing test data")
with open('/opt/ml/processing/output_test/test_data', 'wb') as test_file:
    write_numpy_to_dense_tensor(test_file, test_set[0], test_set[1])

print("Writing validation data")
# Convert the valid data into the format required by the SageMaker KMeans algorithm
numpy.savetxt('/opt/ml/processing/output_valid/valid-data.csv', valid_set[0], delimiter=',', fmt='%g')