%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% evaluate the number of common nearest neighbours

%% input parameters
delta_t				= 4e-3;
src_file			= 'L:\array\tcf\rescaled\z001_t001.mat';
n_subset			= 1e6;
n_neighbours_max	= 10;
font_size			= 24;
P_loss				= 3e-2;

%% load particles
part		= load(src_file);

% evaluate displacements of subset
subset		= 1 : min(n_subset, part.n_part);
part.x		= part.x(:, subset);
part.u		= part.u(:, subset);
part.n_part = length(subset);

re_part_A	= part.x;
re_part_B	= part.x + part.u*delta_t;


% resample
n_A			= size(part.x, 2);
n_B			= floor(n_A*(1-P_loss));
iBinA		= 1 : n_A;
lost		= randperm(n_A, n_A-n_B);
iBinA(lost)			= [];
re_part_B(:,lost)	= [];

%% evaluate nearest neighbours
P_BinA		= nan(n_neighbours_max, n_neighbours_max);

for k = 1 : n_neighbours_max
	%% find the k nearest neighbours in sets A and B
	fprintf('k=%d, searching...\n', k);
	[inn_A, r_A]	= knnsearch(re_part_A.', re_part_A.', 'K', k+1);
	[inn_B, r_B]	= knnsearch(re_part_B.', re_part_B.', 'K', k+1);
	% remove self
	inn_A(:,1)		= [];
	inn_B(:,1)		= [];
	
	%% establish whether particle is in nearest neighbours
	b_BinA			= false(n_B, k);
	for ipB = 1 : n_B
		ipA			= iBinA(ipB);
		for col = 1 : k
			b_BinA(ipB,col)	= ismember(iBinA(inn_B(ipB,col)), inn_A(ipA,:));
		end
	end
	P_BinA(k,1:k)	= sum(b_BinA, 1) ./ n_B;
end

%% show results
figure(1);
clf;
k_range		= 1 : n_neighbours_max;
[nsize,nidx]= ndgrid(k_range);
hd			= imagesc(k_range, k_range, 1-P_BinA.');
set(hd, 'AlphaData', ~isnan(P_BinA.'));
axis equal tight xy;
caxis([0 0.05]);
cb			= colorbar;
colormap(jet);
set(gca,'xtick',k_range, 'ytick', k_range);

titstr		= sprintf('$P_{loss} = %0.1f\\%%$', P_loss*100);
title(titstr, 'interpreter', 'latex');
title(cb, 'P(missing)', 'interpreter', 'latex');
xlabel('Neighbourhood Size $K$', 'interpreter', 'latex');
ylabel('Neighbour Index $n$', 'interpreter', 'latex');
set(gca,'fontsize',font_size,'ticklabelinterpreter','latex');
set(cb,'fontsize',font_size,'ticklabelinterpreter','latex');

figstr		= sprintf('figures/Pmissing_map_N%0.0e_Ploss%0.1e.fig', n_A, P_loss);
drawnow;
saveas(gcf, figstr);

%% plot on graph
figure(2);
clf;
for nsize = 1 : n_neighbours_max
	nidx	= 1 : nsize;
	legstr	= sprintf('$K = %d$', nsize);
	plot(nidx, 1-P_BinA(nsize, nidx), 'o-', 'DisplayName', legstr);
	hold on;
end
P_K_lost	= 1 - P_BinA(sub2ind([n_neighbours_max n_neighbours_max], k_range, k_range));
pfit		= polyfit(k_range, P_K_lost, 1);
pval		= polyval(pfit, k_range);
plot(k_range, pval, 'k--', 'LineWidth', 2, 'DisplayName', 'Fit $n=K$');

titstr		= sprintf('$P_{loss} = %0.1f\\%%$', P_loss*100);
title(titstr, 'interpreter', 'latex');
xlabel('Neighbour Index $n$', 'interpreter', 'latex');
ylabel('P(missing)', 'interpreter', 'latex');
pbaspect([4/3 1 1]);
set(gca,'fontsize',font_size,'ticklabelinterpreter','latex');
set(legend, 'location', 'northwest', 'interpreter', 'latex', 'numcolumns', 2);

figstr		= sprintf('figures/Pmissing_N%0.0e_Ploss%0.1e.fig', n_A, P_loss);
drawnow;
saveas(gcf, figstr);
