function x = tv(y,lambda)
%
% iteratively reweighted least squares approach for total-variation
% image smoothing
% includes generic subfunction for solving optimizations like
% minimize( ||A*x-b||_1 + ||D*x-c||^2 )
%
% written middle of 2008, posted online jan 9, 2009
% author: justin domke  email: (author's last name)@cs.umd.edu
%
% demo of how to use this (segment a random image)
% y = rand(50);
% x = tv(y,1);


[ly lx] = size(y);

siz = length(y);
N = reshape(1:ly*lx,ly,lx);
nconst = ly*(lx-1) + (ly-1)*lx;

% set up the matrix A

% A = sparse(nconst,siz^2);
% c = 1;
% for i=1:siz-1
%     for j=1:siz
%         n1 = N(i,j);
%         n2 = N(i+1,j);
%         A(c,n1) =  1;
%         A(c,n2) = -1;
%         c = c + 1;
%     end
% end
% for i=1:siz
%     for j=1:siz-1
%         n1 = N(i,j);
%         n2 = N(i,j+1);
%         A(c,n1) =  1;
%         A(c,n2) = -1;
%         c = c + 1;
%     end
% end

% faster, more unclear way
I = zeros(2*nconst,1);
J = zeros(2*nconst,1);
V = zeros(2*nconst,1);
c = 1;
n = 1;
for i=1:ly-1
    for j=1:lx
        n1 = N(i,j);
        n2 = N(i+1,j);
        I(n) = c;
        J(n) = n1;
        V(n) = 1;
        n = n + 1;
        I(n) = c;
        J(n) = n2;
        V(n) = -1;
        n = n + 1;
        c = c + 1;
    end
end
for i=1:ly
    for j=1:lx-1
        n1 = N(i,j);
        n2 = N(i,j+1);
        I(n) = c;
        J(n) = n1;
        V(n) = 1;
        n = n + 1;
        I(n) = c;
        J(n) = n2;
        V(n) = -1;
        n = n + 1;
        c = c + 1;
    end
end
A = sparse(I,J,V);

y = y(:);

myA = lambda*A;
myb = zeros(nconst,1);

myD = speye(ly*lx);
myc = y;

% function to display everything during minimization
function fun(x)
    x = reshape(x,ly,lx);
    l = ceil(x*1000);
    myim = [repmat(reshape(y,ly,lx),[1 1 3]); ...
            repmat(reshape(x,ly,lx),[1 1 3]); ...
            double(label2rgb(max(0,l),'jet','k','shuffle'))/255 ];
    imshow(myim)
    drawnow
end
% function to do nothing
function nofun(x)
end

x = irls_solveL1L2(myA,myb,myD,myc,1e-10,@fun);
% use this if you prefer no output
%x = irls_solveL1L2(myA,myb,myD,myc,1e-10,@nofun);

x = reshape(x,ly,lx);

end

function x = irls_solveL1L2(A,b,D,c,tol,fun)

% x = irls_solveL1L2(A,b)
% minimize( ||A*x-b||_1 + ||D*x-c||^2 )

% based on p 196-197 in Linear Programming: Foundations and Extensions
% by Robert Vanberbei

if nargin < 5
    tol = 1e-5;
end

[m n] = size(A);

x = D\c;

old_norm = inf;

DtD2 = 2*D'*D;
Dtc2 = 2*D'*c;

tic
for iter=1:inf
    new_norm = norm(A*x-b,1) + sum( (D*x-c).^2 );
    time = toc;
    e = abs(A*x-b);
    Ei = sparse(1:m,1:m,1./max(e,1e-8)); % constant of 1e-10 is not well justified
    
    if mod(iter,10)==0
        if nargin >= 6, fun(x); end
        fprintf('iter: %d  norm: %f  time: %f \n', iter, new_norm, time);
    end
    
    %M = (A'*Ei*A + 2*D'*D);
    %n = (A'*Ei*b + 2*D'*c);
    M = (A'*Ei*A + DtD2);
    n = (A'*Ei*b + Dtc2);
    
    x = gauss_seidel(M'*M/1e5, M'*n/1e5, x, 50);
    
    if old_norm - new_norm < tol; break; end
    old_norm = new_norm;
end

end

function x = gauss_seidel(A, b, x, maxiter)

I = size(A,1);

if nargin < 3 || isempty(x), x=zeros(I,1); end
if nargin < 4, maxiter = inf; end
    
D = sparse(1:I,1:I,diag(A));
U = triu(A,1);
L = tril(A,-1);

%P = (D+L);
%Q = U;

% slightly faster for mysterious reasons
%P = D+U;
%Q = L;

% SOR
w = 1.95;
P = D+w*L;
Q = -w*U + (1-w)*D;
myb = w*(P\b);

for i=1:maxiter
    %x = P\(b-Q*x);
    
    % SOR
    x = P\(Q*x) + myb;
    
    if mod(i,1)==0
        mynorm = norm(A*x-b);
        %fprintf('iter: %d  norm: %f \n', i, mynorm);
        if mynorm < 1e-10
            break
        end
    end
end

end