%%
%
% I used a few references to complete this assignment
% http://www.csse.uwa.edu.au/~pk/research/matlabfns/
%    for the homography denormalization
%
% http://en.wikipedia.org/wiki/RANSAC
%    basic ransac overview
%
% I consulted with Zhimin to figure out some of the problems
%
%

function fitting()

    % part 1
    %house_image1 = imread('house1.jpg');
    house_image2 = imread('house2.jpg');
    house_matches = load('house_matches.txt');
    house_camera1 = load('house1_camera.txt');
    house_camera2 = load('house2_camera.txt');

    %library_image1 = imread('library1.jpg');
    library_image2 = imread('library2.jpg');
    library_matches = load('library_matches.txt');
    library_camera1 = load('library1_camera.txt');
    library_camera2 = load('library2_camera.txt');

    get_figure_number(1);
    
    %part 2
    show_simple_fit(@least_squares_fit, @affine_estimate, house_matches, house_image2);
    show_simple_fit(@least_squares_fit, @affine_estimate, library_matches, library_image2);

    %part 3
    show_ransac(@least_squares_fit, @affine_estimate, house_matches, house_image2, ...
                3, 1000, 20, 20);
    show_ransac(@least_squares_fit, @affine_estimate, library_matches, library_image2, ...
                3, 1000, 20, 35);
    
    %part 4
    show_simple_fit(@svd_fit_unnormal, @projective_estimate, house_matches, house_image2);  
    show_simple_fit(@svd_fit_normal, @projective_estimate, house_matches, house_image2);
    show_simple_fit(@svd_fit_unnormal, @projective_estimate, library_matches, library_image2);
    show_simple_fit(@svd_fit_normal, @projective_estimate, library_matches, library_image2);
    
    %part 5
    show_ransac(@svd_fit_unnormal, @projective_estimate, house_matches, house_image2, ...
                4, 2000, 20, 10);
    show_ransac(@svd_fit_unnormal, @projective_estimate, library_matches, library_image2, ...
                4, 1000, 30, 30);
    
    %part 6
    show_epipolar(@fit_fundamental_unnormal, house_matches, house_image2);
    show_epipolar(@fit_fundamental_normal, house_matches, house_image2);
    show_epipolar(@fit_fundamental_unnormal, library_matches, library_image2);
    show_epipolar(@fit_fundamental_normal, library_matches, library_image2);
    
    % part 7
    show_triangulation(house_camera1, house_camera2, house_matches);
    show_triangulation(library_camera1, library_camera2, library_matches);
    
    % extra credit
    show_dominant_ransac(@svd_fit_unnormal, @projective_estimate, house_matches, house_image2, ...
                4, 1000, 30, 20);

end


%% ---------------------------------------------------------------
% entry points
function show_simple_fit(fit_func, estimate_func, points, image)
    M = fit_func( points(:,1:2), points(:,3:4) );
    
    estimated = estimate_func(M, points(:,1:2), points(:,3:4) );
    distance = distance_between(estimated, points(:,3:4));
    residual = mean( distance );

    figure( get_figure_number() ); imshow(image); hold on;
    show_points(points(:,3:4), estimated, '.b', '+r', [0.5 0 1]);
    title( ['residual:' num2str(residual)]);
end

function show_ransac(fit_func, estimate_func, points, image, initial_set_size, ...
                        num_tries, min_set_size, min_distance)
    [M inliers initial inlier_error] = simple_ransac(fit_func, estimate_func, points, ...
        initial_set_size, num_tries, min_set_size, min_distance);
    
    estimated = estimate_func(M, points(:,1:2), points(:,3:4) );
    distance = distance_between(estimated, points(:,3:4));
    residual = mean( distance );

    figure( get_figure_number() ); imshow(image); hold on;
    show_points(points(:,3:4), estimated, '.b', '+r', [0.5 0 1]);
    show_points(points(inliers,3:4), estimated(inliers,:), '.b', '+y', 'g');
    show_points(points(initial,3:4), estimated(initial,:), '.b', 'sy', 'g');
    title( ['inliers:' num2str(length(inliers)) ' inlier\_residual:' num2str(inlier_error) ...
            ' full\_residual:' num2str(residual)] );
end


function show_dominant_ransac(fit_func, estimate_func, all_points, image, initial_set_size, ...
                        num_tries, min_set_size, min_distance)

    % first set
    points1 = all_points;
    [M1 inliers1 initial1 inlier_error] = simple_ransac(fit_func, estimate_func, points1, ...
        initial_set_size, num_tries, min_set_size, min_distance);    
    estimated1 = estimate_func(M1, points1(:,1:2), points1(:,3:4) );
    
    figure( get_figure_number() ); imshow(image); hold on;
    show_dominant_all(points1, estimated1, 'r', inliers1, initial1);
    pause;
    
    % remove inliers
    points2 = points1;
    points2(inliers1,:) = [];
    
    % second set
    [M2 inliers2 initial2 inlier_error] = simple_ransac(fit_func, estimate_func, points2, ...
        initial_set_size, num_tries, min_set_size, min_distance);    
    estimated2 = estimate_func(M2, points2(:,1:2), points2(:,3:4) );
    
    figure( get_figure_number() ); imshow(image); hold on;
    show_dominant_all(points2, estimated2, 'y', inliers2, initial2);
    pause;
    
    % remove inliers
    points3 = points2;
    points3(inliers2,:) = [];
    
    % second set
    [M3 inliers3 initial3 inlier_error] = simple_ransac(fit_func, estimate_func, points3, ...
        initial_set_size, num_tries, min_set_size, min_distance);    
    estimated3 = estimate_func(M3, points3(:,1:2), points3(:,3:4) );
    
    figure( get_figure_number() ); imshow(image); hold on;
    show_dominant_all(points3, estimated3, 'g', inliers3, initial3);
    
    figure( get_figure_number() ); imshow(image); hold on;
    show_dominant_inliers(points1, estimated1, 'r', inliers1, initial1);
    show_dominant_inliers(points2, estimated2, 'y', inliers2, initial2);
    show_dominant_inliers(points3, estimated3, 'g', inliers3, initial3);
    
end


function show_epipolar(fit_func, matches, image)
    N = size(matches,1);
    
    % first, fit fundamental matrix to the matches
    F = fit_func(matches(:, 1:2), matches(:, 3:4)); % this is a function that you should write
    L = [matches(:,1:2) ones(N,1)] * F; % transform points from 
    % the first image to get epipolar lines in the second image

    % find points on epipolar lines L closest to matches(:,3:4)
    L = L ./ repmat(sqrt(L(:,1).^2 + L(:,2).^2), 1, 3); % rescale the line
    pt_line_dist = sum(L .* [matches(:,3:4) ones(N,1)],2);
    closest_pt = matches(:,3:4) - L(:,1:2) .* repmat(pt_line_dist, 1, 2);

    % find endpoints of segment on epipolar line (for display purposes)
    pt1 = closest_pt - [L(:,2) -L(:,1)] * 10; % offset from the closest point is 10 pixels
    pt2 = closest_pt + [L(:,2) -L(:,1)] * 10;
    
    % error
    distance = distance_between(closest_pt, matches(:,3:4));
    residual = mean( distance );

    % display points and segments of corresponding epipolar lines
    figure( get_figure_number() ); imshow(image); hold on;
    plot(matches(:,3), matches(:,4), '.b');
    line([matches(:,3) closest_pt(:,1)]', [matches(:,4) closest_pt(:,2)]', 'Color', [0.5 0 1]);
    line([pt1(:,1) pt2(:,1)]', [pt1(:,2) pt2(:,2)]', 'Color', 'r');
    title( ['residual: ' num2str(residual)] );
end

function show_triangulation(camera1, camera2, matches)
    [points reproj1 reproj2] = get_triangulation(camera1, camera2, matches);
    
    % get error in reprojections
    distance = distance_between(reproj1, matches(:,1:2));
    residual1 = mean( distance );
    distance = distance_between(reproj2, matches(:,3:4));
    residual2 = mean( distance );
    
    % get camera centers
    [U, S, V] = svd(camera1);
    center1 = V(:,end);
    center1 = unhomogenize(center1');
    [U, S, V] = svd(camera2);
    center2 = V(:,end);
    center2 = unhomogenize(center2');
    
    figure( get_figure_number() );
    
    axis equal;
    hold on;  % don't fall!
    %camproj('perspective');
    
    plot3(points(:,1), points(:,2), points(:,3), '+r');
    plot3(center1(1), center1(2), center1(3), '.b');
    plot3(center2(1), center2(2), center2(3), '.m');

    title(['residual 1:' num2str(residual1) ' residual 2:' num2str(residual2) ]);
end

%% ---------------------------------------------------------------
% graph functions

function show_points(points1, points2, point_style1, point_style2, line_style)
    plot(points1(:,1), points1(:,2), point_style1);
    plot(points2(:,1), points2(:,2), point_style2);
    line([points1(:,1) points2(:,1)]', [points1(:,2) points2(:,2)]', 'Color', line_style);
end

function f = get_figure_number(set_val)
    persistent figure_count;
    
    if(nargin == 1)
        figure_count = 0;
    else
        figure_count = figure_count+1;
        f = figure_count;
    end
end

function show_dominant_all(points, estimated, color, inliers, initial)
    show_points(points(:,3:4), estimated, '.b', '+b', 'b');
    show_dominant_inliers(points, estimated, color, inliers, initial);
end

function show_dominant_inliers(points, estimated, color, inliers, initial)
    show_points(points(inliers,3:4), estimated(inliers,:), '.b', ['+' color], 'b');
    show_points(points(initial,3:4), estimated(initial,:), '.b', ['s' color], 'b');
end

%% ---------------------------------------------------------------
% shared estimation functions

function [M, inliers, initial inlier_error] = simple_ransac(fitting_func, estimate_func, ...
            all_points, num_pts_needed, num_tries, min_set_size, min_distance)
    [rows, npts] = size(all_points);
    best_M = NaN;
    best_error =  inf;

    for i=1 : num_tries
        initial_points = randsample(rows, num_pts_needed);
        M = fitting_func( all_points(initial_points,1:2), all_points(initial_points,3:4) );
        estimated = estimate_func(M, all_points(:,1:2), all_points(:,3:4) );
        distance = distance_between(estimated, all_points(:,3:4));
        
        consensus_set = get_inliers(distance, initial_points, min_distance);

        if( length(consensus_set) >= min_set_size)
            consensus_from = all_points(consensus_set,1:2);
            consensus_to = all_points(consensus_set,3:4);
            M = fitting_func( consensus_from, consensus_to );
            estimated = estimate_func(M, consensus_from, consensus_to );
            error = mean( distance_between(estimated, consensus_to) );
            
            if error < best_error;
                best_error = error
                best_set = consensus_set;
                best_M = M;
                best_initial = initial_points;
            end
        end
    end

    if isnan(best_M)
        M = [];
        inliers = [];
        initial = [];
        inlier_error = inf;
        warning ( 'No matches found- time to explode' );
    else
        M = best_M;
        inliers = best_set;
        initial = best_initial;
        inlier_error = best_error;
    end
end

function I = get_inliers(distance, start_set, min_distance)
    inliers = [];
    num_points = size(distance);
    
    for i=1 : num_points(1)
        in_start_set = 0;
        for j=1 : length(start_set)
            if( i == start_set(j) )
                in_start_set = 1;
            end
        end
        
        if( distance(i) < min_distance || in_start_set)
            inliers = [inliers; i];
        end
    end
    
    I = inliers;
end

function distance = distance_between(points1, points2)
    distance = points1 - points2;
    distance = distance.^2;
    distance = sum(distance, 2);
    distance = sqrt(distance);
end

function [normalized, T] = get_normalize(points)
    num_points = length(points);
    
    center = mean(points(:,1:2));
    center_list = [center(1)*ones(num_points, 1) center(2)*ones(num_points, 1)];
    distance = mean(distance_between(center_list, points(:,1:2)));
    
    scale = sqrt(2)/distance;
    
    T = [scale 0 -scale*center(1); 0 scale -scale*center(2); 0 0 1];
    normalized = (T * points')';
end

function out_p = unhomogenize(in_p)
    [rows, cols] = size(in_p);
    
    out_p = zeros(rows, cols-1);
    
    for i=1 : cols-1
        out_p(:,i) = in_p(:,i)./in_p(:,cols);
    end
end

function out_p = homogenize(in_p)
    [rows, cols] = size(in_p);
    
    out_p = ones(rows, cols+1);
    
    for i=1 : cols
        out_p(:,i) = in_p(:,i);
    end
end

%% ---------------------------------------------------------------
% affine functions

function [Model Affine] = least_squares_fit(from_points, to_points)
    B = [to_points(:,1); to_points(:,2)];
    A = get_affine_equations(from_points(:,1:2));

    Model = A\B;
    
    if( nargout > 1)
        Affine = A;
    end
end

function A = get_affine_equations(image_points)
    num_points = size(image_points);
    num_points = num_points(1);
    
    A = [image_points(:,1)   image_points(:,2)    zeros(num_points, 1) ...
        zeros(num_points, 1) ones(num_points, 1)  zeros(num_points, 1);
        zeros(num_points, 1) zeros(num_points, 1) image_points(:,1) ...
        image_points(:,2)    zeros(num_points, 1) ones(num_points, 1)];
end

function estimated = affine_estimate(M, from_points, to_points)
    A = get_affine_equations( from_points );
    estimated_points = A*M;
    length = size(estimated_points);
    ex = estimated_points(1:length(1)/2,:);
    ey = estimated_points(length(1)/2+1:end,:);
    estimated = [ex ey];
end

%% ---------------------------------------------------------------
% projective functions

function H = svd_fit_normal(from_points, to_points)
    H = svd_fit(from_points, to_points, 1);
end

function H = svd_fit_unnormal(from_points, to_points)
    H = svd_fit(from_points, to_points, 0);
end

function H = svd_fit(from_points, to_points, normalize)
    num_points = length(from_points);
    homo_from = homogenize(from_points);
    homo_to = homogenize(to_points);
    
    % normalize
    if (normalize == 1)
        [homo_from, T_from] = get_normalize(homo_from);
        [homo_to, T_to] = get_normalize(homo_to);
    end
    
    A = get_projective_equations(homo_from, homo_to);
    
    [U, S, V] = svd(A);
    H = V(:,end);
    H = reshape(H, 3, 3)';
    
    % denormalize
    % I don't really know what's going on here
    % I tried the T1'*X*T2 formulation from the lecture and it doesn't
    % work at all. I then looked at the homography stuff here:
    % http://www.csse.uwa.edu.au/~pk/research/matlabfns/
    if (normalize == 1)
        H = inv(T_to) * H * T_from;
        %H = T_to \ H * T_from;
    end
end

function A = get_projective_equations(from_points, to_points)
    num_points = length(from_points);
    
    A = zeros(num_points*3, 9);
    for i=1 : num_points
        X = from_points(i,:);
        x = to_points(i,1);
        y = to_points(i,2);
            
        A(i*3-2,:) = [0 0 0  -X    y*X ];
        A(i*3-1,:) = [  X   0 0 0 -x*X ];
        A(i*3-0,:) = [ -y*X  x*X  0 0 0];
        %A(i*2-1,:) = [0 0 0  X    -y*X ];
        %A(i*2-0,:) = [  X   0 0 0 -x*X ];
    end
end

function estimated = projective_estimate(H, from_points, to_points );
    num_points = length(from_points);
    homo_from = homogenize(from_points);
    
    estimated = (H * homo_from')';
    estimated = unhomogenize(estimated);
end

%% ---------------------------------------------------------------
% fundamental matching

function F = fit_fundamental_normal(from_points, to_points)
    F = fit_fundamental(from_points, to_points, 1);
end

function F = fit_fundamental_unnormal(from_points, to_points)
    F = fit_fundamental(from_points, to_points, 0);
end

function F = fit_fundamental(from_points, to_points, normalize)
    num_points = length(from_points);
    homo_from = homogenize(from_points);
    homo_to = homogenize(to_points);
    
    % normalize
    if (normalize == 1)
        [homo_from, T_from] = get_normalize(homo_from);
        [homo_to, T_to] = get_normalize(homo_to);
    end
    
    A = get_fundamental_equations(homo_from, homo_to);
    
    [U, S, V] = svd(A);
    
    F = V(:,end);
    F = reshape(F, 3, 3)';
    
    %constraint time
    [U, S, V] = svd(F);
    S(3,3) = 0;
    F = U * S * V';
    
    % denormalize
    if (normalize == 1)
        F = T_from' * F * T_to;
    end
end

function A = get_fundamental_equations(from, to)
    num_points = length(from);
  
    A = [from(:,1).*to(:,1)  from(:,1).*to(:,2)  from(:,1)  from(:,2).*to(:,1) ...
         from(:,2).*to(:,2)      from(:,2)       to(:,1)       to(:,2) ...
         ones(num_points, 1) ];
end

%% ---------------------------------------------------------------
% triangulation

function [X, C1_points, C2_points] = get_triangulation(camera1, camera2, in_points)
    num_points = size(in_points, 1);
    
    X = zeros(num_points, 3);
    C1_points = zeros(num_points, 2);
    C2_points = zeros(num_points, 2);
    
    for i=1 : num_points
        A = get_triangulation_equations(camera1, camera2, in_points(i,:));
        [U, S, V] = svd(A);
        point = V(:,end);
        
        % reproject
        C1_reproject = camera1 * point;
        C2_reproject = camera2 * point;
        C1_points(i,:) =  unhomogenize(C1_reproject');
        C2_points(i,:) =  unhomogenize(C2_reproject');
        
        point = unhomogenize(point');
        X(i,:) = point;
    end
end

function c = get_cross_product(p)
    c = [  0   -p(3)  p(2);
          p(3)   0   -p(1);
         -p(2)  p(1)   0  ];
end

function A = get_triangulation_equations(camera1, camera2, points)  
    point1 = homogenize(points(:,1:2));
    point2 = homogenize(points(:,3:4));
    
    cross1 = get_cross_product(point1);
    cross2 = get_cross_product(point2);
    
    eq1 = cross1 * camera1;
    eq2 = cross2 * camera2;
    
    A = [eq1 ; eq2];
    %A = [point1(1)*camera1(3,:) - camera1(1,:);
    %     point1(2)*camera1(3,:) - camera1(2,:);
    %     point2(1)*camera2(3,:) - camera2(1,:);
    %     point2(2)*camera2(3,:) - camera2(2,:)];
end

%%