% -------------------------------------------------------------------------
% 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        : compute_artihm_tree.m
% -------------------------------------------------------------------------


clear all
close all


% MATLAB data files
% files = {'animal_farm.mat', 'cones.mat', 'teddy.mat', 'tsukuba_extracted4.mat'};
files = {'tsukuba_extracted4.mat'};
num_files = length(files);


% Tree parameters
% max_distances = [ 3  2 1];
max_distances = [10  5  2];
max_offsets   = [31 15 7];
max_depth     = length(max_distances);


% Loop over all files
ith_file = 1;
for file_name_cell = files
    
    file_name         = file_name_cell{1};
    [~, data_name, ~] = fileparts(file_name);
    disp(['File : ' 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
    [r, c, num_images] = size(in);

    % 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);
    distortions                 = zeros(1, max_depth);
    
    % Compute the contour trees for each layer
    label_verts_tree = cell(num_layers, 1);
    for l = 1 : num_layers-1
        if l < num_layers
            label_verts_tree{l} = estimate_contours_tree(layers_stack(:,:,l), max_distances, max_offsets);
        else
            last_layer          = padarray(layers_stack(:,:,l), [1 1]);
            label_verts_tree{l} = estimate_contours_tree(last_layer, [20 20 20], max_offsets);
        end
    end % end of l = 1 : num_layers

    % Obtain the bitstreams
    [bitstream, layer_limits] = tree2bitstream_v3(label_verts_tree, max_offsets, width, height);
    
    field_bits = log2(max_offsets + 1) + 1;
    
    % Compress the bitstream
    [substreams] = arithm_enc_stream(bitstream, layer_limits, field_bits);
    
    compressed_bits = 0;
    for depth = 1 : length(substreams)
        compressed_bits = compressed_bits + length(substreams{depth})*8;
        
        
        % Compute distortion
        for l = 1 : num_layers-1
            if l < num_layers
                verts                = tree2contour(label_verts_tree{l}, depth);
                approx_layers(:,:,l) = contour2shape(verts, width, height);
            else
                verts                = tree2contour(label_verts_tree{l}, depth);
                last_layer           = contour2shape(verts, width+2, height+2);
                approx_layers(:,:,l) = last_layer(2:end-1,2:end-1);
            end
        end
        distortions(depth) = get_distortion(layers_stack(:,:,1:num_layers-1), approx_layers);
        
        disp(['Depth : ' num2str(depth) ' - bits : ' num2str(compressed_bits) ' - distortion : ' num2str(distortions(depth))])
    end
    
    % Reconstruct the original stream
    rec_bitstream = arithm_dec_stream(substreams, field_bits);
    
end