%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% perform self-calibration on particle images 
%
% uses stereo-matched locations of particles generated by track_movie_shadowgraph.m
% and identifies those which are unambiguously matched. Then, new
% calibrations are generated from the old ones that minimise 
% the disparity between observations and triangulated particles
%
% The new calibrations are constrained to be similar to the old
% calibration. It is possible to select which parameters in the calibration
% should be varied.
%
addpath(genpath([pwd '/../core']));
addpath([pwd '/../../misc']);

%% input parameters
root_dir				= 'F:\mixer\lpt\run29\';
match_file				= '/tracks/tracers.selfcal.210630.mat';
calib_file				= [root_dir '/tracks/camera-model.selfcal2.210630.mat'];

residual_percentiles	= [5 25 50 75 95];
perturb_model_flags		= struct('b_move_camera',	false, ...
								 'b_rotate_camera',	true, ...
                                 'b_bundle_calibration', false);
									
font_size				= 20;
f_cutoff_nd				= 0.1;	% non-dimensional cutoff frequency
n_min_length            = 10;

n_i_bins                = 16;
n_j_bins                = 16;
n_dep_bins              = 16;

%% load stereo matching data
match_file_full			= [root_dir match_file];
fprintf('loading %s\n', match_file_full);
fs						= importdata(match_file_full);
old_setup				= fs.setup;
tracks				    = fs.tracks;

img_max_dist			= fs.opts.triangulation.R_tolerance_img;

% old calibration
n_cameras				= old_setup.n_cameras;
old_cameras				= old_setup.camera;

%% filter tracks
% apply a temporal filter to position in image space to reduce
% random errors
fprintf('filtering tracks at f/f_s = %0.2f\n', f_cutoff_nd);
frame_rate              = fs.opts.frame_rate;
f_cutoff                = f_cutoff_nd * frame_rate;
omega	                = 2*pi*f_cutoff;
lambda	                = 1/omega^6;
% spline smoothing parameter

n_tracks				= length(tracks);
filtered                = {};
for n = 1 : n_tracks
    %% filter
    trk                 = tracks(n);
    t_ok	            = trk.t(1, ~trk.ext);
    if length(t_ok) <= n_min_length
        continue;
    end

    re_ok               = trk.re_pos(:, ~trk.ext);
    n_t                 = length(trk.t);
    
    im_filt             = zeros(2, n_t, n_cameras);
    for c = 1 : n_cameras
        im_meas         = squeeze(trk.im_pos(:,c,~trk.ext));
        assert(~any(isnan(im_meas),'all'))
        im_filt(:,:,c)  = ssplnd(t_ok(:), im_meas.', trk.t(:), lambda).';
    end
    assert(~any(isnan(im_filt),'all'))
    im_filt             = permute(im_filt, [1 3 2]);

    %% save in predictor field
    trk.im_pred         = im_filt;
    filtered{end+1}     = trk;
end
filtered                = cell2mat(filtered);
n_filtered              = length(filtered);

%% identify matched points from tracks

% extract those particles which are unambiguously matched
filtered				= reshape(filtered, [1 1 n_filtered]);
% get image space coordinates
im_valid				= arrayfun(@(x) x.im_pred(:,:,~x.ext), filtered, 'UniformOutput', false);
im_valid                = cell2mat(im_valid);
% permute
im_valid				= permute(im_valid, [1 3 2]);

% re-triangulate using reference camera model
[re_old,~,im_old_res]   = triangulate(old_cameras, im_valid);

%% recalibrate with bundle self calibration
fprintf('Recalibrating\n');
new_setup				= old_setup;
if perturb_model_flags.b_bundle_calibration
    % bundle calibration
    [new_setup.camera,im_res_old,im_res_new] = bundle_self_calibration(old_setup.camera, im_valid);
    im_rms_old  = rms(im_res_old,2);
    im_rms_new  = rms(im_res_new,2);
else
    % independent re calibration
    for c = 1 : n_cameras
	    %% recalibrate
	    old_cal				= old_setup.camera(c).calib;
	    im_meas				= im_valid(:,:,c);
	    im_proj				= pinhole_model(re_old, old_cal, true);
	    old_res				= sqrt(sum((im_proj - im_meas).^2,1));
	    [new_cal, new_res]	= pinhole_fit_constrained(old_cal, im_meas, re_old, perturb_model_flags);
        % update calibration
	    new_setup.camera(c).calib= new_cal;
    end
end
[re_new,~,im_new_res]   = triangulate(new_setup.camera, im_valid);

%% measure new and old residuals
% define bins for mean residual
i_axis                  = old_setup.camera(1).i_axis;
j_axis                  = old_setup.camera(1).j_axis;
i_lim                   = sort(i_axis([1 end]));
j_lim                   = sort(j_axis([1 end]));
i_edges                 = linspace(i_lim(1), i_lim(2), n_i_bins+1)';
j_edges                 = linspace(j_lim(1), j_lim(2), n_j_bins+1)';
s_edges                 = zeros(n_dep_bins+1, n_cameras);
im_res_map              = zeros(n_i_bins, n_j_bins, n_dep_bins, 2);
e2b                     = @(x) (x(2:end)+x(1:end-1))/2;
i_bins                  = e2b(i_edges);
j_bins                  = e2b(j_edges);
s_bins                  = e2b(s_edges);
[im,jm]                 = ndgrid(i_bins,j_bins);

for c = 1 : n_cameras
    %% calculate difference to previous
    old_cal				= old_setup.camera(c).calib;
    new_cal             = new_setup.camera(c).calib;
	% shift, rotation
	xc_shift			= new_cal.xc - old_cal.xc;
	Omega				= old_cal.R \ new_cal.R;
	omega_vec			= vrrotmat2vec(Omega);

    

    %% check residual
    old_proj = pinhole_model(re_old, old_setup.camera(c).calib);
	new_proj = pinhole_model(re_new, new_setup.camera(c).calib);
	im_meas  = im_valid(:,:,c);

    old_res  = im_meas - old_proj;
    new_res  = im_meas - new_proj;

    %% quantize residuals
    % define bins in depth coordinate
    GR                  = old_cal.G * old_cal.R;
    e_ray               = GR \ [mean(i_axis); mean(j_axis); 1];
    e_ray               = e_ray/norm(e_ray);
    x_lb                = linprog(+e_ray, old_setup.gl_bb_A, old_setup.gl_bb_b);
    x_ub                = linprog(-e_ray, old_setup.gl_bb_A, old_setup.gl_bb_b);
    s_lim               = sort([dot(e_ray,x_lb) dot(e_ray,x_ub)]);
    s_edges(:,c)        = linspace(s_lim(1), s_lim(2), n_dep_bins+1);
    % quantize into bins
    re_dep              = sum(e_ray .* re_old, 1);
    [~,ib]              = histc(im_meas(1,:)', i_edges);
    [~,jb]              = histc(im_meas(2,:)', j_edges);
    [~,kb]              = histc(re_dep(:), s_edges(:,c));
    ok                  = ib >= 1 & ib <= n_i_bins & jb >= 1 & jb <= n_j_bins & kb >= 1 & kb <= n_dep_bins;
    im_res_map(:,:,:,1) = accumarray([ib(ok,1) jb(ok,1) kb(ok,1)], old_res(1,ok)', [n_i_bins, n_j_bins, n_dep_bins], @mean);
    im_res_map(:,:,:,2) = accumarray([ib(ok,1) jb(ok,1) kb(ok,1)], old_res(2,ok)', [n_i_bins, n_j_bins, n_dep_bins], @mean);

	%% report 
	fprintf('Camera %d:\n', c);
    fprintf('   old residual:    %0.2f %0.2f px\n', rms(old_res,2));
    fprintf('   new residual:    %0.2f %0.2f px\n', rms(new_res,2)); 
    fprintf('   shift:           %0.3f %0.3f %0.3f mm\n', xc_shift);
	fprintf('   rotation:        %0.3f mrad\n',  1000 * omega_vec(4));
	fprintf('   about:           %0.3f %0.3f %0.3f\n', omega_vec(1:3));
end

%% calculate statistics
% number of matches
n_matches				= arrayfun(@(x) size(x.re_pos,2), filtered);
n_matches_total			= sum(n_matches);
n_valid					= arrayfun(@(x) nnz(~x.ext), filtered);
n_valid_total			= sum(n_valid);

% statistics of residuals
residual_rms_old		= rms(im_old_res, 1);
residual_rms_new        = rms(im_new_res, 1);

residual_edges			= linspace(0, img_max_dist, 51);
residual_bins			= (residual_edges(2:end) + residual_edges(1:end-1))/2;
% histograms
residual_hist_old		= histc(im_old_res, residual_edges);
residual_hist_old		= residual_hist_old(1:end-1, :);
residual_pdf_old		= bsxfun(@rdivide, residual_hist_old, trapz(residual_bins, residual_hist_old, 1));
residual_hist_new		= histc(im_new_res, residual_edges);
residual_hist_new		= residual_hist_new(1:end-1, :);
residual_pdf_new		= bsxfun(@rdivide, residual_hist_new, trapz(residual_bins, residual_hist_new, 1));

%% plot residuals
figure(1);
clf;
markers					= {'r-','g-','b-'; 'r--','g--','b--'};
legstr					= cell(2,n_cameras);
for c = 1 : n_cameras
	plot(residual_bins, residual_pdf_old(:,c), markers{1,c});
	hold on;
    plot(residual_bins, residual_pdf_new(:,c), markers{2,c});
	hold on;
	legstr{1,c}			= sprintf('C%d old', c);
    legstr{2,c}			= sprintf('C%d new', c);
end

% make pretty
L							= legend(legstr(:));
set(L,	'Location', 'South', ...
			'Interpreter', 'Latex', ...
			'fontsize', font_size);
make_plot_pretty(	'Size of residual, $r$ / px', ...
						'$P(r)$', ...
						'PDF of residuals in image plane', font_size);
xlim(residual_edges([1 end]));
ylim([0 ylim*[0;1]]);

%% plot residual map
figure(10);
clf;
scale=20;
layout=tiledlayout(2,2);
for c = 1 : n_cameras
    nexttile;
    cmap = jet(n_dep_bins);
    for k = 1: n_dep_bins
        hquiv = quiver3(im, jm, s_edges(k,c) + 0*im, ...
                    im_res_map(:,:,k,1)*scale, im_res_map(:,:,k,2)*scale, 0*im, 0);
        set(hquiv, 'color', cmap(k,:));
        hold on;
    end
    daspect([1 1 1/20])
end

%% plot residual images
figure(2);
clf;
layout=tiledlayout(2,2);

% 3d plot of positions
nexttile
plot3(re_old(1,:), re_old(2,:), re_old(3,:), 'r.');
hold on;
plot3(re_new(1,:), re_new(2,:), re_new(3,:), 'b.');
axis equal tight xy;

view(3);
camup([0 1 0])
axis(old_setup.gl_bb_lim(:));
xlabel('$x$ [mm]','interpreter','latex')
ylabel('$y$ [mm]','interpreter','latex')
zlabel('$z$ [mm]','interpreter','latex')

% residuals
for c = 1 : n_cameras
	nexttile;

	old_proj = pinhole_model(re_old, old_setup.camera(c).calib);
	new_proj = pinhole_model(re_new, new_setup.camera(c).calib);
    im_meas  = im_valid(:,:,c);
    old_res  = old_proj - im_meas;
    new_res  = new_proj - im_meas;

	plot(old_proj(1,:), old_proj(2,:), 'r.')
	hold on;
	plot(new_proj(1,:), new_proj(2,:), 'b.')
	plot(im_meas(1,:), im_meas(2,:), 'go')
    quiver(new_proj(1,:), new_proj(2,:), -new_res(1,:), -new_res(2,:), 0, 'k')

	axis equal tight xy;

    titstr=sprintf('camera %d', c);
    title(titstr);
end


%% save new calibration
fprintf('Saving calibration to %s\n', calib_file);
if exist(calib_file, 'file')
	ack						= input('Calibration file already exists, overwrite? (y/n) ', 's');
	if strcmp(ack, 'y')
		save(calib_file, '-struct', 'new_setup');
	end
else
	save(calib_file, '-struct', 'new_setup');
end
