%%  Audio Declipping Example
%   
%   Audio declipping is a problem proposed in Audio Inpaining Toolbox and
%   in [2]. This is an example of solving the problem with fast omp using
%   Gabor dictionary and ompGabor implemented in SMALLbox [1].
%
%   [1] I. Damnjanovic, M. E. P. Davies, and M. P. Plumbley "SMALLbox - an 
%   evaluation framework for sparse representations and dictionary 
%   learning algorithms," V. Vigneron et al. (Eds.): LVA/ICA 2010, 
%   Springer-Verlag, Berlin, Germany, LNCS 6365, pp. 418-425
%   [2] A. Adler, V. Emiya, M. G. Jafari, M. Elad, R. Gribonval, and M. D.
%   Plumbley, “Audio Inpainting,” submitted to IEEE Trans. Audio, Speech,
%   and Lang. Proc., 2011, http://hal.inria.fr/inria-00577079/en/.

%
%   Centre for Digital Music, Queen Mary, University of London.
%   This file copyright 2011 Ivan Damnjanovic.
%
%   This program is free software; you can redistribute it and/or
%   modify it under the terms of the GNU General Public License as
%   published by the Free Software Foundation; either version 2 of the
%   License, or (at your option) any later version.  See the file
%   COPYING included with this distribution for more information.
%
%%

clear all;
%   Defining the solvers to test in Audio declipping scenario

% First solver omp2 - DCT+DST dictionary  with no additional constraints

SMALL.solver(1) = SMALL_init_solver('ompbox', 'omp2', '', 0);
SMALL.solver(1).add_constraints = 0;

% Second solver omp2 - DCT+DST dictionary  with additional constraints

SMALL.solver(2) = SMALL_init_solver('ompbox', 'omp2', '', 0);
SMALL.solver(2).add_constraints = 1;

% Third solver omp2 - Gabor dictionary with no additional constraints

SMALL.solver(3) = SMALL_init_solver('ompbox', 'omp2Gabor', '', 0);
SMALL.solver(3).add_constraints = 0;

% Fourth solver omp2- Gabor dictionary with no additional constraints

SMALL.solver(4) = SMALL_init_solver('ompbox', 'omp2Gabor', '', 0);
SMALL.solver(4).add_constraints = 1;

%   Defining the Problem structure

SMALL.Problem = generateAudioDeclippingProblem('', 0.6, 256, 0.5, @wRect, @wSine, @wRect, @Gabor_Dictionary, 2);

for idxSolver = 1:4
    
    fprintf('\nStarting Audio Declipping of %s... \n', SMALL.Problem.name);
    fprintf('\nClipping level %s... \n', SMALL.Problem.clippingLevel);
    
    start=cputime;
    tStart=tic;
    
    error2=0.001^2;
    coeffFrames = zeros(SMALL.Problem.p, SMALL.Problem.n);
    
    
    
    for i = 1:SMALL.Problem.n
        
        idx = find(SMALL.Problem.M(:,i));
        if size(idx,1)==SMALL.Problem.m
            continue
        end
        Dict = SMALL.Problem.B(idx,:);
        wDict = 1./sqrt(diag(Dict'*Dict));
        
        SMALL.Problem.A  = Dict*diag(wDict);
        
        SMALL.Problem.b1 = SMALL.Problem.b(idx,i);
        
        
        
        % set solver parameters
        
        SMALL.solver(idxSolver).param=struct(...
            'epsilon', error2*size(idx,1),...
            'maxatoms', 128, ...
            'profile', 'off');
        
        % Find solution
        
        SMALL.solver(idxSolver)=SMALL_solve(SMALL.Problem, SMALL.solver(idxSolver));
        
        % Refine solution with additional constraints for declipping scenario
        
        if (SMALL.solver(idxSolver).add_constraints)
            SMALL.solver(idxSolver)=CVX_add_const_Audio_declipping(...
                SMALL.Problem, SMALL.solver(idxSolver), i);
        end
        
        %%
        coeffFrames(:,i) = wDict .* SMALL.solver(idxSolver).solution;
        
        
    end
    
    %%   Set reconstruction function
    
    SMALL.Problem.reconstruct=@(x) AudioDeclipping_reconstruct(x, SMALL.Problem);
    reconstructed=SMALL.Problem.reconstruct(coeffFrames);
    SMALL.Problem = rmfield(SMALL.Problem, 'reconstruct');
    tElapsed=toc(tStart);
    
    SMALL.solver(idxSolver).time = cputime - start;
    fprintf('Solver %s finished task in %2f seconds (cpu time). \n', ...
        SMALL.solver(idxSolver).name, SMALL.solver(idxSolver).time);
    fprintf('Solver %s finished task in %2f seconds (tic-toc time). \n', ...
        SMALL.solver(idxSolver).name, tElapsed);
    
    
    
    %% Plot results
    
    xClipped = SMALL.Problem.clipped;
    xClean   = SMALL.Problem.original;
    xEst1    = reconstructed.audioAllSamples;
    xEst2    = reconstructed.audioOnlyClipped;
    fs       = SMALL.Problem.fs;
    
    figure
    hold on
    plot(xClipped,'r')
    plot(xClean)
    plot(xEst2,'--g')
    plot([1;length(xClipped)],[1;1]*[-1,1]*max(abs(xClipped)),':r')
    legend('Clipped','True solution','Estimate')
end

% % Normalized and save sounds
% normX = 1.1*max(abs([xEst1(:);xEst2(:);xClean(:)]));
% L = min([length(xEst2),length(xEst1),length(xClean),length(xEst1),length(xClipped)]);
% xEst1 = xEst1(1:L)/normX;
% xEst2 = xEst2(1:L)/normX;
% xClipped = xClipped(1:L)/normX;
% xClean = xClean(1:L)/normX;
% wavwrite(xEst1,fs,[expParam.destDir 'xEst1.wav']);
% wavwrite(xEst2,fs,[expParam.destDir 'xEst2.wav']);
% wavwrite(xClipped,fs,[expParam.destDir 'xClipped.wav']);
% wavwrite(xClean,fs,[expParam.destDir 'xClean.wav']);
