function [aff,Ngbs,md] = gaussknn(T,Y,sigma,KNN,flag)
% gaussknn  computes the gaussian affinity matrix of two data sets, and finds
%           KNN+1 nearest neighbors in the second data set for each data point 
%           in the first data set. Since usually T=Y, then gaussknn returns the
%           indices of each point itself and its KNN nearest neighbors.
% 
% Syntax:
%        [aff,Ngbs,md] = gaussknn(T,Y,sigma,KNN,flag)
% 
% Input:
%           T = NxD matrix containing data points in rows.
%           Y = MxD matrix containing data points in rows.
%       sigma = kernel width of heat kernel.
%         KNN = number of nearist neighbors of each data point.
%        flag = indicates whether to make result sparse(setting gaussian 
%               affinity to zero except for nearest neighbors), 
%               0(default) no, otherwise yes.
%        
% Output:
%         aff = NxM gaussian affinity between T and Y.
%        Ngbs = Nx(KNN+1) matrix, the i-th row of which contains the index 
%               of KNN nearest neighbors of T's i-th data point in Y.
%          md = Nx1 matrix containing the average distance of each data
%               point to its KNN+1 nearest neighbors.
%
% History:
%    Weiran Wang        February 24, 2009  recreated.

if ~exist('flag','var') || isempty(flag)
    flag = 0;
end

N = size(T,1); 
M = size(Y,1);

if flag==0 
    t = sum(T.^2,2); 
    y = sum(Y.^2,2)';
    sqd = max(repmat(t,1,M) + repmat(y,N,1) - 2*T*Y',0);

    [sqd2,I] = sort(sqd,2,'ascend');
    Ngbs = I(:,1:(KNN+1));
    md = mean(sqrt(sqd2(:,1:(KNN+1))),2);
    aff = exp(-sqd/(2*sigma^2));
else
    addpath ./MinMaxSelection;
    mem = 1; SEG = floor((mem*1024^3)/(4*M*8)/2);	% This will fit in mem GB RAM
    aff = sparse(N,M);
    Ngbs = zeros(N,KNN+1);
    md = zeros(N,1);
    y = sum(Y.^2,2)';
    i = 1;
    while (i <= N)   % Deal with segments separately and then assemble them.
        idx_seg = i:min(i+SEG-1,N);
        P = length(idx_seg);   % Length of the segment.
        
        T_seg = T(idx_seg,:);
        t_seg = sum(T_seg.^2,2);
        sqd = max(repmat(t_seg,1,M) + repmat(y,P,1) - 2*T_seg*Y',0);
	if exist('minkmex')==3
        	[sqd2,Ngbs_seg] = mink(sqd,KNN+1,2);	% About 4-5 times faster than sort for N=4000, KNN=30.
        	md(idx_seg) = mean(sqrt(sqd2),2);
	else
        	[sqd2,I] = sort(sqd,2,'ascend');
        	Ngbs_seg = I(:,1:(KNN+1));
        	md(idx_seg) = mean(sqrt(sqd2(:,1:(KNN+1))),2);
	end
       	Ngbs(idx_seg,:) = Ngbs_seg;
        ind = repmat((1:P)',1,KNN+1);
        aff(idx_seg,:) = sparse(ind(:),Ngbs_seg(:),...
            exp(-sqd(sub2ind([P,M],ind(:),Ngbs_seg(:)))/(2*sigma^2)),P,M);
        i = i+SEG;
    end
end
