% load train data and labels
data = loadMNISTImages('train-images-idx3-ubyte');
labels = loadMNISTLabels('train-labels-idx1-ubyte');
% shift labels up one digit (matlab arrays start with 1)
labels = labels+1;

% load test data and labels
datatest = loadMNISTImages('t10k-images-idx3-ubyte');
labelstest = loadMNISTLabels('t10k-labels-idx1-ubyte');
labelstest = labelstest+1;

% whiten train and test data
data = whiten(data);
datatest = whiten(datatest);

% load the W matrix you want to use (extracted using trainrica.m)
load('0.mat')


% s is the representation of data in the new basis
s = W*data;
% rec recreates the original image by applying W*W'*data. Used only as 
% visual test
rec = W'*s;
% display_network(rec);

% same for test data
stest = W*datatest;
rectest = W'*stest;

% for each label, creates an array:
% (0,...,1,...,0) where the 1 is in the (i+1)-th position, where i
% is the label
target = zeros(10,60000);
for i=1:60000
    target((labels(i)),i)=1;
end
% train the classifier with the data in the new basis
net = trainSoftmaxLayer(s,target);

% same for test data
targetest = zeros(10,10000);
for i=1:10000
    targetest((labelstest(i)),i)=1;
end

% classifies test data and visualizes a confusion matrix with results
Y=net(stest);
plotconfusion(targetest,Y);