%
%
%

clc, clear, close all

%% Parameteres
nTrials   = 10;										%number of trials of the experiment

% Dictionary learning parameters
toolbox   = 'TwoStepDL';							%dictionary learning toolbox
dicUpdate = 'ksvd';									%dictionary learning updates
dicDecorr = {'iterproj','ink-svd','shrinkgram'};	%dictionary decorrelation methods
minCoherence = linspace(0.1,1,10);					%coherence levels

iterNum   = 20;				%number of iterations
epsilon   = 1e-6;			%tolerance level
dictSize  = 512;			%number of atoms in the dictionary
percActiveAtoms = 5;		%percentage of active atoms

% Test signal parameters
signal    = audio('music03_16kHz.wav'); %audio signal
blockSize = 256;						%size of audio frames
overlap   = 0.5;						%overlap between consecutive frames

% Dependent parameters
nActiveAtoms = fix(blockSize/100*percActiveAtoms); %number of active atoms

% Initial dictionaries
gaborParam = struct('N',blockSize,'redundancyFactor',2,'wd',@rectwin);
gaborDict = Gabor_Dictionary(gaborParam);
initDicts = {[],gaborDict};

%% Generate audio approximation problem
signal			 = buffer(signal,blockSize,blockSize*overlap,@rectwin);	%buffer frames of audio into columns of the matrix S
SMALL.Problem.b  = signal.S;
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
nDecorrAlgs = length(dicDecorr);		%number of decorrelation algorithms
nCorrLevels = length(minCoherence);		%number of coherence levels
nInitDicts  = length(initDicts);		%number of initial dictionaries

SMALL.DL(nTrials,nInitDicts,nCorrLevels,nDecorrAlgs) = SMALL_init_DL(toolbox); %create dictionary learning structures
for iTrial=1:nTrials
	for iInitDicts=1:nInitDicts
		for iCorrLevels=1:nCorrLevels
			for iDecorrAlgs=1:nDecorrAlgs
				SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).toolbox = toolbox;
				SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).name = dicUpdate;
				SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).profile = true;
				SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).param = ...
					struct( 'data',SMALL.Problem.b,...
					'Tdata',nActiveAtoms,...
					'dictsize',dictSize,...
					'iternum',iterNum,...
					'memusage','high',...
					'solver',solver,...
					'decFcn',dicDecorr{iDecorrAlgs},...
					'coherence',minCoherence(iCorrLevels),...
					'initdict',initDicts(iInitDicts));
				
				SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs) = ...
					SMALL_learn(SMALL.Problem,SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs));
				save('SMALL_DL','SMALL');
			end
		end
	end
end

%% Evaluate coherence and snr of representation for the various methods
sr = zeros(size(SMALL.DL));				%signal to noise ratio
mu1 = zeros(nTrials,nInitDicts,nCorrLevels,nDecorrAlgs,blockSize);	%cumulative coherence
mu2 = zeros(nTrials,nInitDicts,nCorrLevels,nDecorrAlgs,blockSize);	%cumulative coherence
dic(size(SMALL.DL)) = dictionary;		%initialise dictionary objects
for iTrial=1:nTrials
	for iInitDicts=1:nInitDicts
		for iCorrLevels=1:nCorrLevels
			for iDecorrAlgs=1:nDecorrAlgs
				%Sparse representation
				SMALL.Problem.A = SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).D;
				tempSolver = SMALL_solve(SMALL.Problem,solver);
				%calculate snr
				sr(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs) = ...
					snr(SMALL.Problem.b,SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).D*tempSolver.solution);
				%calculate mu
				dic(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs) = ...
					dictionary(SMALL.DL(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).D);
				mu1(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs,:) = ...
					dic(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).cumcoherence;
				mu2(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs,:) = ...
					dic(iTrial,iInitDicts,iCorrLevels,iDecorrAlgs).cumcoherence(2);
			end
		end
	end
end

%% Plot results
minMu = sqrt((dictSize-blockSize)/(blockSize*(dictSize-1)));	%lowe bound on coherence
initDictsNames = {'Data','Gabor'};
dicDecorrNames = {'IPR','INK-SVD'};
lineStyles     = {'k.-','r*-','b+-'};
for iInitDict=1:nInitDicts
	figure, hold on, grid on
	title([initDictsNames{iInitDict} ' Initialisation']);
	plot([1 1],[0 25],'k-');
	for iDecorrAlgs=1:nDecorrAlgs-1
		coherenceLevels = squeeze(mean(mu(:,iInitDict,:,iDecorrAlgs,1),1));
		meanSNRs		= squeeze(mean(sr(:,iInitDict,:,iDecorrAlgs),1));
		stdSNRs			= squeeze(std(sr(:,iInitDict,:,iDecorrAlgs),0,1));
		errorbar(coherenceLevels,meanSNRs,stdSNRs,lineStyles{iDecorrAlgs});
	end
	plot([minMu minMu],[0 25],'k--')
	
	set(gca,'YLim',[0 25],'XLim',[0 1.4]);
	legend([{'\mu_{max}'},dicDecorrNames,{'\mu_{min}'}]);
	xlabel('\mu');
	ylabel('SNR (dB)');
end

%% 
mu2 = squeeze(mean(mu2,1));
mu = squeeze(mean(mu,1));
