close all;
clear;
clc;

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

%% Loading...
CONTRAST_NORMALIZE = 'local_cn';
ZERO_MEAN = 1;
COLOR_IMAGES = 'gray';
I = CreateImages2('datasets/standard', ...
    CONTRAST_NORMALIZE, ZERO_MEAN, COLOR_IMAGES);

% number of images
N = length(I);

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

% iterations of the ADMM
MAXITER = 100;

% load dictionary trained using main.m
load fruits_dictionary.mat;

% specifiy whether to train the dictionary on the corrupted image
dict_update = 0;

% best Lagrangian parameter for inpainting
if dict_update
	lambda = 0.04;
else
	lambda = 0.06;
end

% size of filter
n = sqrt(size(D,1));

% number of filters
m = size(D,2);

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

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

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

noisyI = cell(1,N);
sz = cell(1,N);
M = cell(1,N);
sumpsnr = 0;
for i = 1 : N
	sz{i} = size(I{i});

	% noisy image
	noisyI{i} = I{i} + noiseSD * randn(sz{i});

	% random mask
	M{i} = zeros(sz{i});
	M{i}(rand(sz{i}) < 0.5) = 1;
	
	% degraded image
	noisyI{i} = M{i}.*noisyI{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
x = 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);
	x{i} = y{i} / n^2;
	u{i} = 0 * y{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 = x{i} + u{i};
		
		param = [];
		param.L = L;
		param.lambda = lambda / rho;
		param.mode = 2;
		
		DtX = D'*sig;
		alpha{i} = mexLasso(sig, G, DtX, param);
		
		% x update via local Laplacian (taking into account the mask now)
		Dalpha = D*alpha{i};
		V1 = 1/rho * y{i} + Dalpha - u{i};
		summed = M{i}.*mycol2im(V1, sz{i}, n);
		averaged = summed ./ (n^2 + rho);
		V2 = myim2col(averaged, n);
		x{i} = V1 - V2;
		
		u{i} = u{i} + x{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(x{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
	
	if dict_update
		% all slices concatenated
		x_flat = cell2mat(x) + cell2mat(u);
		
		% all needles concatenated
		alpha_flat = cell2mat(alpha);
		
		% random subset of samples
		samples = randperm(size(x_flat,2),ceil(percentOfSamples*size(x_flat,2)));
		x_flat = x_flat(:,samples);
		alpha_flat = alpha_flat(:,samples);
		
		% keeps track of which slices were used to initialize an atom
		unused_sigs = 1 : size(x_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 = x_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) = x_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 = x_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 = x_flat(:,idxs) - D*alpha_flat(:,idxs) + D(:,i)*alpha_flat(i,idxs);
			
			alpha_flat(i,idxs) = E'*D(:,i);
			
		end
	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

