function [order, orientation, time, recovered_page] = assemble_shreds(shred_directory)
% This is a nearest neighbor search method for recovering a shredded
% document. Written by Gautam Machiraju, Winter 2019-2020 for EE 368.

% start timer
tic;

%-----------------
% HYPERPARAMETERS
%-----------------
e=12; % comparison region
w=3; % empty pixel wiggle room
mode = 'diff';

%---------------
% PREPROCESSING
%---------------

% load data + preprocess shreds
S = dir(fullfile(shred_directory,'*.png')); % pattern to match filenames.
shreds = struct([]);
disp('begin preprocessing...')

for i = 1:1:length(S) % remove this original STUFF for submission****
    if strcmp(S(i).name, 'original.png') == 1
        original = imread(strcat(shred_directory, '/', S(i).name));
    else
        shreds(i).filename = S(i).name;
        shreds(i).name = num2str(i); % ID =1-100, not 0-99
%         disp(shreds(i).filename)
%         disp(shreds(i).name)

        img = imread(strcat(shred_directory, shreds(i).filename));
        img_gray = rgb2gray(img);
        level = graythresh(img_gray); % otsu
        img_bw = imbinarize(img_gray,level);
        
        % Hough transform: shred rotation
        %---------------------------------
        % Find all the edges in the image
        BW = edge(img_bw,'canny');
        % Find the Hough information about the lines
        [H, theta, rho] = hough(BW);
        % First column is rho index, second is theta index
        P = houghpeaks(H,1);
        % Required rotatation
        theta_peak = theta(P(1,2));
        % Rotate image accordingly
        img = imrotate(img, theta_peak);
        img_bw = imrotate(img_bw, theta_peak);
        img_hsv = rgb2hsv(img);
        
        % crop black zones 
        %------------------
        mask = img_bw > 300; % mask = bright stuff
        % Fill in the body
        mask = imfill(mask, 'holes'); % mask is whole solid body.
        % OR it in with the zeros
        mask = mask | (img_bw == 0); % mask now includes pure zeros.
       % mask_3c = repmat(mask, [1,1,3]); % 3-channel mask
%         % Crop: extract pixels that are not masked
%         img = uint8(double(img) .* ~mask_3c); 
%         img_bw = nonzeros(double(img_bw) .* ~mask);
        okind = find(~mask);
        [ii,jj] = ind2sub(size(mask), okind);
        ymin=1; ymax=length(img); xmin=min(jj); xmax=max(jj);
        
        img_crop = imcrop(img, [xmin,ymin,xmax-xmin+1,ymax-ymin+1]);
        img_bw_crop = imcrop(img_bw, [xmin,ymin,xmax-xmin+1,ymax-ymin+1]);
        img_hsv_crop = rgb2hsv(img_crop); 
        
%         % enforce dims by padding if too short of a shred
%         if length(img_crop) < H
%             img_crop(end+1,:,:) = 1;
%         end
%         if length(img_bw_crop) < H
%             img_bw_crop(end+1,:) = 1;
%         end
%         if length(img_hsv_crop) < H
%             img_hsv_crop(end+1,:,:) = 1;
%         end
%         disp([i, length(img_crop), length(img_bw_crop), length(img_hsv_crop)])
%         
        % cache: cropped, uncropped, dims
        shreds(i).img_crop = img_crop;
        shreds(i).img_hsv_crop = img_hsv_crop;
        shreds(i).img_bw_crop = img_bw_crop;
        shreds(i).img = img;
        shreds(i).img_hsv = img_hsv;
        shreds(i).img_bw = img_bw;
        sz = size(img_bw_crop); shreds(i).width = sz(2); shreds(i).height = sz(1);
    end
end
    
% resize strip lengths
maxH = max([shreds(:).height]); % should be 3022 
for k = 1:1:length(shreds)
    shreds(k).img_hsv_crop = imresize(shreds(k).img_hsv_crop, [maxH, shreds(k).width]);
    sz = size(shreds(k).img_hsv_crop);
    shreds(k).height = sz(1);
end

%---------------
% MAIN ROUTINE
%---------------
disp('begin main routine...')

% seed shred
seed_img = shreds(1).img_hsv_crop;
seed_ids = [str2num(shreds(1).name)];
seed_orients = [0]; % use seed as is.... can rotate whole thing over if needed?
disp(strcat('added shred ID #', num2str(str2num(shreds(1).name)), ' / index #', num2str(str2num(shreds(1).name)), ', to the image!'))

% other shreds to search: 2-100
active = zeros(1, 100); active(str2num(shreds(1).name)) = 1; % keeps track of what shreds are used

while ~all(active == 1)
    
    % reset search standards for each shred
    best_cost = inf;
    
    for j = 1:1:100
        if active(j) == 0 % checking only inactive shreds
            shred_j = shreds(j).img_hsv_crop; 
            shred_jr = imrotate(shred_j,180);

            % trimming off the edge column before concatenation (0-padded)
            LHS = cat(2, seed_img(:,end-e-w+1:end-w,:), shred_j(:,1+w:e+w,:));
            RHS = cat(2, shred_j(:,end-e-w+1:end-w,:), seed_img(:,1+w:e+w,:));
            LHSr = cat(2, seed_img(:,end-e-w+1:end-w,:), shred_jr(:,1+w:e+w,:));
            RHSr = cat(2, shred_jr(:,end-e-w+1:end-w,:), seed_img(:,1+w:e+w,:));

            % search 1: rotation and direction of stitching
            [min_cost, orientation_id] = alignment_cost_fn(LHS,RHS,LHSr,RHSr, e, mode);

            % search 2: nearest neighbor cacheing
            if min_cost < best_cost
                best_index = j;
                best_shred_name = str2num(shreds(j).name);
                best_cost = min_cost;
                %best_shred_hsv = shred_j; % shreds(j).img_hsv
                best_shred = shreds(j).img_hsv_crop; %shred_j
                best_orientation_id = orientation_id;
            end
        end
    end
    
    % update based on search!
    [seed_ids, seed_orients, seed_img] = extend_shreds(best_orientation_id, best_shred, best_index, best_shred_name, seed_ids, seed_orients, seed_img);
    active(best_shred_name) = 1;
    disp(strcat('added shred ID #', num2str(best_shred_name), ' / index #', num2str(best_index), ', to the image!'))
end

% save image
recovered_page = hsv2rgb(seed_img);
imwrite(recovered_page, 'recovered.png');

orientation = seed_orients;
order = seed_ids;

% stop timer
time = toc;
end


%----------------
% HELPER METHODS
%----------------

function [min_cost, orientation_id] = alignment_cost_fn(LHS,RHS,LHSr,RHSr, e, mode)
    if strcmp(mode, 'diff') == 1
        a = hsv_dist(LHS(:,1:e,:), LHS(:,end-e+1:end,:)); 
        b = hsv_dist(RHS(:,1:e,:), RHS(:,end-e+1:end,:));
        c = hsv_dist(LHSr(:,1:e,:), LHSr(:,end-e+1:end,:));
        d = hsv_dist(RHSr(:,1:e,:), RHSr(:,end-e+1:end,:));
        
        costs = [a,b,c,d];
        min_cost = min(costs);
        orientation_id = find(costs == min_cost, 1);   
    end
end

function [seed_ids, seed_orients, seed_img] = extend_shreds(best_orientation_id, best_shred, best_index, best_shred_name, seed_ids, seed_orients, seed_img)
    if best_orientation_id == 2 % RHS
        seed_ids = [[best_shred_name], seed_ids]; %[int2str(best_shred_name), seed_ids];
        seed_orients = [[0], seed_orients];
        seed_img = cat(2, best_shred, seed_img);
    elseif best_orientation_id == 1 % LHS
        seed_ids = [seed_ids, [best_shred_name]]; % [seed_ids, int2str(best_shred_name)];
        seed_orients = [seed_orients, [0]];
        seed_img = cat(2, seed_img, best_shred);
    elseif best_orientation_id == 4 % RHSr
        seed_ids = [[best_shred_name], seed_ids]; % [strcat(int2str(best_shred_name),'*'), seed_ids];
        seed_orients = [[1], seed_orients];
        seed_img = cat(2, imrotate(best_shred,180), seed_img);
    elseif best_orientation_id == 3 % LHSr
        seed_ids = [seed_ids, [best_shred_name]]; % [seed_ids, strcat(int2str(best_shred_name),'*')];
        seed_orients = [seed_orients, [1]];
        seed_img = cat(2, seed_img, imrotate(best_shred,180));
    end
end

function [sum_dists] = hsv_dist(a, b)
    % d({h,s,v}) below should be in range [0,1]; a,b already normalized [0,1], so
    % no need for normalization.
    dh = abs(a(:,:,1)-b(:,:,1)); 
    ds = abs(a(:,:,2)-b(:,:,2));
    dv = abs(a(:,:,3)-b(:,:,3));
    dists = sqrt(dh.^2 + ds.^2 + dv.^2);
    sum_dists = sum(dists, 'all'); % used to be dists
end


