%% This script does gradient denoising using some convex regularisation priors.

addpath('../solutions/','../Lin1D/');


%% set up some prior image regularisation functionals.
%
TK1 = @(s,T) (s.^2)/2;
TV =  @(s,T) (abs(s));
sTV = @(s,T) (T*sqrt(s.^2 + T^2) - T^2);
PM1 = @(s,T) (T^2/2*(log(1 + (s/T).^2) ));
%Huber = @(s,T) (if (s < T) s/^2/2; else T*s -T^2/2; end  );

dTK1 = @(s,T) (s);
dTV =  @(s,T) (ones(size(s)));
dsTV = @(s,T) ((s*T)./sqrt(s.^2 + T^2));
dPM1 = @(s,T) ((s*T^2)./(s.^2 + T^2));

kapTK1 = @(s,T) (ones(size(s)));
kapTV =  @(s,T) (1./s);
kapsTV = @(s,T) ((T)./sqrt(s.^2 + T^2));
kapPM1 = @(s,T) ((T^2)./(s.^2 + T^2));
%%
N = 256;
impath = '../../StandardTestImages/';
ln = imread([impath 'Lena512.png'],'png'); ln = double(imresize(ln,[N N]));
hs = imread([impath 'house.png'],'png'); hs = double(imresize(hs,[N N]));
pt = imread([impath 'Cameraman256.png'],'png'); pt = double(imresize(pt,[N N]));
ot = imread([impath 'ot.pgm']); ot = double(imresize(ot,[N,N]));
sl = phantom(N);
xt = pt; %pt;%hs;%sl
xt = xt./max(xt(:)); % normalise to max of 1;
Tf = 0.2; % fraction to apply to max to get threshold
figure(1); clf;
subplot(2,3,1); colormap(gray); imagesc(xt); title('original f_{true}');colorbar;
immax = max(max(xt));
%% ------------------------- noise image ----------------------------------
sig = 0.05;
noiselevel = 0.05;

sig = 0;
xb = xt;

xb = xb + noiselevel*immax.*randn(size(xt));
figure(1);
subplot(2,3,2); colormap(gray); imagesc(xb); title('noisy g');colorbar;
%% seperate figure for slides...
figure(100);clf;
subplot(1,2,1);colormap(gray); imagesc(xt); title('original f_{true}');colorbar;
subplot(1,2,2);colormap(gray); imagesc(xb); title('noisy g');colorbar;
%% -------------------- estimate edge threshold ----------------------
[gx,gy] = gradient(xb);
gxb = sqrt(gx.^2 + gy.^2);
eps = Tf*max(max(gxb));
disp(['setting edge threshold to ',num2str(eps)]);
Rfunc = 'sTV';
switch Rfunc
    case 'PM1'
        kapfunc = @(s,T) kapPM1(s,T);
        Psifunc = @(s,T) PM1(s,T);
        disp(['setting regularisation to Perona Malik'])
        
    case 'sTV'
        kapfunc = @(s,T) kapsTV(s,T);
        Psifunc = @(s,T) sTV(s,T);
        disp(['setting regularisation to Smoothed TV'])
    otherwise
        disp(['Unrecognised Regularistion functional',Rfunc]);
end
%% ---------------- scale space ------------------------------------------
f1d = reshape(xb,[],1);
SSeps = eps/5; % scale space uses smaller value of threshold...
dt = 0.25;
Verbose = false; % suppress 2+too much reporting
for k = 1:100
    Alap = AnisoLaplacian(reshape(f1d,N,N),kapfunc,SSeps,Verbose);
    f1d = f1d + dt*Alap*f1d;
    figure(3); clf; imagesc(reshape(f1d,N,N));colormap(gray); title(['Scale Space : Iteration ',num2str(k)]);colorbar;
    pause(0.1);
end


%% ------------------------- denoise --------------------------------------
ExplicitScheme = false; % true;
alpha = 2e-1; % 2e-2; % 1e-4;
%tau=min([0.25,1/alpha]); % choose step length. Tricky!
tau = 2;
tol=1e-6;

err0=1e15;
tic;
niter = 150;
err = zeros(niter,1);
Psival= err;
toterr = err;
x1 = zeros(size(xb));  % initialise zero
%x1 = randn(size(xb));  % initialise random
%x1 = ln                % initialise with another image!


xr = x1;
xr1d = reshape(xr,[],1);
g1d  = reshape(xb,[],1);
Solver = 'GradientImplicit';%'LaggedFixedPoint',%'LaggedGaussNewton'; %'GradientExplicit';
for n = 1:niter
    %        [flapl,fkap] = SparseLapl(xr,@(f)PeronaMalik2(f,eps));
    [Alapl,gim,fkap] = AnisoLaplacian(reshape(xr1d,N,N),kapfunc,eps,Verbose);
    switch Solver
        case 'GradientExplicit'
            xr1d = xr1d + tau*(g1d - xr1d + alpha*Alapl*xr1d);            
        case 'GradientImplicit';            
            B = ((1+tau)*speye(length(g1d)) - tau*alpha*Alapl);
            xr1d = B\(xr1d + tau*g1d);
        case 'LaggedGaussNewton'
            BGN = (speye(length(g1d)) - alpha*Alapl);
            r = (g1d - xr1d + alpha*Alapl*xr1d);
            xr1d = xr1d + BGN\r;
        case 'LaggedFixedPoint'
            BGN = (speye(length(g1d)) - alpha*Alapl);
            xr1d = BGN\g1d;
            %case 'come back to this!'
            % solve
            %        xr1 = xr1+tau*pcg( @(x) ATACim(x,alpha,KTKD,sig,n1,n2),gr,tol,kryit);
            %        xr = reshape(xr1,n1,n2);
            %       yr = xr;
            %       gr = g - reshape(yr,[],1);
        otherwise
            disp(['Unknown Solver '],Solver);
    end
    err(n) = 0.5*sum((g1d-xr1d).^2);
    figure(1);
    [gx,gy] = gradient(reshape(xr1d,N,N));
    gg = sqrt(gx.^2 + gy.^2);
    
    
    %   TVerr(n) =  sum(sum(gg));%alpha*norm(gxr,1);
    %   TVerr(n) = eps^2/2 *sum(sum( log( 1 + (gg./eps)^2)));
    Psival(n) = sum(Psifunc(xr1d,eps));
    
    eps = Tf*max(max(gg)); disp(['setting edge threshold to ',num2str(eps)]);
    
    toterr(n) = err(n)+alpha*Psival(n);
    disp(['Max kappa ',num2str(max(fkap(:))), ' datafit error ',num2str(err(n)),' Psi ',num2str(Psival(n)), ' total ',num2str(toterr(n))]);

    figure(1);
    xrim = reshape(xr1d,N,N);
    subplot(2,3,4); colormap(gray); imagesc(xrim);title(['reconstructed f^{(n)} it',num2str(n)]);colorbar;
    subplot(2,3,5); colormap(gray); imagesc(xb-xrim);title(['residual (g-f^{(n)}) it',num2str(n)]);colorbar;
    subplot(2,3,3); colormap(gray); imagesc(gg);title(['gradient \nabla f^{(n)} it',num2str(n)]);colorbar;
    subplot(2,3,6); colormap(gray); imagesc(fkap);title(['\kappa it',num2str(n)]);colorbar;
    drawnow;
    if(toterr(n) >= err0)
        disp('convergence reached');
        break;
    else
        err0 = toterr(n);
        tol=tol/2;
 %       kryit = kryit*2;
 %       tau=sqrt(tau);
    end
    pause(0.5);
end
toc;
subplot(2,3,4); colormap(gray); imagesc(xrim);title('final reconstructed f_{recon}');colorbar;
subplot(2,3,3); colormap(gray); imagesc(xrim-xt);title('final estimate error (f_{true}-f_{recon})'); colorbar;
subplot(2,3,5); colormap(gray); imagesc(xb-xrim);title(['final residual (g-f_{recon}) it',num2str(n)]);colorbar;
subplot(2,3,6); histogram(xb(:)-xrim(:));

figure(2);clf; hold on;
semilogy([1:n],err(1:n));
semilogy([1:n],Psival(1:n),'r');
semilogy([1:n],toterr(1:n),'k');
legend('data error','Prior','Posterior');


%% as a function call...
maxit = 100;
%ExplicitScheme = true;
%alpha = 1;
%Solver = 'LaggedGaussNewton'; %'GradientImplicit';'GradientExplicit';

tol=1e-6;
figure(11); clf;
subplot(2,3,1); colormap(gray); imagesc(xt); title('original f_{true}');colorbar;
subplot(2,3,2); colormap(gray); imagesc(xb); title('noisy g');colorbar;

tic;
[xr2,nit2,err2,Psival2,toterr2]= GradientDenoise(xb,x1,alpha,tol,maxit,kapfunc,Psifunc,eps,Tf,Solver,11);
toc;
xr2im = reshape(xr2,N,N);
subplot(2,3,4); colormap(gray); imagesc(xr2im);title('final reconstructed f_{recon}');colorbar;
subplot(2,3,3); colormap(gray); imagesc(xr2im-xt);title('final estimate error (f_{true}-f_{recon})'); colorbar;
subplot(2,3,5); colormap(gray); imagesc(xb-xr2im);title(['final residual (g-f_{recon}) it',num2str(n)]);colorbar;
subplot(2,3,6); histogram(xb(:)-xr2im(:));


figure(12);clf; hold on;
semilogy([1:nit2],err2(1:nit2));
semilogy([1:nit2],Psival2(1:nit2),'r');
semilogy([1:nit2],toterr2(1:nit2),'k');
legend('data error','Prior','Posterior');


%% run for a range of paramters
npix = length(xb(:));
maxit = 100;
%ExplicitScheme = true;
%Solver = 'LaggedGaussNewton'; %'GradientImplicit';'GradientExplicit';

la = [-3:0.5:3];
na = length(la);
alphaDP = exp(la);
tol=1e-6;

% figure(21); clf;
% subplot(2,3,1); colormap(gray); imagesc(xt); title('original f_{true}');colorbar;
% subplot(2,3,2); colormap(gray); imagesc(xb); title('noisy g');colorbar;

xDP = zeros(npix,na);
nitDP = zeros(na,1);
errDP = zeros(maxit,na);
PsivalDP = errDP; toterrDP = errDP;
DP = zeros(na,1);
rerrDP = DP;
PsiDP = DP;
nsigma = noiselevel*immax;
for k = 1:na
    disp(['Calling GradientDenoise with alpha=',num2str(alphaDP(k))]);
    tic;
    [xDP(:,k),nitDP(k),errDP(:,k),PsivalDP(:,k),toterrDP(:,k)]= GradientDenoise(xb,x1,alphaDP(k),tol,maxit,kapfunc,Psifunc,eps,Tf,Solver,0);
    toc;
    DP(k) = errDP(nitDP(k),k)/npix - nsigma^2;
    rerrDP(k) = errDP(nitDP(k),k);
    PsiDP(k) = PsivalDP(nitDP(k),k);
end


%% look at results;
figure(20);clf;
a0 = binsearch(alphaDP,DP);
subplot(1,2,1);
semilogx(alphaDP,DP,'r');title('DP values');
hold on;
semilogx([alphaDP(1),alphaDP(end)],[0,0],'k');
semilogx([a0,a0],[DP(1),DP(end)],'k');


subplot(1,2,2);
plot(rerrDP,PsiDP,'r');title('L-Curve');

%alphaind = 8;
%nitval = nitDP(alphaind);


figure(21); clf;
disp(['Running denoise with optimal DP value alpha=',num2str(a0)]);
subplot(2,3,1); colormap(gray); imagesc(xt); title('original f_{true}');colorbar;
subplot(2,3,2); colormap(gray); imagesc(xb); title('noisy g');colorbar;
tic;
[xDP0,nitDP0,errDP0,PsivalDP0,toterrDP0]= GradientDenoise(xb,x1,a0,tol,maxit,kapfunc,Psifunc,eps,Tf,Solver,0);
toc;
xrDPim = reshape(xDP0,N,N);
subplot(2,3,4); colormap(gray); imagesc(xrDPim);title('final reconstructed f_{recon}');colorbar;
subplot(2,3,3); colormap(gray); imagesc(xrDPim-xt);title('final estimate error (f_{true}-f_{recon})'); colorbar;
subplot(2,3,5); colormap(gray); imagesc(xb-xrDPim);title(['final residual (g-f_{recon}) it',num2str(n)]);colorbar;
subplot(2,3,6); histogram(xb(:)-xrDPim(:));


figure(22);clf; hold on;
semilogy([1:nitDP0],errDP0(1:nitDP0));
semilogy([1:nitDP0],PsivalDP0(1:nitDP0),'r');
semilogy([1:nitDP0],toterrDP0(1:nitDP0),'k');
legend('data error','Prior','Posterior');

rerrDP0 = errDP0(nitDP0);
PsiDP0 = PsivalDP0(nitDP0);

figure(20);subplot(1,2,2);hold on;
plot(rerrDP0,PsiDP0,'b+');


%%
function [aLap,gim,kap2d,kap1x,kap1y] = AnisoLaplacian(im,kapfunc,T,Verbose)

[n1,n2] = size(im);
N = n1; % assume square.
oneN = ones(N+1,1);
D1x = spdiags([oneN -oneN],-1:0,N+1,N);
D2x = -D1x'*D1x; 
% The following allows interpolation to pixel mid-points;
intx1d = spdiags([0.5*oneN 0.5*oneN],-1:0,N+1,N);

D1x2d = kron(speye(N),D1x);
D1y2d = kron(D1x,speye(N));
intx2d = kron(speye(N),intx1d);
inty2d = kron(intx1d,speye(N));


h = reshape(im,[],1);
gx1d = D1x2d*h; gx2d = reshape(gx1d,n1+1,n2);
gy1d = D1y2d*h; gy2d = reshape(gy1d,n1,n2+1);
gim = sqrt(gx2d(1:n1,:).^2 + gy2d(:,1:n2).^2);

% we can get an alternative gradient image using interpolation operators.
gg1d = sqrt((intx2d'*gx1d).^2 + (inty2d'*gy1d).^2);
gg2d = reshape(gg1d,n1,n2);

% define kappa
kap2d = kapfunc(gim,T);
if(Verbose)
    disp(['N = ',num2str(N)]);
    [nk1,nk2] = size(kap2d);
    disp(['Nk1=',num2str(nk1),' Nk2=',num2str(nk2)]);
    
end
kap1d = reshape(kap2d,[],1);
kap1x = spdiags(intx2d*kap1d,0:0,(n1+1)*n2,(n1+1)*n2);
kap1y = spdiags(inty2d*kap1d,0:0,n1*(n2+1),n1*(n2+1));
aLap = -(D1x2d'*kap1x*D1x2d + D1y2d'*kap1y*D1y2d);
end % end of function

%%
function [xr1d,n,err,Psival,toterr]= GradientDenoise(xb,x1,alpha,tol0,maxit,kapfunc,Psifunc,eps0,Tf,Solver,figno)

[n1,n2] = size(xb);
N = n1; % assume square;
tau=min([0.25,1/alpha]); % choose step length. Tricky!
err0=1e15;
tol = tol0;
niter = maxit;
err = zeros(niter,1);
Psival = err;
toterr = err;
xr = x1;  % initialise 
eps = eps0;
xr1d = reshape(xr,[],1);
g1d  = reshape(xb,[],1);
n = 1;
for n = 1:niter
    [Alapl,gim,fkap] = AnisoLaplacian(reshape(xr1d,N,N),kapfunc,eps,false);
%    if ExplicitScheme        
%        xr1d = xr1d + tau*(g1d - xr1d + alpha*Alapl*xr1d);       
%    else
%        B = ((1+tau)*speye(length(g1d)) - tau*alpha*Alapl);
%       xr1d = B\(xr1d + tau*g1d); % can this be right!?
%    end  
switch Solver
    case 'GradientExplicit'
        xr1d = xr1d + tau*(g1d - xr1d + alpha*Alapl*xr1d);
    case 'GradientImplicit';
        B = ((1+tau)*speye(length(g1d)) - tau*alpha*Alapl);
        xr1d = B\(xr1d + tau*g1d);
    case 'LaggedGaussNewton'
        BGN = (speye(length(g1d)) - alpha*Alapl);
        r = (g1d - xr1d + alpha*Alapl*xr1d);
        xr1d = xr1d + BGN\r;
    case 'LaggedFixedPoint'
        BGN = (speye(length(g1d)) - alpha*Alapl);
        xr1d = BGN\g1d;
        %case 'come back to this!'
        % solve
        %        xr1 = xr1+tau*pcg( @(x) ATACim(x,alpha,KTKD,sig,n1,n2),gr,tol,kryit);
        %        xr = reshape(xr1,n1,n2);
        %       yr = xr;
        %       gr = g - reshape(yr,[],1);
    otherwise
        disp(['Unknown Solver '],Solver);
end
    
    err(n) = 0.5*sum((g1d-xr1d).^2);
    xrim = reshape(xr1d,N,N);
    [gx,gy] = gradient(xrim);
    gg = sqrt(gx.^2 + gy.^2);
    
    Psival(n) = sum(Psifunc(xr1d,eps));
    eps = Tf*max(max(gg)); 
    %disp(['setting edge threshold to ',num2str(eps)]);
    
    toterr(n) = err(n)+alpha*Psival(n);
%    disp(['Max kappa ',num2str(max(fkap(:))), ' datafit error ',num2str(err(n)),' Psi ',num2str(Psival(n)), ' total ',num2str(toterr(n))]);
    if(figno)
        figure(figno);
        subplot(2,3,4); colormap(gray); imagesc(xrim);title(['reconstructed f^{(n)} it',num2str(n)]);colorbar;
        subplot(2,3,5); colormap(gray); imagesc(xb-xrim);title(['residual (g-f^{(n)}) it',num2str(n)]);colorbar;
        subplot(2,3,3); colormap(gray); imagesc(gg);title(['gradient \nabla f^{(n)} it',num2str(n)]);colorbar;
        subplot(2,3,6); colormap(gray); imagesc(fkap);title(['\kappa it',num2str(n)]);colorbar;
        
        drawnow;
        pause(0.5);
    end
    if(toterr(n) >= err0)
        disp('convergence reached');
        break;
    else
        err0 = toterr(n);
        tol=tol/2;
    end
end % end of iterations
end % end of function

%% 
function x0 = binsearch(x,y)
% search for a zero in y sampled at x and return interpolated value
nx = length(x);
lo = 1; mid = floor(nx/2); hi = nx;

while( hi > lo+1)
    if y(mid) < 0
        lo = mid;
    else
        hi = mid;
    end
    mid = floor((hi+lo)/2);
end
x0 = x(lo) - y(lo)*(x(hi)-x(lo))/(y(hi) - y(lo));
end