clc, clear, close all

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

% Dictionary learning parameters
toolbox   = 'TwoStepDL';							%dictionary learning toolbox
dicUpdate = 'mocod';								%dictionary learning updates
zeta	  = logspace(-2,2,10); 
eta  	  = logspace(-2,2,10); 

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
nInitDicts  = length(initDicts);		%number of initial dictionaries
nZetas = length(zeta);
nEtas  = length(eta);

SMALL.DL(nTrials,nInitDicts,nZetas,nEtas) = SMALL_init_DL(toolbox); %create dictionary learning structures
for iTrial=1:nTrials
	for iInitDicts=1:nInitDicts
		for iZetas=1:nZetas
			for iEtas=1:nEtas
				SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).toolbox = toolbox;
				SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).name = dicUpdate;
				SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).profile = true;
				SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).param = ...
					struct('data',SMALL.Problem.b,...
					'Tdata',nActiveAtoms,...
					'dictsize',dictSize,...
					'iternum',iterNum,...
					'memusage','high',...
					'solver',solver,...
					'initdict',initDicts(iInitDicts),...
					'zeta',zeta(iZetas),...
					'eta',eta(iEtas));
				SMALL.DL(iTrial,iInitDicts,iZetas,iEtas) = ...
					SMALL_learn(SMALL.Problem,SMALL.DL(iTrial,iInitDicts,iZetas,iEtas));
			end
		end
	end
end

%% Evaluate coherence and snr of representation for the various methods
sr = zeros(size(SMALL.DL));				%signal to noise ratio
mu = zeros(iTrial,iInitDicts,iZetas,iEtas);	%coherence
dic(size(SMALL.DL)) = dictionary;		%initialise dictionary objects
for iTrial=1:nTrials
	for iInitDicts=1:nInitDicts
		for iZetas=1:nZetas
			for iEtas=1:nEtas
				%Sparse representation
				SMALL.Problem.A = SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).D;
				tempSolver = SMALL_solve(SMALL.Problem,solver);
				%calculate snr
				sr(iTrial,iInitDicts,iZetas,iEtas) = ...
					snr(SMALL.Problem.b,SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).D*tempSolver.solution);
				%calculate mu
				dic(iTrial,iInitDicts,iZetas,iEtas) = ...
					dictionary(SMALL.DL(iTrial,iInitDicts,iZetas,iEtas).D);
				mu(iTrial,iInitDicts,iZetas,iEtas) = ...
					dic(iTrial,iInitDicts,iZetas,iEtas).coherence;
			end
		end
	end
end

save('MOCOD.mat')

%% Plot results
minMu = sqrt((dictSize-blockSize)/(blockSize*(dictSize-1)));	%lowe bound on coherence
initDictsNames = {'Data','Gabor'};
lineStyles     = {'k.-','r*-','b+-'};
for iInitDict=1:nInitDicts
	figure, hold on, grid on
	title([initDictsNames{iInitDict} ' Initialisation']);
	coherenceLevels = squeeze(mean(mu(:,iInitDict,:,:),1));
	meanSNRs		= squeeze(mean(sr(:,iInitDict,:,:),1));
	%stdSNRs				= squeeze(std(sr(:,iInitDict,iZetas,iEtas),0,1));
	subplot(2,2,1)
	surf(eta,zeta,coherenceLevels);
	set(gca,'Xscale','log','Yscale','log','ZLim',[0 1.4]);
	view(gca,130,20)
	xlabel('\eta');
	ylabel('\zeta');
	zlabel('\mu');
	title('Coherence')
	
	subplot(2,2,2)
	surf(eta,zeta,meanSNRs);
	set(gca,'Xscale','log','Yscale','log','ZLim',[0 25]);
	view(gca,130,20)
	xlabel('\eta');
	ylabel('\zeta');
	zlabel('SNR (dB)');
	title('Reconstruction Error')
	
	subplot(2,2,[3 4])
	mus = mu(:,iInitDict,:,:);
	mus = mus(:);
	SNRs = sr(:,iInitDict,:,:);
	SNRs = SNRs(:);
	[un idx] = sort(mus);
	plot([1 1],[0 25],'k')
	hold on, grid on
	scatter(mus(idx),SNRs(idx),'k+');
	plot([minMu minMu],[0 25],'k--')
	set(gca,'YLim',[0 25],'XLim',[0 1.4]);
	xlabel('\mu');
	ylabel('SNR (dB)');
	legend([{'\mu_{max}'},'MOCOD',{'\mu_{min}'}]);
	title('Coherence-Reconstruction Error Tradeoff')
	
% 	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
