% Code to learn the face pose recognition model
% It uses matconvnet - http://www.vlfeat.org/matconvnet/
% 
% Author: Aravindh Mahendran
% aravindh.mahendran@new.ox.ac.uk

clear all;
close all;

run C:\Users\aravindh\Software\matconvnet\matlab\vl_setupnn.m

load('facedata.mat'); % loads X for the face images, Y for their ground truth labels
% X is 30 x 32 x 1 x 277 dimensions and
% Y is 1 x 1 x 4 x 277 dimensions
% X_test is 30 x 32 x 1 x 139
% Y_test is 139 x 1 dimensions containing integers between 1,2,3,4

% hyperparameters for the training
lr = 0.03;
numIterations = 600;
numHiddenUnits = 3;

% Initialize the parameters for the network
% Matconvnet is a convolutional neural network toolbox
% The inner product layers are designed by making kernels with width and
% height = 1. These kernels operate on data which is also width and height
% = 1
f1 = 0.05 * randn(1, 1, 30*32, numHiddenUnits, 'single');
b1 = zeros(numHiddenUnits, 1, 'single');
f3 = 0.05 * randn(1, 1, numHiddenUnits, size(Y, 3), 'single');
b3 = zeros(size(Y,3), 1, 'single');

loss = zeros(numIterations, 1);
for i=1:numIterations
    loss(i) = 0;
    for j=1:size(X, 4)
        x_cur = reshape(X(:,:,1,j), [1, 1, 30*32, 1]);
        y_cur = Y(1,1,:,j);
        
        % forward pass
        h1 = vl_nnconv(x_cur, f1, b1);
        h2 = vl_nnsigmoid(h1);
        h3 = vl_nnconv(h2, f3, b3);
        h4 = vl_nnsigmoid(h3);
        [loss_cur, dzdh4] = vl_nnsqeuclidean(h4, y_cur, [], single(1));
        loss(i) = loss(i) + loss_cur;
        
        % backward pass
        dzdh3 = vl_nnsigmoid(h3, dzdh4);
        [dzdh2, dzdf3, dzdb3] = vl_nnconv(h2, f3, b3, dzdh3);
        dzdh1 = vl_nnsigmoid(h1, dzdh2);
        [dzdx_cur, dzdf1, dzdb1] = vl_nnconv(x_cur, f1, b1, dzdh1);
        
        % Update the parameters
        f1 = f1 - lr * dzdf1; b1  = b1 - lr * dzdb1;
        f3 = f3 - lr * dzdf3; b3 = b3 - lr * dzdb3;
    end
    if(mod(i, 10) == 1)
        fprintf(1, 'Iteration %d\n', i);
        figure(1);
        plot(loss(1:i));
        xlabel('Number of Iterations');
        ylabel('Loss');
        
        figure(2);
        subplot(1,3,1); imagesc(reshape(f1(:,:,:,1), [30, 32, 1])); axis image; colormap gray;
        subplot(1,3,2); imagesc(reshape(f1(:,:,:,2), [30, 32, 1])); axis image; colormap gray;
        subplot(1,3,3); imagesc(reshape(f1(:,:,:,3), [30, 32, 1])); axis image; colormap gray;
        
        drawnow;
        
        pause(0.5);
    end
end

% Visualize the learnt network
figure;
subplot(1,3,1); imagesc(reshape(f1(:,:,:,1), [30, 32, 1])); axis image; colormap gray;
subplot(1,3,2); imagesc(reshape(f1(:,:,:,2), [30, 32, 1])); axis image; colormap gray;
subplot(1,3,3); imagesc(reshape(f1(:,:,:,3), [30, 32, 1])); axis image; colormap gray;

squeeze(f3)

% Let's test it on the test set.
numTestData = size(X_test, 4);
X_test = reshape(X_test, [1, 1, 30*32, numTestData]);
h1 = vl_nnconv(X_test, f1, b1);
h2 = vl_nnsigmoid(h1);
h3 = vl_nnconv(h2, f3, b3);
h4 = vl_nnsigmoid(h3);
correct = 0;
for i=1:numTestData
    [~, result_cur] = max(squeeze(h4(1,1,:,i)));
    if(result_cur == Y_test(i))
        correct = correct + 1;
    end
end
accuracy = correct / numTestData;