clear all;

numberOfPoints = 2; % number of user selected points
sampleRadius = 32;  % radius of subImage around selected point
pickPoints = 0;     % switch to pick points or to use default points

% read in all the images and square them
image1 = double(imread('canyon1-2.png'));
image1sq = image1.^2;
size1h = size(image1, 1);
size1w = size(image1, 2);

image2 = double(imread('canyon2-2.png'));
image2sq = image2.^2;
size2h = size(image2, 1);
size2w = size(image2, 2);

userX = []; % user selected points in first image
userY = []; % user selected points in first image

if (pickPoints)
    % display unaltered image
    figure('Name', 'Raw Images', 'Position', [0 0 1200 400]);
    subplot(1, 2, 1); imagesc(image1); axis square; colormap gray; xlabel('Select Points To Match In This Image'); hold on;
    subplot(1, 2, 2); imagesc(image2); axis square; colormap gray; xlabel('Second Image'); hold on;



    % for number of points loops and stores inputted points in array

    for i = 1:numberOfPoints
       [tempy, tempx] = ginput(1);
       tempx = round(tempx);
       tempy = round(tempy);
       userX = [userX ; tempx];
       userY = [userY ; tempy];
       plot(tempy, tempx, 'r+');
       t = text(tempy + 10, tempx, num2str(i));
       set(t, 'Color', [1 0 0]);
    end
else
   userX = [150 200 150];
   userY = [150 200 200];  
end

% userX = [400 500];
% userY = [400 500];

% attempt to find matching point for each point in second image
% find by taking sub image around selected point, and using minimum
% squared error template matching
% if matching point is found, mark it with an x

Xoffset1 = [];  % list of x coordinates of points in image1
Xoffset2 = [];  % list of x coordinates of points in image2
Yoffset1 = [];  % list of y coordinates of points in image1
Yoffset2 = [];  % list of y coordinates of points in image2

image1F = fft2(image1);
image2F = fft2(image2);

image1sqF = fft2(image1sq);
image2sqF = fft2(image2sq);

mask = zeros(size1h, size1w);
mask(1:sampleRadius, 1:sampleRadius) = 1.0;
maskF = fft2(mask);

for i = 1:numberOfPoints

    % fix this!  both coordinates have to match for some reason
    % examples -> userX(i) == userY(i).  Why?
    subImage = image1(userX(i):(userX(i) + sampleRadius - 1), userY(i):(userY(i) + sampleRadius - 1));

    totsq = sum(sum(subImage.^2));

    subImagePad = zeros(size1h, size1w);
    subImagePad(1:sampleRadius, 1:sampleRadius) = subImage;

    subImgF = fft2(subImagePad);

    cc1F = image1F.*conj(subImgF);
    cc2F = image2F.*conj(subImgF);

    cc1 = abs(ifft2(cc1F));
    cc2 = abs(ifft2(cc2F));

    cct1 = abs(ifft2(image1sqF.*conj(maskF))) - 2 * cc1 + totsq;
    cct2 = abs(ifft2(image2sqF.*conj(maskF))) - 2 * cc2 + totsq;

    [m1, indx1] = min(cct1(:))
    [m2, indx2] = min(cct2(:))

    [xoff1, yoff1] = ind2sub(size(cct1), indx1)
    [xoff2, yoff2] = ind2sub(size(cct2), indx2)
    
    Xoffset1 = [Xoffset1; xoff1];
    Xoffset2 = [Xoffset2; xoff2];
    Yoffset1 = [Yoffset1; yoff1];
    Yoffset2 = [Yoffset2; yoff2];
    
     clear cc1;
     clear cc2;
    clear cct1;
    clear cct2;
    clear cc1F;
    clear cc2F;

end

clear image1F;
clear image2F;

figure('Name', 'Resulting Matched Locations', 'Position', [0 0 700 300]);
subplot(1,2,1); imagesc(image1); axis square; colormap gray;
hold on;

for i = 1:numberOfPoints
    plot(Yoffset1(i), Xoffset1(i), 'rx');
    t = text(Yoffset1(i) + 10, Xoffset1(i), num2str(i));
    set(t, 'Color', [1 0 0]);
end

hold off;

subplot(1,2,2); imagesc(image2); axis square; colormap gray;
hold on;
for i = 1:numberOfPoints
    plot(Yoffset2, Xoffset2, 'rx');
    t = text(Yoffset2(i) + 10, Xoffset2(i), num2str(i));
    set(t, 'Color', [1 0 0]);
end
hold off;

% At this point, we have identified numberOfPoints points on
% image1, and found the matching point in image2

% we want to find the transformation matrix based on these points
% that will transform second image to match the first image

image1pts = [Xoffset1 ; Yoffset1];
matA = [Xoffset2' Yoffset2'; Yoffset2' -Xoffset2'; ones(numberOfPoints, 1)' zeros(numberOfPoints, 1)';zeros(numberOfPoints, 1)' ones(numberOfPoints, 1)']';

% solve linear equation
 tempA = matA \ image1pts;
% 
oldA = [tempA(1) -tempA(2) tempA(4); tempA(2) tempA(1) tempA(3); 0 0 1]

p2 = [Xoffset1'; Yoffset1'; ones(numberOfPoints, 1)'];
p1 = [Xoffset2'; Yoffset2'; ones(numberOfPoints, 1)'];

newA = 2 * inv(p1*p1'+p1*p1')*(p1 * p2')
A = newA';
temp = A(2,3);
A(2,3) = A(1,3);
A(1,3) = temp;
A(1,2) = -A(1,2);
A(2,1) = -A(2,1);
A = A
%A = (p1 * p2') * inv(p1 * p1');


% now we apply the transformation to the corners
% of image2 to create a new image to hold the transformed image2
corners = A * [1 1 size2w size2w; 1 size2h 1 size2h; 1 1 1 1];
[newX, newY] = ndgrid(min([corners(1,:) 0]) : max([corners(1,:) size1w]), min([corners(2,:) 0]) : max([corners(2,:) size1h]));
[newW newH] = size(newX);

% transform image2
newImage2 = A \ [newX(:) newY(:) ones(newW * newH, 1)]';

% interpolate new image
clear transformedImage2;

tempX = reshape(newImage2(1,:), newW, newH)';
tempY = reshape(newImage2(2,:), newW, newH)';
transformedImage = interp2(image2, tempX, tempY, '*bilinear');

newImageOffset = round([max([corners(1,:) 0]) max([corners(2,:) 0])]);

finalImage(1 : newH, 1 :newW) = transformedImage;
%finalImage(27 : 27 + newH - 1, 125 : 125 + newW -1) = transformedImage;

finalSize = [1 size1h 1 size1w];

if (tempA(3,1) < 0)
    finalSize(1,1) = -round(tempA(3,1));
    finalSize(1,2) = -round(tempA(3,1)) + size1h - 1;
end
if (tempA(4,1) < 0)
    finalSize(1,3) = -round(tempA(4,1));
    finalSize(1,4) = -round(tempA(4,1)) + size1w - 1;
end

%    finalImage(finalSize(1,1):finalSize(1,2), finalSize(1,3):finalSize(1,4)) = image1;


finalImage(1:size1h, 1:size1w) = image1;

%transformedImage(1 : size1h, 1 : size1w) = image1;

figure('Name', 'Stitched Image', 'Position', [0 0 350 350]);
imagesc(finalImage); axis equal; colormap gray