clear

%% Parameteres

% Dictionary learning parameters
toolbox   = 'TwoStepDL'; %dictionary learning toolbox
dicUpdate = {'ksvd','mailhe'}; %dictionary updates
iterNum   = 20; %number of iterations

% Test signal parameters
signal    = audio('music03_16kHz.wav'); %audio signal
blockSize = 256; %size of audio frames
dictSize  = 512; %number of atoms in the dictionary
overlap   = 0.5; %overlap between consecutive frames
sigma     = 1e6; %snr of noise (set to be negligible so that the problem becomes approximation rather than denoising)
percActiveAtoms = 5; %percentage of active atoms

% Dependent parameters
nActiveAtoms = fix(blockSize/100*percActiveAtoms); %number of active atoms
epsilon      = 1/sigma; %error constraint for sparse representation step (corresponds to noise applied to signals)
minMu = sqrt((dictSize-blockSize)/(blockSize*(dictSize-1))); %target coherence (based on coherence lower bound)
minCoherence = [0.1:0.1:1];

% Initial dictionaries
dctDict = dictionary('dct',blockSize,dictSize);
dctDict = dctDict.phi;
gaborParam = struct('N',blockSize,'redundancyFactor',2,'wd',@rectwin);
gaborDict = Gabor_Dictionary(gaborParam);

%% Generate audio denoising problem with low noise (audio representation)
SMALL.Problem = generateAudioDenoiseProblem(signal.s,[],blockSize,...
	dictSize,overlap,sigma); % generate representation problem
SMALL.Problem.b1 = SMALL.Problem.b; % copy signals from training set b to test set b1 (needed for later functions)

% omp2 sparse representation solver
ompParam = struct('X',SMALL.Problem.b,'epsilon',epsilon,'maxatoms',nActiveAtoms); %parameters
solver = SMALL_init_solver('ompbox','omp2',ompParam,false); %solver structure


%% Test
name = dicUpdate{1}; %use ksvd update
SMALL.DL(1:36) = SMALL_init_DL(toolbox,name); %create dictionary learning structures

% learn with random initialisation and no decorrelation
SMALL.DL(1).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','none'); %parameters for the dictionary learning
SMALL.DL(1) = SMALL_learn(SMALL.Problem,SMALL.DL(1)); %learn dictionary
%save('SMALL','SMALL');

% learn with random initialisation and mailhe decorrelation
for iMu=1:10
	SMALL.DL(1+iMu).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
		'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
		'decFcn','mailhe','coherence',minCoherence(iMu)); %parameters for the dictionary learning
	SMALL.DL(1+iMu) = SMALL_learn(SMALL.Problem,SMALL.DL(1+iMu)); %learn dictionary
	%save('SMALL','SMALL');
end

% learn with random initialisation and tropp decorrelation
SMALL.DL(12).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','tropp','coherence',minCoherence); %parameters for the dictionary learning
SMALL.DL(12) = SMALL_learn(SMALL.Problem,SMALL.DL(12)); %learn dictionary

% Learn with dct initialisation and no decorrelation
SMALL.DL(13).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','none','initdict',dctDict); %parameters for the dictionary learning
SMALL.DL(13) = SMALL_learn(SMALL.Problem,SMALL.DL(13)); %learn dictionary

% learn with dct initialisation and mailhe decorrelation
for iMu=1:10
	SMALL.DL(13+iMu).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
		'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
		'decFcn','mailhe','coherence',minCoherence(iMu),'initdict',dctDict); %parameters for the dictionary learning
	SMALL.DL(13+iMu) = SMALL_learn(SMALL.Problem,SMALL.DL(13+iMu)); %learn dictionary
	%save('SMALL','SMALL');
end

% learn with dct initialisation and tropp decorrelation
SMALL.DL(24).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','tropp','coherence',minCoherence,'initdict',dctDict); %parameters for the dictionary learning
SMALL.DL(24) = SMALL_learn(SMALL.Problem,SMALL.DL(24)); %learn dictionary


% Learn with gabor initialisation and no decorrelation
SMALL.DL(25).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','none','initdict',gaborDict); %parameters for the dictionary learning
SMALL.DL(25) = SMALL_learn(SMALL.Problem,SMALL.DL(25)); %learn dictionary

% learn with gabor initialisation and mailhe decorrelation
for iMu=1:10
	SMALL.DL(25+iMu).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
		'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
		'decFcn','mailhe','coherence',minCoherence(iMu),'initdict',gaborDict); %parameters for the dictionary learning
	SMALL.DL(25+iMu) = SMALL_learn(SMALL.Problem,SMALL.DL(25+iMu)); %learn dictionary
	%save('SMALL','SMALL');
end

% learn with gabor initialisation and tropp decorrelation
SMALL.DL(36).param = struct('data',SMALL.Problem.b,'Tdata',nActiveAtoms,...
	'dictsize',dictSize,'iternum',iterNum,'memusage','high','solver',solver,...
	'decFcn','tropp','coherence',minCoherence,'initdict',gaborDict); %parameters for the dictionary learning
SMALL.DL(36) = SMALL_learn(SMALL.Problem,SMALL.DL(36)); %learn dictionary

%% Evaluate coherence and snr of representation for the various methods
sigNoiseRatio = zeros(36,1);
mu = zeros(36,1);
for i=1:36
	SMALL.Problem.A = SMALL.DL(i).D;
	tempSolver = SMALL_solve(SMALL.Problem,solver);
	sigNoiseRatio(i) = snr(SMALL.Problem.b,SMALL.DL(i).D*tempSolver.solution);
	dic(i) = dictionary(SMALL.DL(i).D);
	mu(i) = dic(i).coherence;
end

%% Plot results
minMu = sqrt((dictSize-blockSize)/(blockSize*(dictSize-1)));

figure,
%subplot(3,1,1)
hold on, grid on
title('Data Initialisation')
plot([1 1],[0 25],'k-');
plot(mu(1),sigNoiseRatio(1),'ks');
plot(mu(12),sigNoiseRatio(12),'kd');
plot(mu(2:11),sigNoiseRatio(2:11),'k*-');
plot([minMu minMu],[0 25],'k--')
set(gca,'YLim',[0 25],'XLim',[0 1.4]);
legend({'\mu_{max}','K-SVD','Grassmannian','INK-SVD','\mu_{min}'});
xlabel('\mu');
ylabel('SNR (dB)');

figure
%subplot(3,1,2)
hold on, grid on
title('DCT Initialisation')
plot([1 1],[0 25],'k-');
plot(mu(13),sigNoiseRatio(13),'ks');
plot(mu(24),sigNoiseRatio(24),'kd');
plot(mu(14:23),sigNoiseRatio(14:23),'k*-');
plot([minMu minMu],[0 25],'k--')
set(gca,'YLim',[0 25],'XLim',[0 1.4]);
legend({'\mu_{max}','K-SVD','Grassmannian','INK-SVD','\mu_{min}'});
xlabel('\mu');
ylabel('SNR (dB)');

figure
%subplot(3,1,3)
hold on, grid on
title('Gabor Initialisation')
plot([1 1],[0 25],'k-');
plot(mu(25),sigNoiseRatio(25),'ks');
plot(mu(36),sigNoiseRatio(36),'kd');
plot(mu(26:35),sigNoiseRatio(26:35),'k*-');
plot([minMu minMu],[0 25],'k--')
set(gca,'YLim',[0 25],'XLim',[0 1.4]);
legend({'\mu_{max}','K-SVD','Grassmannian','INK-SVD','\mu_{min}'});
xlabel('\mu');
ylabel('SNR (dB)');

% minMu = sqrt((dictSize-blockSize)/(blockSize*(dictSize-1)));
% maxSNR = max(sigNoiseRatio);
%
% figure, subplot(2,2,1)
% snrMat = buffer(sigNoiseRatio(1:9),3);
% bar(snrMat');
% title('Signal to noise ratio')
% xlabel('Initial dictionary')
% ylabel('SNR (dB)')
% set(gca,'XTickLabel',{'data','dct','gabor'});
% legend('none','Mailhe','Tropp')
% grid on
%
% subplot(2,2,2), grid on
% snrMat = buffer(sigNoiseRatio(10:18),3);
% bar(snrMat');
% title('SNR - Mailhe Update')
% xlabel('Initial dictionary')
% ylabel('SNR (dB)')
% set(gca,'XTickLabel',{'data','dct','gabor'},'YLim',[0 maxSNR+1]);
% legend('none','Mailhe','Tropp')
% grid on
%
% subplot(2,2,3), hold on, grid on
% title('Coherence')
% muMat = buffer(mu(1:9),3);
% line([0.5 3.5],[1 1],'Color','r');
% bar(muMat');
% line([0.5 3.5],[minMu minMu],'Color','k');
% set(gca,'XTick',1:3,'XTickLabel',{'data','dct','gabor'},'YLim',[0 1.05])
% legend('\mu_{max}','none','Mailhe','Tropp','\mu_{min}')
% ylabel('\mu')
% xlabel('Initial dictionary')
%
% subplot(2,2,4), hold on, grid on
% title('Coherence - Mailhe Update')
% muMat = buffer(mu(10:18),3);
% line([0.5 3.5],[1 1],'Color','r');
% bar(muMat');
% line([0.5 3.5],[minMu minMu],'Color','k');
% set(gca,'XTick',1:3,'XTickLabel',{'data','dct','gabor'},'YLim',[0 1.05])
% legend('\mu_{max}','none','Mailhe','Tropp','\mu_{min}')
% ylabel('\mu')
% xlabel('Initial dictionary')
