% -------------------------------------------------------------------------
% 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 = {'cones.mat', 'teddy.mat', 'tsukuba_extracted4.mat'};
files = {'animal_farm.mat'};


% 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
    [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
    
    %lambdas                 = [0.01 0.1 1 10];
    lambdas = 0.01 * ones(1, 1);
    tic;
    [quad_stack, quad_bits] = quad_tree_lossy_segmentation(layers_stack(:,:,1:num_layers-1), lambdas);
    t_elapsed = toc;
    
    disp(['Elapsed time per lambda: ' num2str(t_elapsed / length(lambdas))])
    disp([' \lambda = ' num2str(lambdas(1))])
    
    %% Compute distortions
    num_points = length(lambdas);
    quad_dists = zeros(num_points, 1);
%     for ith_point = 1 : num_points
%         approx_layers         = quad_stack{ith_point};
%         quad_dists(ith_point) = get_distortion(layers_stack(:,:,1:num_layers-1), approx_layers);
%         
%         disp(['Lambda                    : ' num2str(lambdas(ith_point))]);
%         disp(['Bits                      : ' num2str(quad_bits{ith_point})]);
%         disp(['Obtained distortion       : ' num2str(quad_dists(ith_point))]);
        
%         fprintf(fid, '\n');
%         fprintf(fid, 'Depth                     : %f\n', depth);
%         fprintf(fid, 'Bits                      : %f\n', req_bits(depth));
%         fprintf(fid, 'Obtained distortion       : %f\n', distortions(depth));
%         fprintf(fid, 'Quad bits                 : %f\n', quad_bits(depth));
%         fprintf(fid, 'Quad distortion           : %f\n', quad_dists(depth));
        
%     end
    
    %% Save results
%     results_file_name = [folder_name '/quad_tree_' file_name];
%     save(results_file_name, 'quad_stack', 'lambdas', 'quad_bits', 'quad_dists');
    
    
    %%
    
%     fclose(fid);
    
end % end of file_name_cell = files


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