# loocv_starter.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
# package MLDatasets. 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("MLDatasets");
#
# Step 2: Each time you work on this assignment, assuming MLDatasets
# is installed, you'll want to include this file. To do so, in your
# julia environment, use this call:
#
#   include("loocv_starter.jl")
#
# Then see the methods below to load the data. The first time you run
# load_all_data() you'll be asked if you want to download a dataset
# from [a long website name]. Type 'y' to download it upon prompting.
#
# Further remarks: if this code fails to work, see the other file
# loocv_starter_csv.jl.
#
# Author: John Duchi <jduchi@stanford.edu> and ENGR-108 team.
# Date: November 2024

using MLDatasets
using LinearAlgebra

# (A_train, y_train, A_test, y_test) = load_all_data()
#
# 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.
function load_all_data()
  # Loads the training data and test data of 28-by-28 images
  trainset = MNIST(:train);
  testset = MNIST(:test);
  (im_dim_1, im_dim_2, N_train) = size(trainset.features);
  N_test = size(testset.features, 3);

  # First, find the features that are always zero in each example.
  pixel_sums = sum(trainset.features, dims = 3);
  all_zeros = (pixel_sums .== 0)[:];

  dim = im_dim_1 * im_dim_2;

  # Make training data matrix of raw features, size N_train-by-784, then
  # remove pixels corresponding to all-zeros.
  A_train = reshape(trainset.features, (dim, N_train))';
  A_train = A_train[:, .! all_zeros];
  y_train = trainset.targets;
  
  # Make test data matrix of raw features, size N_test-by-784, then
  # remove pixels corresponding to all-zeros.
  A_test = reshape(testset.features, (dim, N_test))';
  A_test = A_test[:, .! all_zeros];
  y_test = testset.targets;
  
  return (A_train, y_train, A_test, y_test);
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