# loocv_starter_csv.jl
#
# Starter file for Leave One Out Cross Validation on the MNIST
# dataset, which provides the functions load_all_data() and
# convert_to_binary().
#
# Remarks: the first time you use this file, you may need to add the
# packages CSV and DataFrames.  To do this, you'll need to use the
# following steps:
#
# Step 1: run the following two lines of code. They may take a little while.
# You will only need to do step 1 once.
#
#   using Pkg;
#   Pkg.add("CSV");
#   Pkg.add("DataFrames");
#
# Step 2: Download the data file at the course website in
#
#   https://web.stanford.edu/class/engr108/data/mnist.zip
#
# Unzip that file into the directory in which you'll do your work.
#
# Step 3: Each time you work on this assignment, assuming CSV and
# DataFrames are installed, you'll want to include this file. To do
# so, in your julia environment, use this call:
#
#   include("loocv_starter_csv.jl")
#
# Further remarks: if this code fails to work, see the other file
# loocv_starter.jl.
#
# Author: John Duchi <jduchi@stanford.edu> and ENGR-108 team.
# Date: November 2024

using CSV;
using DataFrames;

# (A_train, y_train, A_test, y_test) = load_all_data(path_to_files)
#
# Loads the MNIST digit-recognition training data and training data
# into two matrices, denoted A_train and A_test, where A_train is of
# size N_train-by-d and A_test is of size N_test-by-d, where d is the
# dimension of the vectorized images in the MNIST dataset. The MNIST
# images are originally 28-by-28 grayscale images, where pixels take
# values in [0, 1], (0 being black, 1 white) but because some pixels
# are always zero in the training data, we eliminate those.
#
# The i-th row A_train[i, :] of the training matrix (respectively test)
# is thus a transposed d-vector containing the pixel values for image i.
#
# The vectors y_train and y_test are integer vectors, where the i-th
# entry y_train[i] is a number in {0, 1, ..., 9} representing the class.
#
# Note: the argument 'path_to_files' should be a string with the path
# to the directory where mnist_train.csv and mnist_test.csv are
# stored. The path should be an absolute path.
function load_all_data(path_to_files::String = "./")
  # load training data
  yAtrain = Matrix(DataFrame(CSV.File(string(path_to_files,
                                             "./mnist_train.csv"))));
  # Get training labels
  ytrain = yAtrain[:, 1];
  # Normalize data
  Atrain = yAtrain[:, 2:end] / 256;
  # First, find the features that are always zero in each example.
  pixel_sums = sum(Atrain, dims = 1);
  all_zeros = (pixel_sums .== 0)[:];
  Atrain = Atrain[:, .! all_zeros];

  # Load test data
  yAtest = Matrix(DataFrame(CSV.File(string(path_to_files,
                                            "./mnist_test.csv"))));
  # Get test labels
  ytest = yAtest[:, 1];
  # Make test data matrix of raw features, size N_test-by-784, then
  # remove pixels corresponding to all-zeros. Also normalize to be in [0, 1].
  Atest = yAtest[:, 2:end] / 256;
  Atest = Atest[:, .! all_zeros];
  return (Atrain, ytrain, Atest, ytest);
end

# y_binary = convert_to_binary(y, target)
#
# Converts the given vector y of labels in {0, 1, ..., 9} into a
# binary vector of {-1, 1} entries, where the positive labels
# correspond to the target class (input target), and the negatives
# correspond to anything else. The input target defaults to 0.
function convert_to_binary(y::Vector{Int64}, target::Int64 = 0)
  if (target < 0 || target > 9)
    error("Attempting to label images of class ", target,
          " not in 0, 1, 2, ..., 9.")
  end
  y_bin = (y .== 0);
  y_bin = (2 * y_bin .- 1);
  return y_bin;
end