%%================================================================
%% Step 0a: Load data
%  Here we provide the code to load natural image data into x.
%  x will be a 784 * 600000 matrix, where the kth column x(:, k) corresponds to
%  the raw image data from the kth 12x12 image patch sampled.
%  You do not need to change the code below.

function [xHat, k, xZCAWhite, xPCAWhite] = pca(x, retained_variance_target)
epsilon = 1e-1; 

%%================================================================
%% Step 1a: Implement PCA to obtain xRot
%  Implement PCA to obtain xRot, the matrix in which the data is expressed
%  with respect to the eigenbasis of sigma, which is the matrix U.

fprintf('Calculate covariance matrix of size %d x %d\n', size(x, 1), size(x, 1));
sigma = x * x' / size(x, 2);

fprintf('Perform Singular Value Decomposition\n');
[U,S,V] = svd(sigma);

fprintf('Perform PCA transformation\n');
wpca = U' * diag(sqrt(1./(diag(S) + epsilon)));
wzca = U * wpca;
%%================================================================
%% Step 2: Find k, the number of components to retain
%  Write code to determine k, the number of components to retain in order
%  to retain at least 99% of the variance.

fprintf('Find K for variance target of %f\n', 100*retained_variance_target);
retained_variance = 1.0;
lambda = diag(S);
k = length(lambda);
while(retained_variance > retained_variance_target)
	retained_variance = sum(lambda(1:k))./sum(lambda);
	k = k - 1;
end
k = k

%%================================================================
%% Step 3: Implement PCA with dimension reduction
%  Now that you have found k, you can reduce the dimension of the data by
%  discarding the remaining dimensions. In this way, you can represent the
%  data in k dimensions instead of the original 144, which will save you
%  computational time when running learning algorithms on the reduced
%  representation.
% 
%  Following the dimension reduction, invert the PCA transformation to produce 
%  the matrix xHat, the dimension-reduced data with respect to the original basis.
%  Visualise the data and compare it to the raw data. You will observe that
%  there is little loss due to throwing away the principal components that
%  correspond to dimensions with low variation.

fprintf('Calculate xHat\n');
xrot = U' * x;
xHat = U(:, 1:k) * xrot(1:k, :);

%%================================================================
%% Step 4a: Implement PCA with whitening and regularisation
%  Implement PCA with whitening and regularisation to produce the matrix
%  xPCAWhite. 

%%% YOUR CODE HERE %%%
fprintf('Calculate xPCAWhite\n');
xPCAWhite = wpca * x;

%%================================================================
%% Step 5: Implement ZCA whitening
%  Now implement ZCA whitening to produce the matrix xZCAWhite. 
%  Visualise the data and compare it to the raw data. You should observe
%  that whitening results in, among other things, enhanced edges.

%%% YOUR CODE HERE %%%
fprintf('Calculate xZCAWhite\n');
xZCAWhite = wzca * x;