%
% Estimates the class prior in a semi-supervised fashion by matching
%  the divergences
% The best prior is selected based on grid search and the matrices used
%   for cross-validation is pre-computed.
%
% USAGE:
% 	[theta_est, theta_list, PD] = pe_prior_est_grid(x, y, xm, K)
%
% INPUTS:
% 	x:	The (labelled) training data, d x ntr, where ntr is the number of training samples
% 	y:	The class labels of the vectors in x
% 	xm:	The test samples
% 	K:	The number of folds for cross validation (default 5)
%
% OUTPUTS:
% 	theta_est:	The estimated theta value which minimizes the Pearson divergence.
% 	theta_list:	The list of theta values for which the Pearson divergence is estimated
%	PD:		The estimated Pearson divergence for each value in theta_list
%
% Written by Christo du Plessis (christo@sg.cs.titech.ac.jp)
% Updated: 2012/08/25
function [theta_est, theta_list, PD] = pe_prior_est_grid(x, y, xm, K)

	theta_est = [];		 

	% get the sizes
	d = size(x, 1);
	n = size(x, 2);
	nm = size(xm, 2);

	% get the number of classes
	class_labels = sort(unique(y));
	c = length(class_labels);
	
	% check that there are only two classes
	if (length(class_labels)~=2)
		error('Grid search is only implemented for two classes');
	end
	
	% Default to 5 folds if no folds are specified
	if (nargin<4)
		K = 5;
	end 

	% Set the hyperparameters to be evaluated during cross-validation
	lambda_list = linspace(1E-2, 5, 30);
	sigma_med_dist = compMedDist([x xm]);
	sigma_list = linspace(1, 50, 30)*sigma_med_dist;

	% For each CV parameter, calculate the matrix
	for sigma_index = 1:length(sigma_list)
		sigma = sigma_list(sigma_index);
		
		fprintf('\t sigma %d %g\n', sigma_index, sigma);

		% Calculating the CV matrices is the only part that is (somewhat) time-inefficient
		%  This is remedied by calculating the matrices for each lambda value (done below)
		% Further speedups are possible, but not worth the effort (IMHO).
		if (0)

			for lambda_index = 1:length(lambda_list)
				lambda = lambda_list(lambda_index);
				
				fprintf('sigma_index %d, lambda_index %d\n', sigma_index, lambda_index);
				
				% Calculate the cross-validation matrix for this specific sigma and
				%  lambda value

				[A_CV{sigma_index, lambda_index}] = calc_CV_matrix(x, y, xm, K, sigma, lambda);
			end

		end
		
		
		% Calculate the Pearson divergence matrices for specific lambda and sigma values
		for lambda_index = 1:length(lambda_list)
			lambda = lambda_list(lambda_index);
				
			% Calculate the PD matrix 
			[A_PD{sigma_index, lambda_index}] = calc_PD_matrix(x, y, xm, sigma, lambda);
		end

		% Precompute the cross-validation matrices
		[A_CV_list] = calc_CV_matrix_lambda_list(x, y, xm, K, sigma, lambda_list);
		for lambda_index = 1:length(lambda_list)
			A_CV{sigma_index, lambda_index} = A_CV_list{lambda_index};
		end
	end

	
	theta_list = linspace(0, 1, 40);
	
	% search for an arbitrary number of values
	best_PD = inf;
	best_theta = [];
	
	for i=1:length(theta_list);

		% set the theta value
		theta(1, 1) = theta_list(i);
		theta(2, 1) = 1 - theta_list(i);
		
		% Calculate the PD score
		PD_score = estimate_PD(theta, A_CV, A_PD, sigma_list, lambda_list);
		
		PD(i) = PD_score;
	end

	[best_PD, idx] = min(PD);
	theta_est = theta_list(idx);
end

% Calculate the PD for a specific theta value
function [pd, sigma_chosen, lambda_chosen] = estimate_PD(theta, A_CV, A_PD, sigma_list, lambda_list)

	pd = 0;

	% get the sizes
	nsigma = size(A_CV, 1);
	nlambda = size(A_CV, 2);
	
	cv_score = zeros(nsigma, nlambda);

	% calculate the Pearson divergence for each sigma and lambda value
	for sigma_index = 1:nsigma
	
		for lambda_index = 1:nlambda
			cv_score(sigma_index, lambda_index) = theta'*A_CV{sigma_index, lambda_index}*theta;
		end	
	end

	% calculate the minimum cv_score (no ~ for old matlab)
	[tilde, idx_sigma, idx_lambda] = min2(cv_score);
	sigma_chosen = sigma_list(idx_sigma);
	lambda_chosen = lambda_list(idx_lambda);

	% calculate the Pearson divergence
	pd = theta'*A_PD{idx_sigma, idx_lambda}*theta - 1/2;
end

% Calculates the matrix for the cross-validation step for a specific sigma and lambda value
% 
function [A_CV] = calc_CV_matrix(x, y, xm, K, sigma, lambda)

	% size of the test dataset
	nm = size(xm, 2);

	% get the number of classes
	class_labels = sort(unique(y));
	c = length(class_labels);

	A_CV = zeros(c);

	% split up the labelled x values into the values for each class
	for i=1:c
		idx = find(y == class_labels(i));
		xc{i} = x(:, idx); % per class x values
		nl(i) = size(xc{i}, 2);

		cv_split_xh{i} = floor([0:nl(i)-1]*K./nl(i))+1;
	end

	% calculate the folds to split the xm dataset
	cv_split_xm = floor([0:nm-1]*K/nm)+1;

	% calculate the CV matrix
	for m=1:K
		for n = 1:K

			% calculate the center for the Gaussian basis
			x_ce = [];
			for i=1:c
				idx_tr = cv_split_xh{i}~=m;
				idx_te = cv_split_xh{i}==m;

				xc_tr{i} = xc{i}(:, idx_tr);
				xc_te{i} = xc{i}(:, idx_te);

				x_ce = [x_ce xc_tr{i}];
			end

			nb_t = size(x_ce, 2) + 1;

			% calculate the Phi matrix for the training and test set of the numerator
			xm_tr = xm(:, cv_split_xm~=n);
			xm_te = xm(:, cv_split_xm==n);

			nm_tr = size(xm_tr, 2);
			nm_te = size(xm_te, 2);

			Phi_xm_tr = GaussBasisA(xm_tr, x_ce, sigma)';
			Phi_xm_te = GaussBasisA(xm_te, x_ce, sigma)';

			% calculate Phi and h for the denominator samples
			H_tr = [];
			H_te = [];
			for i=1:c
				Phi_x_tr = GaussBasisA(xc_tr{i}, x_ce, sigma)';
				Phi_x_te = GaussBasisA(xc_te{i}, x_ce, sigma)';

				H_tr(:, i) = mean(Phi_x_tr, 2);
				H_te(:, i) = mean(Phi_x_te, 2);
			end

			% calculate the G matrix
			G_tr = Phi_xm_tr*Phi_xm_tr'/nm_tr;
			G_te = Phi_xm_te*Phi_xm_te'/nm_te;

			% calculate the inverse of (G_tr + lambda R)
			R = eye(nb_t); R(1,1) = 0;
			GlRinv = inv(G_tr + lambda*R);

			% calculate the A matrix
			Jt1 = (1/2)*H_tr'*GlRinv*G_te*GlRinv*H_tr;
			Jt2 = -H_tr'*GlRinv*H_te;

			% add the matrices
			A_CV = A_CV + Jt1 + Jt2;
		end
	end

	% A_CV ought to be the sum at this point
end

% Calculates the matrix for the cross-validation step for a specific sigma and a lambda_list
% Should give a significant speedup over the previous method
function [A_CV] = calc_CV_matrix_lambda_list(x, y, xm, K, sigma, lambda_list)

	% size of the test dataset
	nm = size(xm, 2);

	% get the number of classes
	class_labels = sort(unique(y));
	c = length(class_labels);

	for ll=1:length(lambda_list)
		A_CV{ll} = zeros(c);
	end

	% split up the labelled x values into the values for each class
	for i=1:c
		idx = find(y == class_labels(i));
		xc{i} = x(:, idx); % per class x values
		nl(i) = size(xc{i}, 2);

		cv_split_xh{i} = floor([0:nl(i)-1]*K./nl(i))+1;
	end

	% calculate the folds to split the xm dataset
	cv_split_xm = floor([0:nm-1]*K/nm)+1;

	% calculate the CV matrix
	for m=1:K
		for n = 1:K

			% calculate the center for the Gaussian basis
			x_ce = [];
			for i=1:c
				idx_tr = cv_split_xh{i}~=m;
				idx_te = cv_split_xh{i}==m;

				xc_tr{i} = xc{i}(:, idx_tr);
				xc_te{i} = xc{i}(:, idx_te);

				x_ce = [x_ce xc_tr{i}];
			end

			nb_t = size(x_ce, 2) + 1;

			% calculate the Phi matrix for the training and test set of the numerator
			xm_tr = xm(:, cv_split_xm~=n);
			xm_te = xm(:, cv_split_xm==n);

			nm_tr = size(xm_tr, 2);
			nm_te = size(xm_te, 2);

			Phi_xm_tr = GaussBasisA(xm_tr, x_ce, sigma)';
			Phi_xm_te = GaussBasisA(xm_te, x_ce, sigma)';

			% calculate Phi and h for the denominator samples
			H_tr = [];
			H_te = [];
			for i=1:c
				Phi_x_tr = GaussBasisA(xc_tr{i}, x_ce, sigma)';
				Phi_x_te = GaussBasisA(xc_te{i}, x_ce, sigma)';

				H_tr(:, i) = mean(Phi_x_tr, 2);
				H_te(:, i) = mean(Phi_x_te, 2);
			end

			% calculate the G matrix
			G_tr = Phi_xm_tr*Phi_xm_tr'/nm_tr;
			G_te = Phi_xm_te*Phi_xm_te'/nm_te;

			% calculate the inverse of (G_tr + lambda R)
			R = eye(nb_t); R(1,1) = 0;
			
			
			for ll=1:length(lambda_list)
				GlRinv = inv(G_tr + lambda_list(ll)*R);

				% calculate the A matrix
				Jt1 = (1/2)*H_tr'*GlRinv*G_te*GlRinv*H_tr;
				Jt2 = -H_tr'*GlRinv*H_te;

				% add the matrices
				A_CV{ll} = A_CV{ll} + Jt1 + Jt2;
			end
		end
	end

	% A_CV ought to be the sum at this point
end

function [A_PD] = calc_PD_matrix(x, y, xm, sigma, lambda)
	
	% size of the test dataset
	nm = size(xm, 2);

	% get the number of classes
	class_labels = sort(unique(y));
	c = length(class_labels);

	A_PD = zeros(c);

	% get the x values for each class
	for i=1:c
		idx = find(y == class_labels(i));
		xc{i} = x(:, idx); % per class x values
		nl(i) = size(xc{i}, 2);
	end
	
	% set the basis functions
	x_ce = x;
	b = size(x_ce, 2) + 1;

	% calculate H
	for i=1:c
		Phi_x = GaussBasisA(xc{i} , x_ce, sigma)';
		H(:, i) = mean(Phi_x, 2);
	end
	
	% calculate G and set R
	Phi_xm = GaussBasisA(xm, x_ce, sigma)';
	G = Phi_xm*Phi_xm'/nm;
	R = eye(b); R(1,1) = 0;
	
	GlRinv = inv(G + lambda*R);

	% calcualte the A matrix
	A1 = H'*GlRinv*H;
	A2 = - 1/2*H'*GlRinv*G*GlRinv*H;
	
	A_PD = A1 + A2; 
end


function [Phi] = GaussBasisA(z, c, sigma)
% from LSIF code

	Phi_tmp=-(repmat(sum(c.^2,1),size(z,2),1) ...
	+repmat(sum(z.^2,1)',1,size(c,2))-2*z'*c)/2;

	Phi=exp(Phi_tmp/(sigma^2));

	% Add the 1 basis vector
	Phi = [ones(size(Phi, 1), 1) Phi];
end

function [mval, idx1, idx2] = min2(A)

	mval = inf;
	idx1 = 1;
	idx2 = 1;

	for i=1:size(A,1)
		for j=1:size(A,2)
			if (A(i,j)<mval)
				mval = A(i,j);
				idx1 = i;
				idx2 = j;
			end
		end
	end
end