function [order,orientation,time,recovered_page] = assemble_shreds(shred_directory)
% Assemble a page from the shreds from the given directory
%{
INPUT:
    shred_directory     string, directory containing all N pngs, one for
                            each strip.
OUTPUT:
    order               Nx1 double, vector containing the index of each 
                            strip (the yy value) in the order that they
                            should be stitched together for the page. The
                            vector cannot contain duplicates.
    orientation         Nx1 logical, whether or not each strip should be
                            flipped 180 degrees.
    time                scalar double, time for the method to run
    recovered_page      WxHx3 double, the recovered image
%}

    tic
    [shreds,flipped] = preprocess(shred_directory);
    [g_orders,g_shreds] = groupshreds(shreds);
    [order,score,flipped,shreds] = iterflip(shreds,g_orders,g_shreds,flipped);
    for k = 1:6
        [g_orders,g_shreds] = regroupshreds(g_orders,g_shreds,shreds);
        [order_n,score_n,flipped_n,shreds_n] = iterflip(shreds,g_orders,g_shreds,flipped);
        if score_n > score
            order = order_n;
            score = score_n;
            flipped = flipped_n;
            shreds = shreds_n;
        end
    end
    recovered_page = makepage(shreds,order,1);
    orientation = flipped(order)';
    time = toc;
    
end

%% Functions

function [g_orders,g_shreds] = groupshreds(shreds)
% Sort shreds into groups
    N = size(shreds,1);
    list = 1:N;
    
%% Get similarity scores
    max_group_size = 12;
    scores = NaN(N,N);
    for i = 1:N
        for j = (i+1):N
            scores(i,j) = getscore(shreds(i,:),shreds(j,:),false);
            scores(j,i) = scores(i,j);
        end
    end
    
%% Sort shreds into groups by discretization
    groups = zeros(N,1);
    n_groups = 0;
    while any(groups == 0)
        ungrouped = find(groups == 0);
        k = ungrouped(randi(length(ungrouped)));
        bin_nums = discretize(scores(k,:),20);
        strongest = find(bin_nums >= 18);
        g_already = groups(strongest);
        strongest(g_already ~= 0) = [];
        if length(strongest) > max_group_size
            strongest = strongest(1:max_group_size);
        end 
        n_groups = n_groups + 1;
        groups([k,strongest]) = n_groups;
    end

%% Form the groups
    g_orders = cell(n_groups,1); % list of shreds in the group in order
    g_pairs = zeros(n_groups,2); % shreds on the left and right edge
    g_shreds = cell(n_groups,3); % shreds consisting of just the edge pairs
    for g = 1:n_groups
        g_list = list(groups==g);
        g_order = ordershreds(shreds(g_list,:));
        g_orders{g} = g_list(g_order);
        g_pairs(g,:) = g_orders{g}([1,end]);
        g_shreds{g,1} = [shreds{g_pairs(g,1),1} shreds{g_pairs(g,2),1}];
        g_shreds{g,2} = [shreds{g_pairs(g,1),2} shreds{g_pairs(g,2),2}];
        g_shreds{g,3} = [shreds{g_pairs(g,1),3} shreds{g_pairs(g,2),3}];
        g_shreds{g,4} = shreds{g_pairs(g,1),4};
        g_shreds{g,4}(shreds{g_pairs(g,2),4}) = 1;
    end
    
end

function [g_orders,g_shreds] = regroupshreds(g_orders_p,g_shreds_p,shreds)
% Given a set of groups, merge pairs of groups together and re-sort
    Ng = size(g_shreds_p,1);
    
%% Get similarity scores
    scores = NaN(Ng,Ng);
    for i = 1:Ng
        for j = (i+1):Ng
            scores(i,j) = getscore(g_shreds_p(i,:),g_shreds_p(j,:),false);
            scores(j,i) = scores(i,j);
        end
    end
    
%% Pair off groups 
    groups = zeros(Ng,1);
    n_groups = 0;
    while any(groups == 0)
        ungrouped = find(groups == 0);
        if numel(ungrouped) > 1
            k = ungrouped(randi(length(ungrouped)));
            [~,strongest] = max(scores(k,ungrouped));
            strongest = ungrouped(strongest);
        else
            strongest = [];
            k = ungrouped;
        end
        n_groups = n_groups + 1;
        groups([k,strongest]) = n_groups;
    end
    
%% Merge groups and resort shreds in groups
    g_orders = cell(n_groups,1); % list of shreds in the group in order
    g_pairs = zeros(n_groups,2); % shreds on the left and right edge
    g_shreds = cell(n_groups,3); % shreds consisting of just the edge pairs
    for g = 1:n_groups
        g_list = [];
        for idx = (find(groups==g))'
            g_list = [g_list g_orders_p{idx}];
        end
        g_order = ordershreds(shreds(g_list,:));
        g_orders{g} = g_list(g_order);
        g_pairs(g,:) = g_orders{g}([1,end]);
        g_shreds{g,1} = [shreds{g_pairs(g,1),1} shreds{g_pairs(g,2),1}];
        g_shreds{g,2} = [shreds{g_pairs(g,1),2} shreds{g_pairs(g,2),2}];
        g_shreds{g,3} = [shreds{g_pairs(g,1),3} shreds{g_pairs(g,2),3}];
        g_shreds{g,4} = shreds{g_pairs(g,1),4};
        g_shreds{g,4}(shreds{g_pairs(g,2),4}) = 1;
    end
end

function [order,score,x_groups,g_order] = ordergroups(shreds,g_orders,g_shreds)
% Get the best order of the groups

    N = size(shreds,1);
    n_groups = size(g_shreds,1);
    order = [];
    g_order = ordershreds(g_shreds);
    x_groups = zeros(n_groups,1);
    for g = 1:n_groups
        order = [order g_orders{g_order(g)}];
        for o = order
            x_groups(g) = x_groups(g) + size(shreds{o,3},2);
        end
    end
    score = scoreorder(shreds,order);
    
end

function [order,score,flipped,shreds,x_groups] = iterflip(shreds,g_orders,g_shreds,flipped)
% Randomly test different orientations of groups and iterate 
    
    [order,score,x_groups] = ordergroups(shreds,g_orders,g_shreds);
    n_groups =  size(g_shreds,1);
    iter = 1;
    while iter < 30
        g = randi([1,n_groups],1); % choose random group
        [shreds_f,g_orders_f,g_shreds_f,flipped_f] = flipgroup(shreds,g_orders,g_shreds,flipped,g);
        [order_f,score_f,x_groups_f] = ordergroups(shreds_f,g_orders_f,g_shreds_f);
        if score_f - score > 0
            % keep the flip
            shreds = shreds_f;
            g_orders = g_orders_f;
            g_shreds = g_shreds_f;
            flipped = flipped_f;
            x_groups = x_groups_f;
            order = order_f;
            score = score_f;
        end
        iter = iter + 1;
    end
    
end

function [shreds,g_orders,g_shreds,flipped] = flipgroup(shreds,g_orders,g_shreds,flipped,g)
    
    % Flip the shreds
    idx = g_orders{g};
    for i = idx
        shreds{i,1} = rot90(shreds{i,1},2);
        shreds{i,2} = rot90(shreds{i,2},2);
        shreds{i,3} = rot90(shreds{i,3},2);
        flipped(i) = ~flipped(i);
    end
    % Reverse the order
    g_orders{g} = fliplr(g_orders{g});
    % Flip the group edge shreds
    g_shreds{g,1} = rot90(g_shreds{g,1},2);
    g_shreds{g,2} = rot90(g_shreds{g,2},2);
    g_shreds{g,3} = rot90(g_shreds{g,3},2);

end

function [order,scores] = ordershreds(shreds)
% Sort shreds into optimal order based on similarity score
    N = size(shreds,1);
    scores = NaN(N,N);
    if N == 1
        order = 1;
        return
    end
    
%% Get similarity scores
    for i = 1:N
        for j = 1:N
            if i == j
                continue
            end
            scores(i,j) = getscore(shreds(i,:),shreds(j,:),true);
        end
    end
    
%% Get the optimal order
    order_scr = zeros(N,N);
    orders = nan(N,N);
    for start = 1:N
        % Find order for each possible starting shred
        order_idx = 2;
        cur = start;
        orders(start,1) = start;
        while any(isnan(orders(start,:)))
            [m,idx] = maxk(scores(cur,:),N-1);
            for i = 1:(N-1)
                if ~any(orders(start,:) == idx(i))
                    orders(start,order_idx) = idx(i);
                    order_scr(start,order_idx) = m(i);
                    order_idx = order_idx + 1;
                    cur = idx(i);
                    break;
                end
            end
        end
    end
    % Choose the most optimal permutation
    [~,c] = max(sum(order_scr,2));
    order = orders(c,:);
    
end

function score = scoreorder(shreds,order)
% Get the similarity score of an order of shreds
    N = size(shreds,1);
    score = 0;
    for i = 1:(N-1)
        score = score + getscore(shreds(order(i),:),shreds(order(i+1),:),true);
    end
end

function score = getscore(shred1,shred2,right)
% Get the similarity score of two shreds

    score_hist = comparehoriz(shred1{4},shred2{4});
    score_color = comparecolor(shred1{1},shred2{1});
    
    if right
        score_right = compareright(~shred1{3},~shred2{3});
        score = score_hist + score_color + score_right;
    else
        score = score_hist + score_color;
    end

end

function score = compareright(b1,b2)
% Get the match score of two shreds 

    x1 = b1(:,end,:); % rightmost of shred1
    y1 = b2(:,1,:); % leftmost of shred2
    x2 = b1(:,end-1,:); % 2nd from right of shred2
    y2 = b2(:,2,:); % 2nd of shred 2
    
    a = dot(x1,y1);
    b = dot(y1,(1-x1));
    c = dot(x1,(1-y1));
    d = dot((1-x1),(1-y1));
    
    score11 = (a+d)/(a+b+c+d);

    a = dot(x1,y2);
    b = dot(y2,(1-x1));
    c = dot(x1,(1-y2));
    d = dot((1-x1),(1-y2));
    
    score12 = (a+d)/(a+b+c+d);
    
    a = dot(x2,y1);
    b = dot(y1,(1-x2));
    c = dot(x2,(1-y1));
    d = dot((1-x2),(1-y1));
    
    score21 = (a+d)/(a+b+c+d);
    
    a = dot(x2,y2);
    b = dot(y2,(1-x2));
    c = dot(x2,(1-y2));
    d = dot((1-x2),(1-y2));
    
    score22 = (a+d)/(a+b+c+d);
    
    score = max([score11 score12 score21 score22]);
    
end

function score = comparecolor(s1,s2)
% Get the color similarity of two shreds

    x = mean(s1,2);
    y = mean(s2,2);
    
    c_norm = (x(:,1,1) - y(:,1,1)).^2 + (x(:,1,2) - y(:,1,2)).^2 ...
        + (x(:,1,3) - y(:,1,3)).^2;
    
    score = 1 - mean(c_norm);
    
end

function score = comparehoriz(h1,h2)
% Get the L1-norm of two horizontal projections

    score = 1 - (sum(abs(h1-h2))./numel(h1));

end

function [shreds,flipped] = preprocess(shred_directory)
% Pre-process all of the shreds
    N = 100;
    listing = dir([shred_directory,'*.png']);
    flipped = false(N,1);
    shreds = cell(N,4);
    for i = 1:N
        s = im2double(imread([shred_directory,listing(i).name]));
        s_g = medfilt2(rgb2gray(s));

        % Use a hough transform to properly align each shred with the
        % vertical
        SE = strel('square',2');
        box1 = s_g > 0.001;
        edges = (imdilate(box1,SE) ~= imerode(box1,SE));
        [H,T,~] = hough(edges,'Theta',-10:0.01:10,'RhoResolution',0.5);
        P = houghpeaks(H,4);
        theta = mean(T(P(:,2)));
        box2 = imrotate(box1,theta);

        % Crop the black regions from the shred, cut off an extra pixel on
        % each side for errors from rotation
        col = sum(box2,1);
        h = size(box2,1)-20;
        first = find(col>=h,1,'first');
        last = find(col>=h,1,'last');

        % Final rotate and crop
        s = imrotate(s,theta);
        s = s(3:(end-2),first:last,:);

        % Add padding so all shreds are the same height
        if i == 1
            h_shred1 = size(s,1) + 5;
            q = round(h_shred1/2);
        end
        p = h_shred1 - size(s,1);
        s = padarray(padarray(s,[5 0 0],'symmetric','pre'),[p 0 0],'symmetric','post');

        % Compare to the top quarter of the first shred for orientation
        if i > 1
            err = immse(shreds{1,1}(1:q,1:5,:),s(1:q,1:5,:));
            s_f = rot90(s,2);
            if immse(shreds{1,1}(1:q,1:5,:),s_f(1:q,1:5,:)) < err
                flipped(i) = true;
                s = s_f;
            end
        end
        
        s = imgaussfilt(s,1);
        s_g = rgb2gray(s);
        
        % Get binarized shreds
        s_b = adaptivethreshold(s_g,size(s_g,2),0.04);
        s_b = imclose(s_b,strel('square',3));
        s_b = medfilt2(s_b);

        % Get thresholded histograms
        SE = strel('square',5);
        th = 8;
        s_h = sum(~s_b,2);
        s_h = imclose(s_h > th,SE);   
        
        shreds{i,1} = s;
        shreds{i,2} = s_g;
        shreds{i,3} = s_b;
        shreds{i,4} = s_h;
    end
end

function page = makepage(shreds,order,set)
% Compile the page
% Set 1: color
% Set 2: grayscale
% Set 3: binary

    sz = size(shreds{order(1),set});
    if set == 1
        page = zeros(sz(1),sz(2)*100,3);
    else
        page = zeros(sz(1),sz(2)*100);
    end
    next = 1;
    for i = 1:numel(order)
        sz = size(shreds{order(i),set});
        rng = next:(next-1+sz(2));
        page(1:sz(1),rng,:) = shreds{order(i),set};
        next = rng(end) + 1;
    end
    if next <= size(page,2)
        page(:,next:end,:) = [];
    end
end

function bw = adaptivethreshold(IM,ws,C,tm)
%ADAPTIVETHRESHOLD An adaptive thresholding algorithm that seperates the
    %foreground from the background with nonuniform illumination.
    %  bw=adaptivethreshold(IM,ws,C) outputs a binary image bw with the local 
    %   threshold mean-C or median-C to the image IM.
    %  ws is the local window size.
    %  tm is 0 or 1, a switch between mean and median. tm=0 mean(default); tm=1 median.
    %
    %  Contributed by Guanglei Xiong (xgl99@mails.tsinghua.edu.cn)
    %  at Tsinghua University, Beijing, China.
    %
    %  For more information, please see
    %  http://homepages.inf.ed.ac.uk/rbf/HIPR2/adpthrsh.htm
    if (nargin<3)
        error('You must provide the image IM, the window size ws, and C.');
    elseif (nargin==3)
        tm=0;
    elseif (tm~=0 && tm~=1)
        error('tm must be 0 or 1.');
    end
    IM=mat2gray(IM);
    if tm==0
        mIM=imfilter(IM,fspecial('average',ws),'replicate');
    else
        mIM=medfilt2(IM,[ws ws]);
    end
    sIM=mIM-IM-C;
    bw=im2bw(sIM,0);
    bw=imcomplement(bw);
end