% MEDIAN-DISTANCE VERSION
% Semi-supervised estimation of class-prior using LSDD.
% Class priors are denoted by \xi 
% Xte is the unlabeled samples from the test dataset
% Xtr is the labeled samples from the training dataset. This should be a 
%  cell array with a cell for each class.
% Note: At this moment the assumption is that each class in the training dataset has the 
%  same amount of samples.
% xi_list -> A list of xi values that should be evaluated
function [xi_best, xi_list, LSDE] =  LSDDPriorEstMedian(Xte, Xtr, xi_list)

	% set the xi list if not given
	if (nargin<3 || isempty(xi_list))
		xi_list = linspace(0, 1, 200);
	end

	% check that the sizes of the cell array are the same
	Xall = [Xte];
	nclasses = length(Xtr);
	for i=1:nclasses
		ntr(i) = size(Xtr{i}, 2);
		Xall = [Xall Xtr{i}];
	end
	if (min(ntr)~=max(ntr))
		error('Different sizes for training data not (yet) supported');
	end
	ntr = ntr(1);
	nte = size(Xte, 2);
	d = size(Xte, 1);

	% select the basis functions
	ntotal = size(Xall, 2);
	b = min(300, ntotal);
	center_index = randperm(ntotal);

	C = Xall(:, center_index(1:b));
	% estimate the median distance
	[~, median_dist] = CalculateDist2(C, Xall);
	
	% get the individual squared distances
	XteC_dist2 = CalculateDist2(C, Xte);
	CC_dist2 = CalculateDist2(C, C);
	for i=1:nclasses
		XtrC_dist2{i} = CalculateDist2(C, Xtr{i});
	end

	% split the data into different folds. n1_cv stores the number of samples per fold
	fold=5;
	cv_fold=[1:fold];
	cv_split_tr = floor([0:ntr-1]*fold./ntr)+1;
	cv_split_te = floor([0:nte-1]*fold./nte)+1;
	cv_index_tr = cv_split_tr(randperm(ntr));
	cv_index_te = cv_split_te(randperm(nte));
	for k=cv_fold
	  nte_cv(k)=sum(cv_index_tr==k);
	  ntr_cv(k)=sum(cv_index_te==k);
	end % k

	% set the hyper-parameter lists
	sigma_list=median_dist*linspace(0.75,3,20); % Candidates of Gaussian width
	lambda_list=logspace(-3,1,10); % Candidates of regularization parameter

	% set the CV scores to 0
	score_cv = zeros(length(sigma_list), length(lambda_list), length(cv_fold), length(xi_list));

	% perform cross-validation
	for sigma_index=1:length(sigma_list)
		sigma=sigma_list(sigma_index);

		% the H matrix only needs to be calculated once per sigma value
		H=(sqrt(pi)*sigma)^d*exp(-CC_dist2/(4*sigma^2));

		% calculate the h1 and h2 matrices for each fold
		for k=cv_fold
			hte_cv(:, k) = sum(exp(-XteC_dist2(:,cv_index_te==k)/(2*sigma^2)),2);
			for j = 1:nclasses
				G_cv(:, j, k) = sum(exp(-XtrC_dist2{j}(:,cv_index_tr==k)/(2*sigma^2)),2);
			end
		end

		for lambda_index=1:length(lambda_list)
			lambda=lambda_list(lambda_index);

			for k=cv_fold

				% first part - the training fold
				htrain_p1 = sum(hte_cv(:,cv_fold~=k),2)/sum(nte_cv(cv_fold~=k));
				htrain_p2 = sum(G_cv(:, :, cv_fold~=k),3)/sum(ntr_cv(cv_fold~=k));

				% second part - the test fold
				htest_p1 = sum(hte_cv(:,cv_fold==k),2)/sum(nte_cv(cv_fold==k));
				htest_p2 = sum(G_cv(:, :, cv_fold==k),3)/sum(ntr_cv(cv_fold==k));

				% Calculate the solution for all xi
				Hlambda = H+lambda*eye(b);
				%HlamIinv = inv(H+lambda*eye(b));
				
				a1 = linsolve(Hlambda, htrain_p1);
				a2 = linsolve(Hlambda, htrain_p2(:, 1));
				a3 = linsolve(Hlambda, htrain_p2(:, 2));

				for xi_index = 1:length(xi_list)
					xi = xi_list(xi_index);					

					% calculate the theta values
					htest = htest_p1 - htest_p2*[xi; 1-xi];

					thetah = a1 - xi*a2 - (1-xi)*a3;

					% calculate the score
					score = thetah'*H*thetah -2*thetah'*htest;
					score_cv(sigma_index, lambda_index, k, xi_index) = score;
				end

			end % k
		end
	end % sigma

	%
	sigma = []; lambda = [];

	% calculate the score for each candidate xi
	for xi_index=1:length(xi_list)
		xi = xi_list(xi_index);

		% get the best lambda and signa value
		[sigma_index, lambda_index] = min2(mean(score_cv(:, :, :, xi_index), 3));
		sigma = sigma_list(sigma_index);
		lambda = lambda_list(lambda_index);

		% step 1: calculate the matrices
		H=(sqrt(pi)*sigma)^d*exp(-CC_dist2/(4*sigma^2));
		hte = mean(exp(-XteC_dist2/(2*sigma^2)),2);

		for j = 1:nclasses
			G(:, j) = mean(exp(-XtrC_dist2{j}/(2*sigma^2)),2); 
		end

		% step 2: calculate theta value
		h = hte - G*[xi; 1-xi];
		thetah=mylinsolve(H+lambda*eye(b), h);

		% step 3: calcualte the divergence
		LSDE(xi_index) = 2*thetah'*h - thetah'*H*thetah;
	end

	% select the best xi
	[~, xi_idx] = min(LSDE);
	xi_best = xi_list(xi_idx);
end

% Calculates the squared distance between the column vectors in matrix X and the column vectors in matrix Y.
% X -> input matrix of dxn
% Y -> matrix of centers, dxm
% XC_dist2 -> output matrix of size nxm
% med_dist -> the median distance between the vectors
function [XC_dist2, med_dist] = CalculateDist2(X, Y)

	[d nx]=size(X);
	[d nc]=size(Y);
	X2=sum(X.^2,1);
	Y2=sum(Y.^2,1);

	% calculate the square
	XC_dist2 = repmat(Y2, nx, 1)+repmat(X2',1,nc)-2*X'*Y;

	% calculate the median distance
	if (nargout>1)
		med_dist = median(sqrt(XC_dist2(:)));
	end
end

% get the minimum value of a matrix and the two indices
% idx1 -> the row index
% idx2 -> the column indec
function [idx1, idx2, mval] = min2(A)
	
	fbest = inf;
	idx1 = 1;
	idx2 = 1;
	
	for i=1:size(A, 1)
		for j=1:size(A, 2)
			if (A(i,j)<fbest)
				idx1 = i;
				idx2 = j;
				fbest = A(i,j);
			end
		end
	end
	
	mval = A(idx1, idx2);	
end


function x=mylinsolve(A,b)
	%Solve Ax=b (A and b are both dense)

	[x,flag]=pcg(A,b);
	%R=chol(A); x=R\(R'\b);
	%x=A\b;
	%x=gmres(A,b);
	%x=bicg(A,b);
	%x=bicgstab(A,b);
	%x=cgs(A,b);
	%x=minres(A,b);
	%x=qmr(A,b);
	%x=symmlq(A,b);
	%opts.SYM=true;x=linsolve(A,b,opts);
	%x=inv(A)*b;
end

