%% Reliable Effective linear Learning

%%
% Project by Edward Chao (ID: 20474367), Prarthana Bhattacharyya (ID:
% 20714026) and Hari Govind V.K. (ID : 20743574).
%
% Based on paper by Agarwal et al [1].
% 

%%
% *Problem Formulation*
% 
% Much of the work in machine learning focuses on linear learning problems
% of the form:
% 
% $$ min_{\textbf{w} \in R^d} {\sum_{i=1}^{n}{l(\textbf{w}^T\textbf{x}_i; y_i) + \lambda R(\textbf{w})}}, $$
% 
% where $\textbf{x}_i$ is the feature vector of the $i$-th example, $y_i$ is the
% label, $\textbf{w}$ is the linear predictor, $l$ is a loss function and
% $R$ is a regularizer. 
%
% In order to optimize the objective function above when the number of
% examples $n$ is too large, only online or only batch learning algorithms 

% Clear all data
clc; clear all; close all;

%Setup parameters
data_length = 3072;
num_classes = 10;
svm_buffer = 1;
step_size = 0.001;
num_iters = 15;
alpha=0.3;
%Load CIFAR-10 data
load('data_batch_1.mat');
data2 = load('data_batch_2.mat');
data3 = load('data_batch_3.mat');
data4 = load('data_batch_4.mat');
data5 = load('data_batch_5.mat');
data = cat(1,data,data2.data,data3.data,data4.data,data5.data);

labels = cat(1,labels,data2.labels,data3.labels,data4.labels,data5.labels);

%% Load CIFAR-10 test data
test_data = load('test_batch.mat');
%Might want to make these random values smaller
random_W = rand(num_classes, data_length);

%% Only online learning
tic
[online_W1, ll_online1] = onlineDescent(random_W,data_length,num_classes,svm_buffer,step_size*0.1,data,labels,alpha);
disp("Online 1 done")
toc

[online_W2, ll_online2] = onlineDescent(online_W1,data_length,num_classes,svm_buffer,step_size*0.1,data,labels,alpha);
disp("Online 2 done")

[online_W3, ll_online3] = onlineDescent(online_W2,data_length,num_classes,svm_buffer,step_size*0.1,data,labels,alpha);
disp("Online 3 done")

%% Online and batch learning combined

[combined_W, ll_combined, al_combined] = batchDescent(online_W3,data_length,num_classes,svm_buffer,step_size,data,labels,num_iters,alpha,test_data.data,test_data.labels);
disp("Combined done")
toc

%% Only batch learning
tic
[batch_W, ll_batch, al_batch] = batchDescent(random_W,data_length,num_classes,svm_buffer,step_size,data,labels,num_iters,alpha,test_data.data,test_data.labels);
disp("Batch only done")
toc

%% Test the weights

% Test batch results
data = im2double(data);
test_scores = batch_W * data';

%Find predicted classes
[M, I] = max(test_scores,[],1);

%Compare predicted classes to ground truth
evaluation = I-(double(labels')+1);
idx = evaluation==0;
num_corr_batch = sum(idx);

%Calculate accuracy
corr_perc_batch = num_corr_batch / length(labels');
fprintf("Batch learning accuracy: %f\n",corr_perc_batch);

test_scores = online_W3 * data';
[M, I] = max(test_scores,[],1);

evaluation = I-(double(labels')+1);
idx = evaluation==0;
num_corr_online = sum(idx);
corr_perc_online = num_corr_online / length(labels');
fprintf("Online learning accuracy: %f\n",corr_perc_online);

test_scores = combined_W * data';
[M, I] = max(test_scores,[],1);

evaluation = I-(double(labels')+1);
idx = evaluation==0;
num_corr_combined = sum(idx);
corr_perc_combined = num_corr_combined / length(labels');
fprintf("Online + batch accuracy: %f\n",corr_perc_combined);

%% Loss over time

figure
hold on

p = polyfit(1:20, ll_batch-min(ll_combined), 4);
p_fun = polyval(p,1:0.1:20);
plot (1:0.1:20, p_fun)
title('Figure 1: Batch vs. Online+Batch Training Suboptimality')
plot(1:20, ll_combined - min(ll_combined))
legend('No online pass', 'With online pass')
xlabel('Iterations')
ylabel('Suboptimality')
hold off

figure
hold on
plot(al_batch .* 100)
title('Figure 2: Batch vs. Online+Batch Test Accuracy')
plot(al_combined .* 100)
legend('No online pass', 'With online pass')
xlabel('Iterations')
ylabel('Accuracy on test set')

function eval = eval_W(W, data, labels)
% Test batch results
data = im2double(data);
test_scores = W * data';

%Find predicted classes
[M, I] = max(test_scores,[],1);

%Compare predicted classes to ground truth
evaluation = I-(double(labels')+1);
idx = evaluation==0;
num_corr = sum(idx);

%Calculate accuracy
eval = num_corr / length(labels');
end

function [W, loss_log] = onlineDescent(initial_W,dl, nc, svmbuff, ss, dataa, label,alpha)
%data_length = 3072;
%num_classes = 10;
%svm_buffer = 1;
%step_size = 0.0001;

data_length = dl;
num_classes = nc;
svm_buffer = svmbuff;
step_size = ss;
data = dataa;
labels = label;

W = initial_W;

%Load CIFAR-10 data
%load('data_batch_1.mat');
num_imgs = length(data);
loss_log = zeros(1,num_imgs);

for i = 1:num_imgs
    %Calculate score of image i
    img = im2double(data(i,:));
    score_i = W*img';
    
    %Get the label of image i
    y = double(labels(i));
    
    tmp = score_i' - score_i(y+1) + svm_buffer;
    maxed = max(0,tmp);
    loss = sum(tmp)-svm_buffer + alpha*norm(W); %Subtract loss contributed by the correct image
    loss_log(i) = loss;
    
    %Compute gradient
    bin_mask = find(maxed);
    bin_array = zeros(1,num_classes);
    bin_array(bin_mask) = 1;
    bin_array(y+1) = 0;
    
    gradient = bin_array' * img + alpha*2*W;
    
    W = W - step_size * gradient;
end
end

function [W, loss_log, acc_log] = batchDescent(initial_W, dl, nc, svmbuff, ss, dataa, label, num_iters,alpha, test_data, test_label)
% data_length = 3072;
% num_classes = 10;
% svm_buffer = 1;
% step_size = 0.01;
data_length = dl;
num_classes = nc;
svm_buffer = svmbuff;
step_size = ss;
data = dataa;
labels = label;
num_iterations = num_iters;

%Might want to make these random values smaller
%W = rand(num_classes, data_length);
W = initial_W;

%Load CIFAR-10 data
%load('data_batch_1.mat');
batch_size = 16;

%num_imgs = length(data);
total_size = length(data);
num_imgs = batch_size;
loss_log = zeros(1,num_iterations);
acc_log = zeros(1,num_iterations);

for j = 1:num_iterations
    total_loss = 0;
    total_gradient = zeros(num_classes, data_length);
    
    %for i = 1:num_imgs
    
    img_num = ceil(total_size.*rand(batch_size,1));
    %Calculate score of image i
    % img = im2double(data(i,:));
    for k=1:length(img_num)
        
        img = im2double(data(img_num(k),:));
        score_i = W*img';
        
        %Get the label of image i
        y = double(labels(img_num(k)));
        
        tmp = score_i' - score_i(y+1) + svm_buffer;
        maxed = max(0,tmp);
        total_loss = total_loss + sum(maxed)-svm_buffer + alpha*norm(W); %Subtract loss contributed by the correct image
        
        %Accumulate gradient
        bin_mask = find(maxed);
        bin_array = zeros(1,num_classes);
        bin_array(bin_mask) = 1;
        bin_array(y+1) = 0;
        total_gradient = total_gradient + bin_array' * img + alpha*2*W;
    end
    average_loss = total_loss / num_imgs;
    average_gradient = total_gradient / num_imgs;
    W = W - step_size * average_gradient;
    loss_log(j) = average_loss;
    acc_log(j) = eval_W(W, test_data, test_label);
end

end
