close all;
clear;
clc;

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

%% Constants
noiseSD = 0/255;

lambda = 0.01;
xi = 0.0025;

n = 11;
m = 100;

L = 20;

MAXITER = 100;
rho = 1;
eta = 0.1;

percentOfSamples = 0.3;

%% Loading...
I = {};
load('datasets\separation\cat4.mat');
I{1} = im2double(imgc);

M = {};
load('datasets\separation\fense2.mat');
M{1} = 0.4*im2double((imresize(fense2,[256,256])));

N = length(I);

noisyI = cell(1,N);
sz = cell(1,N);
for i = 1 : N
    sz{i} = size(I{i});
    noisyI{i} = I{i} + M{i} + noiseSD * randn(sz{i});
end

%% Dictionary Learning
cleanIt = cell(1,N);
cleanIc = cell(1,N);

s = cell(1,N);
alpha = cell(1,N);
Z = cell(1,N);

u = cell(1,N);
V = cell(1,N);

for i = 1 : N    
    s{i} = myim2col(noisyI{i}, n) * 0;
    cleanIc{i} = noisyI{i};
    u{i} = 0 * myim2col(noisyI{i}, n);
    V{i} = zeros(sz{i});
end

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

objective   = zeros(1,MAXITER);
feasibility = zeros(1,MAXITER);
sparsity    = zeros(1,MAXITER);
totTime     = zeros(1,MAXITER);
sumsec      = 0;

op  = @(U) ForwardD(U, [1,1,0]);

for outerIter = 1 : MAXITER
    
    sumnnz      = 0;
    sumnumel    = 0;
    feas        = 0;
    obj         = 0;
    
    tStart = tic;
    
    %% Sparse coding
    G = D'*D;
    for i = 1 : N
        % alpha update
        sig = s{i} + u{i};

        param = [];
        param.L = L;
        param.lambda = lambda / rho;
        param.mode = 2;
        alpha{i} = mexLasso(sig, G, D'*sig, param);
        
        Dalpha = D*alpha{i};
        
        % Z update
        opts.rho_r   = 1;
        opts.beta    = [1 1 0];
        opts.print   = true;
        opts.method  = 'l2';
        opts.max_itr = 100;
        opts.tol     = 1e-10;
        mu           = eta / xi;
        tv = deconvtv(cleanIc{i} + V{i}, 1, mu, opts);
        Z{i} = tv.f;

        % s and cleanIc update
        top = myim2col(noisyI{i}, n) + rho*(Dalpha - u{i});
        bottom = noisyI{i} + eta*(Z{i} - V{i});
        summed = mycol2im(1/rho * top, sz{i}, n) + 1/eta * bottom;
        averaged = summed ./ (1 + n^2/rho + 1/eta);
        s{i} = 1/rho * top - 1/rho * myim2col(averaged, n);
        cleanIc{i} = 1/eta * bottom - 1/eta * averaged;
        
        % u update
        u{i} = u{i} + s{i} - Dalpha;
        
        % V update
        V{i} = V{i} + cleanIc{i} - Z{i};
        
        cleanIt{i} = mycol2im(Dalpha, sz{i}, n);
    
        sumnnz = sumnnz + nnz(alpha{i}) + nnz(alpha{i});
        sumnumel = sumnumel + numel(alpha{i}) + numel(alpha{i});
        feas = feas + norm(s{i} - Dalpha,'fro')^2 + norm(Z{i} - cleanIc{i},'fro')^2;
        [Df1, Df2, Df3] = op(cleanIc{i});
        obj = obj + 0.5*norm(noisyI{i} - cleanIt{i} - cleanIc{i},'fro')^2 ...
            + lambda*sum(abs(alpha{i}(:))) + xi*sum(sqrt(Df1(:).^2 + Df2(:).^2));
    end
    
    sparsity(outerIter) = sumnnz/sumnumel;
    feasibility(outerIter) = feas;
    objective(outerIter) = obj;
    
    %% Dictionary Update
    
    x_flat = cell2mat(s) + cell2mat(u);
    alpha_flat = cell2mat(alpha);
    
    samples = randperm(size(x_flat,2),ceil(percentOfSamples*size(x_flat,2)));
    x_flat = x_flat(:,samples);
    alpha_flat = alpha_flat(:,samples);
    
    
    replaced_atoms = zeros(1,m);
    
    unused_sigs = 1 : size(x_flat,2);
    
    p = [];
    for i = 1 : 1
        p = [p randperm(m)];
    end
    for j = 1 : length(p)
        i = p(j);
        
        idxs = find(alpha_flat(i,:));
        
        if length(idxs) < 1
            maxsignals = 5000;
            perm = randperm(length(unused_sigs));
            perm = perm(1:min(maxsignals,length(perm)));
            
            E = x_flat(:,unused_sigs(perm)) - D*alpha_flat(:,unused_sigs(perm));
            E = sum(E.^2);
            [~,maxidx] = max(E);
            
            D(:,i) = x_flat(:,unused_sigs(perm(maxidx)));
            D(:,i) = D(:,i) / norm(D(:,i));
            assert(norm(D(:,i))>1e-10);
            
            alpha_flat(:,i) = zeros(1,m);
            
            unused_sigs = unused_sigs([1:perm(maxidx)-1,perm(maxidx)+1:end]);
            
            replaced_atoms(i) = 1;
            
            continue;
        end
        
        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
    
    %% Plots
    
    tElapsed = toc(tStart);
    
    sumsec = sumsec + tElapsed;
    
    totTime(outerIter) = sumsec;

    fprintf('Iter = %d, Obj = %.3d, Sparsity = %.3f, Feasibility = %.3f, Total-Time (min) = %.3f \n',...
        outerIter, objective(outerIter), sparsity(outerIter), feasibility(outerIter), totTime(outerIter)/60);

    switchFigure(1);

    subplot(1,3,1);
    plot(sparsity(1:outerIter));
    title('Sparsity');

    subplot(1,3,2);
    semilogy(feasibility(1:outerIter));
    title('Feasibility');

    subplot(1,3,3);
    semilogy(objective(1:outerIter));
    title('Objective');
    
    switchFigure(2);
    randidx = randi(N);
    
    subplot(2,2,1);
    imshow(noisyI{randidx},[]);
    title('Original');
    
    subplot(2,2,2);
    totI = cleanIt{randidx} + cleanIc{randidx};
    imshow(totI,[]);
    title('Restored');
    
    subplot(2,2,3);
    imshow(cleanIt{randidx},[]);
    title('Texture');
    
    subplot(2,2,4);
    imshow(cleanIc{randidx},[]);
    title('Cartoon');
    
    switchFigure(3);
    showDictionary(D);
    title('Trained dictionary');
    
    drawnow;
end

