% demo.m
% A simple demonstration to show how the method works.
% Christo du Plessis
% christo@sg.cs.titech.ac.jp
%
 
	clear;clc;close all;
	
	rand('seed', 1);
	randn('seed', 1);
	
	% set the two distributions
	sigma1 = 1;
	sigma2 = 3;
	mu1 = 0;
	mu2 = 0;

	% set the data size and class balance
	ntr = 50;
	n1te = 80;
	n2te = 20;
	theta = n1te/(n1te + n2te);
	
	% generate the training data
	x1tr = randn(1, ntr)*sigma1 + mu1;
	x2tr =  randn(1, ntr)*sigma2+ mu2;
	xtr = [x1tr x2tr];
	ytr = [ones(1, ntr) 2*ones(1, ntr)];
	
	% generate the test data
	x1te = randn(1, n1te)*sigma1 + mu1;
	x2te = randn(1, n2te)*sigma2 + mu2;
	
	xte = [x1te x2te];
	xte = xte(:, randperm(n1te+n2te));
	
	% calculate the class prior
	[theta_est, theta_list, PD] = pe_prior_est_grid(xtr, ytr, xte);

	% Plot the training distributions and samples
	hfigtr = figure; hold on;
	set(gca,'FontName','Helvetica')
	set(gca,'FontSize',12)
	t = linspace(-7, 9);
	plot(t, normpdf(t, mu1, sigma1), 'r', 'LineWidth', 1.5);
	plot(t, normpdf(t, mu2, sigma2), 'g', 'LineWidth', 1.5);
	scatter(x1tr, normpdf(x1tr, mu1, sigma1), 'or');
	scatter(x2tr, normpdf(x2tr, mu2, sigma2), 'og');
	xlabel('x');
	legend('p_{tr}(x|y=1)', 'p_{tr}(x|y=2)');
	axis tight;
	%legend('boxoff');

	% plot the test (unlabeled distribution and samples)
	hfigte = figure; hold on;
	set(gca,'FontName','Helvetica')
	set(gca,'FontSize',12)
	yte = theta*normpdf(t, mu1, sigma1) + (1-theta)*normpdf(t, mu2, sigma2);
	ysamples = theta*normpdf(xte, mu1, sigma1) + (1-theta)*normpdf(xte, mu2, sigma2);
	plot(t, yte, 'k', 'LineWidth', 1.5);
	plot(t, normpdf(t, mu1, sigma1), 'r', 'LineWidth', 1.5);
	plot(t, normpdf(t, mu2, sigma2), 'g', 'LineWidth', 1.5);
	scatter(xte, ysamples, 'ok');
	legend('p_{te}(x)', 'p_{tr}(x|y=1)', 'p_{tr}(x|y=2)');
	axis tight;

	% 
	hfigpd = figure; hold on;
	set(gca,'FontName','Helvetica')
	set(gca,'FontSize',12)
	plot(theta_list, PD, 'k', 'LineWidth', 1.5);
	[theta_best, idx] = min(PD);
	scatter(theta_list(idx), theta_best , 'filled');
	xlabel('\theta');
	ylabel('PE(\theta)');
	legend('Pearson divergence');
