% -------------------------------------------------------------------------
% MSc Project - MSc in Communications and Signal Processing
% Department of Electrical and Electronic Engineering
% Imperial College London, 2011
%
% Supervisor  : Dr Pier Luigi Dragotti
% Contributor : Andriy Gelman
% Author      : Jon Onativia Bravo
%
% File       : benchmark.m
% -------------------------------------------------------------------------


clear all
close all
clc

%% MATLAB data files
% files = {'animal_farm.mat', 'cones.mat', 'teddy.mat', 'tsukuba_extracted4.mat'};
files = {'teddy.mat', 'tsukuba_extracted4.mat'};
% files = {'animal_farm_32.mat'};

max_distance = 3;
max_offset   = 31;
max_depth    = 1;
lambda       = 0.1;
num_it       = 1;

%% Check if the folder where the results will be saved exists
folder_name = 'results';
if exist(folder_name,'dir') == 0
    mkdir(folder_name);
end

%% Loop over all files
for file_name_cell = files
    
    file_name         = file_name_cell{1};
    [~, data_name, ~] = fileparts(file_name);

    % Load data

    % Disparity vector
    load(file_name, 'D');
    % Images array
    load(file_name, 'in');
    % Labels array
    load(file_name, 'label');

    % Size of the data
    num_images = size(in, 3);

    % Obtain the stack of labels
    layers_stack = generate_layer_stack(label, D);
    
    % Labels parameters
    [height, width, num_layers] = size(layers_stack);
    approx_layers               = zeros(height, width, num_layers-1);
    
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    tic
    for it = 1 : num_it
    % Construct the contours tree
    label_verts_tree = cell(num_layers-1, 1);
    for l = 1 : num_layers-1
        label_verts_tree{l} = estimate_contours_tree(layers_stack(:,:,l), max_distance, max_offset);
    end % end of l = 1 : num_layers

    % Compress the bitstream
    [bitstream, layer_limits] = tree2bitstream_v3(label_verts_tree, max_offset, width, height);
    field_bits = log2(max_offset + 1) + 1;
    [substreams] = arithm_enc_stream(bitstream, layer_limits, field_bits);
    
    % Recontruct the layers
    distortions     = zeros(1, max_depth);
    compressed_bits = 0;
    for depth = 1 : max_depth
        for l = 1 : num_layers-1
            verts                = tree2contour(label_verts_tree{l}, depth);
            approx_layers(:,:,l) = contour2shape(verts, width, height);
        end
    end
    end
    t_cont_tree = toc;
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    tic
    for it = 1 : num_it
        [quad_stack, quad_bits] = quad_tree_lossy_segmentation(layers_stack(:,:,1:num_layers-1), lambda);
    end
    t_quad_tree = toc;
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    
    %% Display results
    t_cont_tree = t_cont_tree / num_it;
    t_quad_tree = t_quad_tree / num_it;
    
    disp(['******** ' file_name ' ********'])
    disp(['contour tree = ' num2str(t_cont_tree)]);
    disp(['quad tree    = ' num2str(t_quad_tree)]);
    
    %% Save results
    results_file_name = [folder_name '/bench_' file_name];
    save(results_file_name, 't_cont_tree', 't_quad_tree');
end
    
% Play a sound when finished
fs = 8000;
Ts = 1 / 8000;
t  = 0:Ts:2;
fo = 220;
so = sin(2*pi*fo*t);

soundsc(so, fs)
