% -------------------------------------------------------------------------
% 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_statistics.m
% -------------------------------------------------------------------------


clear all
close all
clc


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


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

% Entropy encoding parameters
bitstreams    = cell(1, num_files*max_depth);
bits_per_word = 6;
fill_bits_len = ceil(log2(bits_per_word));

% 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);

    % Extract layers and merge

    labels     = unique(label(:,:,:));
    labels     = labels(labels > 0);
    num_layers = length(labels);
    disparity  = ceil( (0:num_images-1)' * D' );

    % The disparity measures the displacement of a layer with respect to the
    % first image (EPI images, camera only moves in the horizontal axis)
    pos_disp = max(disparity(:));
    neg_disp = -min(disparity(:));
    ext_cols = c + pos_disp + neg_disp;

    % The layers_stack will contain one label per image, adding up the layer
    % among all the images disparity compensated
    layers_stack = zeros(r, ext_cols, num_layers);

    for l = 1 : num_layers
        ith_layer                      = zeros(r, ext_cols);

        for i = 1 : num_images
            cur_disp  = disparity(i, l);
            cur_label = label(:,:,i);

            % Indices of the location corresponding to the current layer in the
            % current image
            ind_orig = (cur_label == l);

            % Indices corresponding to the location in the extended layer
            ind_dest  = false(r, ext_cols);
            first_col = neg_disp + 1 + cur_disp;
            last_col  = first_col + c - 1;
            ind_dest(:, first_col:last_col) = ind_orig;

            % Copy 
            ith_layer(ind_dest) = cur_label(ind_orig) / l;
        end

        layers_stack(:,:,l) = ith_layer;    
    end
    
    % ====> Data is available here in "layer_stack" variable
    
    % Labels parameters
    [height, width, num_layer] = size(layers_stack);
    approx_layers              = zeros(height, width, num_layers-1);
    
    % Compute the contour trees for each layer
    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_distances, max_offsets);
    end % end of l = 1 : num_layers

    % Obtain the bitstreams
    [bitstream, layer_limits] = tree2bitstream_v2(label_verts_tree, max_offsets, width, height);
    
    % Reorder subtstreams
    for depth = 1 : max_depth
        if depth == 1
            start_idx = 1;
        else
            start_idx = layer_limits(depth);
        end
        if depth == max_depth
            end_idx = length(bitstream);
        else
            end_idx = layer_limits(depth+1) - 1;
        end
        
        sub_stream = bitstream(start_idx:end_idx);
        len        = length(sub_stream) + fill_bits_len;

        fill_bits  = bits_per_word - mod(len, bits_per_word);
        sub_stream = [ dec2bin(fill_bits, fill_bits_len), sub_stream, false(1, fill_bits) ];
        
        bitstreams{(ith_file-1)*max_depth + depth} = sub_stream;
    end
    
    ith_file = ith_file + 1;
    
end % end of file_name_cell = files

%% 'bitstreams' variable contains all the bitstreams of all the files
bits = cell2mat(bitstreams);

[symb, prob] = get_symbols(bits, bits_per_word);
dict         = huffmandict(symb, prob);
entropy      = - prob * log2(prob');

disp(['Bits per word = ' num2str(bits_per_word)])
disp(['Entroypy      = ' num2str(entropy)])

for ith_file = 1 : num_files
    file_name = files{ith_file};
    disp(['File : ' file_name])
    
    accum_bits = 0;
    for depth = 1 : max_depth
        substream = bitstreams{(ith_file-1)*max_depth + depth};
        substream = apply_dictionary(substream, dict);
        
        accum_bits = accum_bits + length(substream);
        disp(num2str(accum_bits));
    end
    
end

disp('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)
