clear;

%% File Paths
fname = 'AllCIFAR.h5';
addpath(genpath('/Google Drive/Research/code/matlab/toolboxes/WaveletSoftware'));

%% Read in the HDF5 information
hi = hdf5info(fname);
X = hdf5read(hi.GroupHierarchy(1).Datasets(1));

%% Initializations
N = 50000;                           % Number of images to transform
imsize = [32 32];                   % Image dimensions
mean_removal = true;                % Remove mean of individual images?
L = log2(32);                       % Maximum number of decompositions
sq = @(x_) reshape(x_,imsize);      % Helper function for display
[af,sf] = farras;                   % Load in wavelet
PctRetain = 0.05;                    % Percent of w.coeffs. to retain
K = floor(PctRetain*prod(imsize));  % # of w.coeffs. to retain

W = zeros(size(X,1),N);             % Wavelet-ized dataset
WT = W;                             % Trunc. Wavelet-ized dataset

%% Loop over Dataset
for i=1:N
    im = sq(X(:,i))';                   % convert to 2d
    if mean_removal
        im = im - mean(im(:));          % subract mean
    end

    % Wavelet transform
    w = dwt2D(im,5,af);
    w = wcell2mat(w);
    w = w(:);

    % Truncation
    wt = zeros(size(w));
    [~,wo] = sort(abs(w(:)),'descend');
    wt(wo(1:K)) = w(wo(1:K));

    % Assign back to dataset
    W(:,i) = w(:);
    WT(:,i) = wt(:);
    if mod(i,500)==0
        fprintf('%0.2f%% complete.\n',i/N*100);
    end
end

% Write out on the wavelet features
hdf5write('WaveletCIFAR.h5','X',W);

%% Lets compare the wavelet datasets
figure(1);
colormap(parula);
subplot(2,3,1);
    imagesc(abs(W));
    title('Wavelet Coefficients');
    colorbar();
subplot(2,3,2);
    [x,a] = hist(abs(W(:)),100);
    stem(a,x,'-','Marker','.');
    title('Wavelet Magnitude Distribution');
    set(gca,'YScale','log');
subplot(2,3,3);
    imagesc(reshape(mean(abs(W),2),imsize));
    axis image;
    title('Avg. Wavelet Magnitude');
    colorbar();
subplot(2,3,4);
    imagesc(abs(WT));
    title('Truncated Wavelet Coefficients');
    colorbar();
subplot(2,3,5);
    [x,a] = hist(abs(WT(:)),100);
    stem(a,x,'-','Marker','.');
    title('Trunc. Wavelet Magnitude Distribution');
    set(gca,'YScale','log');
subplot(2,3,6);
    imagesc(reshape(mean(abs(WT),2),imsize));
    axis image;
    title('Avg. Trunc. Wavelet Magnitude');
    colorbar();