# Create a SystemML MLContext object
from systemml import MLContext, dml
ml = MLContext(sc)

"""
Explanation: Quick Setup
End of explanation
"""


%%sh
mkdir -p data/mnist/
cd data/mnist/
curl -O http://pjreddie.com/media/files/mnist_train.csv
curl -O http://pjreddie.com/media/files/mnist_test.csv

"""
Explanation: Download Data - MNIST
The MNIST dataset contains labeled images of handwritten digits, where each example is a 28x28 pixel image of grayscale values in the range [0,255] stretched out as 784 pixels, and each label is one of 10 possible digits in [0,9].  Here, we download 60,000 training examples, and 10,000 test examples, where the format is "label, pixel_1, pixel_2, ..., pixel_n".
End of explanation
"""


script_string = """
source("mnist_lenet.dml") as mnist_lenet

# Read training data
data = read($data, format="csv")
n = nrow(data)

# Extract images and labels
images = data[,2:ncol(data)]
labels = data[,1]

# Scale images to [-1,1], and one-hot encode the labels
images = (images / 255.0) * 2 - 1
labels = table(seq(1, n), labels+1, n, 10)

# Split into training (55,000 examples) and validation (5,000 examples)
X = images[5001:nrow(images),]
X_val = images[1:5000,]
y = labels[5001:nrow(images),]
y_val = labels[1:5000,]

# Train
[W1, b1, W2, b2, W3, b3, W4, b4] = mnist_lenet::train(X, y, X_val, y_val, C, Hin, Win)
"""
script = (dml(script_string).input("$data", "data/mnist/mnist_train.csv")
                            .input(C=1, Hin=28, Win=28)
                            .output("W1", "b1", "W2", "b2", "W3", "b3", "W4", "b4"))
W1, b1, W2, b2, W3, b3, W4, b4 = (ml.execute(script)
                                    .get("W1", "b1", "W2", "b2", "W3", "b3", "W4", "b4"))

"""
Explanation: SystemML "LeNet" Neural Network
1. Train
End of explanation
"""


script_string = """
source("mnist_lenet.dml") as mnist_lenet

# Read test data
data = read($data, format="csv")
n = nrow(data)

# Extract images and labels
X_test = data[,2:ncol(data)]
y_test = data[,1]

# Scale images to [-1,1], and one-hot encode the labels
X_test = (X_test / 255.0) * 2 - 1
y_test = table(seq(1, n), y_test+1, n, 10)

# Eval on test set
probs = mnist_lenet::predict(X_test, C, Hin, Win, W1, b1, W2, b2, W3, b3, W4, b4)
[loss, accuracy] = mnist_lenet::eval(probs, y_test)

print("Test Accuracy: " + accuracy)
"""
script = dml(script_string).input(**{"$data": "data/mnist/mnist_train.csv",
                                     "C": 1, "Hin": 28, "Win": 28,
                                     "W1": W1, "b1": b1,
                                     "W2": W2, "b2": b2,
                                     "W3": W3, "b3": b3,
                                     "W4": W4, "b4": b4})
ml.execute(script)

"""
Explanation: 2. Compute Test Accuracy
End of explanation
"""


W1_df = W1.toDF()
b1_df = b1.toDF()
W2_df = W2.toDF()
b2_df = b2.toDF()
W3_df = W3.toDF()
b3_df = b3.toDF()
W4_df = W4.toDF()
b4_df = b4.toDF()
W1_df, b1_df, W2_df, b2_df, W3_df, b3_df, W4_df, b4_df

"""
Explanation: 3. Extract Model Into Spark DataFrames For Future Use
End of explanation
"""


