function [testInfo, predictionInfo] = primalNearestNeighboursPredict(trainX, trainY, testX, modelInfo, params); 
%Predict using the primal kNN algorithm 

numTrainExamples = size(trainX, 1); 
numTestExamples = size(testX, 1); 
numLabels = size(trainY, 2); 

k = params.k; 

tic; 
e = ones(numTrainExamples, 1);
e2 = ones(numTestExamples, 1); 
diagTrainK = sum(trainX.^2, 2); 
diagTestK = sum(testX.^2, 2); 
distanceMatrix = diagTrainK*e2' + e*diagTestK' - 2*(trainX*testX');

maxEntry = max(max(distanceMatrix));
distances = zeros(numTestExamples, k); 
indices = zeros(numTestExamples, k); 
neighbours = zeros(numTestExamples, numLabels, k); 

for i=1:k
    [distances(:, i), indices(:, i)] = min(distanceMatrix, [], 1);
    
    distanceMatrix(sub2ind(size(distanceMatrix), indices(:, i), (1:numTestExamples)')) = maxEntry; 
    neighbours(:, :, i) = trainY(indices(:, i), :); 
end 

testInfo.predictedY = mode(neighbours, 3);
testInfo.rankings = sum(exp(-distances) .* squeeze(sum(neighbours, 2)), 2);
predictTime = toc; 
fprintf('Prediction time: %f s\n', predictTime); 

predictionInfo = struct;
predictionInfo.predictTime = predictTime;