#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

/*
 * The MNIST Data can be downloaded as follows:
 * mkdir -p data/mnist/
 * cd data/mnist/
 * curl -O https://pjreddie.com/media/files/mnist_train.csv
 * curl -O https://pjreddie.com/media/files/mnist_test.csv
 */

source("nn/examples/mnist_2NN.dml") as mnist_2NN

# Read training data
data = read("mnist_data/mnist_train.csv", format="csv")
n = nrow(data)

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

# Scale images to [0,1], and one-hot encode the labels
images = images / 255.0
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
epochs = 5
[W_1, b_1, W_2, b_2, W_3, b_3] = mnist_2NN::train(X, y, X_val, y_val, epochs)

# Read test data
data = read("mnist_data/mnist_test.csv", format="csv")
n = nrow(data)

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

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

# Eval on test set
probs = mnist_2NN::predict(X_test, W_1, b_1, W_2, b_2, W_3, b_3)
[loss, accuracy] = mnist_2NN::eval(probs, y_test)

print("Test Accuracy: " + accuracy)