%% EE368/CS232 2020 Final Project
% Yunhao Bai


shred_directory = '19';
[order, orientation, time, recovered_page] = assemble_shreds(shred_directory);

function [order, orientation, time, recovered_page] = assemble_shreds(shred_directory)
% measure the runtime of function
t1 = clock;

disp(shred_directory);
% transfer the directory to str if necessary
%shred_directory = num2str(shred_directory,'%02d');

%% Part1: image pre-processing: 
% fix the rotation of images, remove the black edges
list = dir(fullfile([shred_directory,'/']));
fileNum = size(list,1);
% remove the outer-most interval number of lines for a cleaner edge
interval = 1; 

imStrips = cell(1,fileNum-3); % -3 for folder with the original.png
for i =3:fileNum-1 % -1 for folder with the original.png
    % image process and binarization
    im = im2double(imread([shred_directory,'/',list(i).name]));
    testIm = rgb2gray(im);
    % binarize the image with 1/255 threshold to distinguish it from the 
    % black background
    level = 1/255;
    imBW = imbinarize(testIm,level);
    [rowNum,~] = size(imBW);
    
    % find the rotation angle by locating the topLeft and bottomLeft corner
    % pixels
    topCorner = 1;
    while isempty(find(imBW(topCorner,:),1,'first'))
        topCorner = topCorner+1;
    end
    topLeft = find(imBW(topCorner,:),1,'first');
    
    bottomCorner = rowNum;
    while isempty(find(imBW(bottomCorner,:),1,'first'))
        bottomCorner = bottomCorner-1;
    end
    bottomLeft = find(imBW(bottomCorner,:),1,'first');
    % rotate the image based on arctangent value of the left edge
    rotAngle = 180*atan((topLeft-bottomLeft)/rowNum)/pi;
    temp = imrotate(im,rotAngle);
    imStrips{i-2} = temp;
    
    testIm = rgb2gray(imStrips{i-2});
    level = 1/255;
    imBW2 = imbinarize(testIm,level);
    a = find(imBW2(round(rowNum/2),:),1,'first');
    b = find(imBW2(round(rowNum/2),:),1,'last');
    % Due to the edge blur and aliasing effect, the edge of this cropped 
    % image always have at 1 almost black line at the very edge, the
    % second line will be partially blackened, the third line, during to
    % the contrast issue, will be slightly lighted. Remove 1 more pixel 
    % width to avoid the black gradient line is preferred, here I label as 
    % interval.
    imCrop = imcrop(imStrips{i-2},[a+interval,0,b-a-2*interval,rowNum]);
    imStrips{i-2} = imCrop;
    %imCrop = imcrop(imStrips{i-2},roi(1).BoundingBox);
    %imshow(imCrop);
    %imwrite(imCrop,[outputDir,list(i).name]);
end

% image strips edge fixation and generate image edge cells for matching
[rowNum,~,channel] = size(rgb2hsv(imStrips{1}));
leftEdges = zeros(fileNum-3,rowNum,channel);
rightEdges = zeros(fileNum-3,rowNum,channel);
imStripsHSV = cell(1,fileNum-3);

brightnessThres = 0;
for i = 1:fileNum-3
    temp = rgb2hsv(imStrips{i});
    [rowNum,colNum,~] = size(temp);
    % left edge
    tempLeftEdge = temp(:,1,:);
    for j = 1:rowNum
        % The Hue channel fixation of left edge
        if tempLeftEdge(j,1,1) == 0
            tempLeftEdge(j,1,1) = temp(j,2,1); 
            % fix the hue region equals zero
        end
        % The brightness channel fixation of left edge
        if tempLeftEdge(j,1,3) > brightnessThres/255
            tempLeftEdge(j,1,3) = (temp(j,2,3)+tempLeftEdge(j,1,3))/2;
        end
    end
    leftEdges(i,:,:) = tempLeftEdge;
    % right edge
    tempRightEdge = temp(:,colNum,:);
    for j = 1:rowNum
        % The Hue channel fixation of right edge
        if tempRightEdge(j,1,1) == 0
            tempRightEdge(j,1,1) = temp(j,colNum-1,1);
        end
        % The brightness channel fixation of right edge
        if tempRightEdge(j,1,3) > brightnessThres/255
            tempRightEdge(j,1,3) = (temp(j,colNum-1,3)+tempRightEdge(j,1,3))/2;
        end
    end
    rightEdges(i,:,:) = tempRightEdge;
    
    % fix the image in HSV channels for furture restoration
    temp(:,1,:) = tempLeftEdge;
    temp(:,colNum,:) = tempRightEdge;
    imStripsHSV{i} = temp;
end

%% Part2: Image matching and stiching
orientation = zeros(1,fileNum-3); 
edgePool = ones(1,fileNum-3);

%randomly select a start edge to contruct the image from
startEdge = randi(100,1);
orientation(startEdge) = 0;
order = [startEdge];
recovered_page = imStripsHSV{order};

% in case of multiple clusters that cannot be directly connect, assign each
% cluster into a cell, then connect at last
cellCounter = 1;
orderCell = {};
recovered_pageCell = {};

while sum(edgePool) ~= 0 
% when every item in the edgePool are used, the sum of edgePool equals zero

% construct the image from the right edge of the starting strip, assign the
% orientation of the starting strip as 0 (which will resulting flipped
% restored image)
    currentRightEdge = startEdge;
    while edgePool(currentRightEdge) ~= 0
        % remove the current right strip from the unused strip pool
        edgePool(currentRightEdge) = 0;
    
        if orientation(currentRightEdge) == 0
            rightEdgesMatrix = repmat(rightEdges(currentRightEdge,:,:),fileNum-3,1);
        else
            % when orientation is 1, the acutual next matching edge is its left
            % edge
            rightEdgesMatrix = fliplr(repmat(leftEdges(currentRightEdge,:,:),fileNum-3,1));
        end
        % right edge to other left edges
        temp = sum(abs(rightEdgesMatrix-leftEdges),2);
        RL_diff = sum(temp,3);
        [minRL, minIdxRL] = min(RL_diff);
        % right edge to flipped other right edges
        temp = sum(abs(rightEdgesMatrix-fliplr(rightEdges)),2);
        RFR_diff = sum(temp,3);
        [minRFR, minIdxRFR] = min(RFR_diff);
        % compare, find and stitch the next image strip to the right
        if minRFR > minRL
            nextRightStrip = minIdxRL;
            if orientation(currentRightEdge) == 0
                % when current image strip has orientation 0, if RL match is 
                % true, then the next right strip also has orientation 0.
                orientation(nextRightStrip) = 0;
            else
                % when current image strip has orientation 1, if RL match is 
                % true, then the next right strip still has orientation 0.
                orientation(nextRightStrip) = 0;
            end
            recovered_page = cat(2,recovered_page,imStripsHSV{nextRightStrip});
        else
            nextRightStrip = minIdxRFR;
            if orientation(currentRightEdge) == 0
                % when current image strip has orientation 0, if flipped R to R
                % match is true, then the next right strip also has orientation 
                % 1.
                orientation(nextRightStrip) = 1;
            else
                % when current image strip has orientation 1, if flipped R to R
                % match is true, then the next right strip also has orientation 
                % 1.
                orientation(nextRightStrip) = 1;
            end
            recovered_page = cat(2,recovered_page,rot90(imStripsHSV{nextRightStrip},2));
        end
    
        % update the next strip
        currentRightEdge = nextRightStrip;   
        % assign the next image Strip on the right side of the original strip
        order = [order,nextRightStrip];
    end

    % construct the image from the left edge of the starting image
    currentLeftEdge = startEdge;
    edgePool(currentLeftEdge) = 1;

    while edgePool(currentLeftEdge) ~= 0
        % remove the next left strip from the unused strip pool
        edgePool(currentLeftEdge) = 0;
        if orientation(currentLeftEdge) == 0
            leftEdgesMatrix = repmat(leftEdges(currentLeftEdge,:,:),fileNum-3,1);
        else
            % when orientation is 1, the acutual next matching edge is its
            % right edge
            leftEdgesMatrix = fliplr(repmat(rightEdges(currentLeftEdge,:,:),fileNum-3,1));
        end
        % right edge to other left edges
        temp = sum(abs(leftEdgesMatrix-rightEdges),2);
        LR_diff = sum(temp,3);
        [minLR, minIdxLR] = min(LR_diff);
        % right edge to flipped other right edges
        temp = sum(abs(leftEdgesMatrix-fliplr(leftEdges)),2);
        LFL_diff = sum(temp,3);
        [minLFL, minIdxLFL] = min(LFL_diff);
        % compare, find and stitch the next image strip to the right
        if minLFL > minLR
            nextLeftStrip = minIdxLR;
            if orientation(currentLeftEdge) == 0
                orientation(nextLeftStrip) = 0;
            else
                orientation(nextLeftStrip) = 0;
            end
            recovered_page = cat(2,imStripsHSV{nextLeftStrip},recovered_page);
        else
            nextLeftStrip = minIdxLFL;
            if orientation(currentLeftEdge) == 0
                orientation(nextLeftStrip) = 1;
            else
                orientation(nextLeftStrip) = 1;
            end
            recovered_page = cat(2,rot90(imStripsHSV{nextLeftStrip},2),recovered_page);
        end
    
        % update the next strip
        currentLeftEdge = nextLeftStrip;   
        % assign the next image Strip on the right side of the original strip
        order = [nextLeftStrip,order];
    end

    % remove the edge image strip as they are looped one more times at the
    % boundaries
    [rowNum,colNum,~] = size(recovered_page);
    firstStrip = order(1);
    [~,removeLeftNum,~] = size(imStripsHSV{firstStrip});
    lastStrip = order(end);
    [~,removeRightNum,~] = size(imStripsHSV{lastStrip});
    recovered_page = imcrop(recovered_page,[removeLeftNum,0, ...
        colNum-removeLeftNum-removeRightNum,rowNum]);
    recovered_pageCell{cellCounter} = recovered_page;

    % remove the edge elements in order as they are looped one more times at the
    % boundaries
    order = order(2:end-1);
    orderCell{cellCounter} = order;

    % update the counter
    cellCounter = cellCounter+1;

    % as some images cannot be biconnected into a pefect "pearl necklace"
    % like image, remove so there will be multiple clusters, after the
    % first big cluster of image be restored, select the next one and do
    % the construction from a new image strip
    startEdge = find(edgePool==1,1);
    % if cannot find any elememt in the edgePool, means all image strips
    % has been ultilized, so break the cycle
    if isempty(startEdge)
        break;
    end
    orientation(startEdge) = 0;
    order = [startEdge];
    recovered_page = imStripsHSV{order};
end

% connect together the orders and recovered_page from each cluster to give
% the final order vector and restored image.
order = [orderCell{:}];
recovered_page = [recovered_pageCell{:}];
%figure;
%imshow(hsv2rgb(recovered_page));

%% Part3: Image post fixation
recovered_page = medfilt3(recovered_page,[1,5,1]);
recovered_page = hsv2rgb(recovered_page);
figure;
imshow(recovered_page);
imwrite(recovered_page,[shred_directory,'_recovered_2.png']);

%% Part4: measure the runtime of the function
t2 = clock;
time = etime(t2,t1);
disp(['The elapsed time is: ',num2str(time), ' seconds.']);
end