close all;
clear;
clc;

addpath mexfiles;
addpath ompbox10;
addpath image_helpers;
addpath vlfeat\toolbox;
addpath(genpath('ksvd'));
addpath(genpath('spams-matlab'));
vl_setup();

%% Constants
% noise added to each image
noiseSD = 0;

% Lagrangian parameter
lambda = 1;

% size of filter
n = 11;

% number of filters
m = 100;

% maximal number of non-zeros per needle (alpha)
L = 5;

% iterations of the ADMM
MAXITER = 500;

% parameter of ADMM
rho = 0.1*n^2;

% percentage of samples used for the dictionary update
percentOfSamples = 0.3;

%% Loading...
% contrast normalization
CONTRAST_NORMALIZE = 'local_cn';
ZERO_MEAN = 1;
COLOR_IMAGES = 'gray';
I = CreateImages('datasets/fruit_100_100', ...
    CONTRAST_NORMALIZE, ZERO_MEAN, COLOR_IMAGES);
I = num2cell(I, [1,2,3]);
I = permute(I,[4,1,2,3]);

% number of images
N = length(I);

noisyI = cell(1,N);
sz = cell(1,N);
sumpsnr = 0;
for i = 1 : N
    sz{i} = size(I{i});
	
	% noisy image
    noisyI{i} = I{i} + noiseSD * randn(sz{i});
    
	% psnr
    p = 20*log10(sqrt(numel(I{i})) / norm(reshape(noisyI{i} - I{i},1,[])));

    sumpsnr = sumpsnr + p;
end

% initial average psnr
initpsnr = sumpsnr / N;

%% Dictionary Learning
y = cell(1,N);
cleanI = cell(1,N);

% slices
s = cell(1,N);

% needles
alpha = cell(1,N);

% dual variables
u = cell(1,N);

% initialize parameters
for i = 1 : N
    y{i} = myim2col(noisyI{i}, n);
    s{i} = y{i} / n^2;
    u{i} = 0 * y{i};
end

D = randn(n^2,m);
for i = 1 : m
    D(:,i) = D(:,i) / norm(D(:,i));
end

% graphs to plot
avgpsnr     = zeros(1,MAXITER);
objective   = zeros(1,MAXITER);
feasibility = zeros(1,MAXITER);
sparsity    = zeros(1,MAXITER);
totTime     = zeros(1,MAXITER);
sumsec      = 0;

for outerIter = 1 : MAXITER
    %% Sparse Coding
    sumpsnr     = 0;
    sumnnz      = 0;
    sumnumel    = 0;
    feas        = 0;
    obj         = 0;
    
    tStart = tic;
    
	% gram of the local dictionary
    G = D'*D;
    for i = 1 : N
		% alpha update via local Lasso
		sig = s{i} + u{i};
		
		param = [];
		param.L = L;
		param.lambda = lambda / rho;
		param.mode = 2;
		
		DtX = D'*sig;
		alpha{i} = mexLasso(sig, G, DtX, param);
		
		% s update via local Laplacian
		Dalpha = D*alpha{i};
		V1 = 1/rho * y{i} + Dalpha - u{i};
		summed = mycol2im(V1, sz{i}, n);
		averaged = summed ./ (n^2 + rho);
		V2 = myim2col(averaged, n);
		s{i} = V1 - V2;
		
		u{i} = u{i} + s{i} - Dalpha;
		
		cleanI{i} = mycol2im(Dalpha, sz{i}, n);
		
		
		sumpsnr = sumpsnr + 20*log10(sqrt(numel(I{i})) / norm(reshape(cleanI{i} - I{i},1,[])));
		sumnnz = sumnnz + nnz(alpha{i});
		sumnumel = sumnumel + numel(alpha{i});
		feas = feas + norm(s{i} - Dalpha,'fro')^2;
		obj = obj + 0.5*norm(noisyI{i} - cleanI{i},'fro')^2 + lambda*sum(abs(alpha{i}(:)));
    end
    
	% update graphs
    avgpsnr(outerIter) = sumpsnr/N;
    sparsity(outerIter) = sumnnz/sumnumel;
    feasibility(outerIter) = feas;
    objective(outerIter) = obj;
    
    
    %% Dictionary Update
	% all slices concatenated
	s_flat = cell2mat(s) + cell2mat(u);
	
	% all needles concatenated
    alpha_flat = cell2mat(alpha);
    
	% random subset of samples
    samples = randperm(size(s_flat,2),ceil(percentOfSamples*size(s_flat,2)));
    s_flat = s_flat(:,samples);
    alpha_flat = alpha_flat(:,samples);
    
	% keeps track of which slices were used to initialize an atom
    unused_sigs = 1 : size(s_flat,2);
    
	% if we want to sweep more than once over all the atoms
    p = [];
    for i = 1 : 1
        p = [p randperm(m)];
    end
	
	% update all the atoms sequentially
    for j = 1 : length(p)
	    % the index of the atom being updated
        i = p(j);
        
		% find slices that use the i-th atom
        idxs = find(alpha_flat(i,:));
        
		% if no slice use the atom, replace it
        if length(idxs) < 1
            maxsignals = 5000;
            perm = randperm(length(unused_sigs));
            perm = perm(1:min(maxsignals,length(perm)));
            
			% error of each slice
            E = s_flat(:,unused_sigs(perm)) - D*alpha_flat(:,unused_sigs(perm));
            E = sum(E.^2);
			
			% slice with highest error
            [~,maxidx] = max(E);
            
			% set the atom to be this slice
            D(:,i) = s_flat(:,unused_sigs(perm(maxidx)));
            D(:,i) = D(:,i) / norm(D(:,i));
            
			% set the coefficients of the replaced atom to zero
            alpha_flat(i,:) = zeros(size(alpha_flat(i,:)));
            
			% remove the chosen slice from the unused signals
            unused_sigs = unused_sigs([1:perm(maxidx)-1,perm(maxidx)+1:end]);
            
            continue;
        end
        
		% k-svd update stage (page 233 in the book "Sparse and Redundant Representations")
        E = s_flat(:,idxs) - D*alpha_flat(:,idxs) + D(:,i)*alpha_flat(i,idxs);
        
        D(:,i) = E*alpha_flat(i,idxs)';
        D(:,i) = D(:,i) / norm(D(:,i));
        
        E = s_flat(:,idxs) - D*alpha_flat(:,idxs) + D(:,i)*alpha_flat(i,idxs);
        
        alpha_flat(i,idxs) = E'*D(:,i);
    end

	% time for iteration
    tElapsed = toc(tStart);
    
	% total time
    sumsec = sumsec + tElapsed;
    totTime(outerIter) = sumsec;
    
	% print statistics
    fprintf('Iter = %d, Obj = %.3d, Sparsity = %.3f, Feasibility = %.3f, Avg-PSNR = %.3f, Total-Time (min) = %.3f \n',...
        outerIter, objective(outerIter), sparsity(outerIter), feasibility(outerIter),...
        avgpsnr(outerIter), totTime(outerIter)/60);
end


%% Plots
switchFigure(1);
showDictionary(D);
title('Trained dictionary'); drawnow()

switchFigure(2);
subplot(1,4,1);
plot(avgpsnr);
hold on;
plot(1:length(avgpsnr), initpsnr * ones(size(avgpsnr)),'r');
hold off;
legend('Estimated image','Initial image', 'location', 'southeast');
title('Average PSNR');

subplot(1,4,2);
plot(sparsity);
title('Sparsity');

subplot(1,4,3);
semilogy(feasibility);
title('Feasibility');

subplot(1,4,4);
semilogy(objective);
title('Objective');
