im = phantom(256);
%theta3 = [0:1:60];
theta3 = [0:8:180];
MultiplicativeNoise =  false;
Positivity =  true;
proxBox = @(x) (min(1,max(0,x)));

noiselevel = 0.01;
g3 = radon(im,theta3);
ndat = length(g3(:));
if MultiplicativeNoise
    g3 = g3.*(1 + noiselevel*randn(size(g3)));
else
    g3 = g3 + noiselevel*max(g3(:))*randn(size(g3));
end
reconUBP3 = iradon(g3,theta3,'linear','none',1,256);
reconFBP3 = iradon(g3,theta3,1,256);
lambda = 1e2*noiselevel;
niter = 20;
tol = 1e-6;
ATg = reshape(reconUBP3,[],1); % back projected data.
x0 = pcg(@(x) ATARadonTK0(x,theta3,256,lambda),ATg,tol,niter );
x1 = pcg(@(x) ATARadonTK1(x,theta3,256,lambda),ATg,tol,niter );
if(Positivity)
    x0 = proxBox(x0);
    x1 = proxBox(x1);
end

%% display
figure(2); clf; hold on;
subplot(2,3,1); imagesc(im);title('Original Image');colorbar;
subplot(2,3,2); imagesc(g3);title('Sinogram space');colorbar;
subplot(2,3,3); imagesc(reconFBP3);title('Filtered Back Projection');colorbar;
subplot(2,3,4); imagesc(reshape(x0,256,256)); title(['Regularised TK0, \alpha=',num2str(lambda)]);colorbar;
subplot(2,3,5); imagesc(reshape(x1,256,256)); title(['Regularised TK1, \alpha=',num2str(lambda)]);colorbar;


%% try anisotropic diffusion
kapPM1 = @(s,T) ((T^2)./(s.^2 + T^2));
eps = 0.5; % ad hoc choice of edge strength 
[aLap,gim,kap2d,kap1x,kap1y] = AnisoLaplacian(reconFBP3,kapPM1,eps,false);
figure(1);clf;
subplot(2,2,1); imagesc(im);
subplot(2,2,2); imagesc(reconFBP3);
subplot(2,2,3); imagesc(gim);
subplot(2,2,4); imagesc(kap2d);

%% solve anisotropic part
x2 = pcg(@(x) ATARadonPM1(x,theta3,256,lambda,aLap),ATg,tol,niter );
if (Positivity)
    x2 = proxBox(x2); % positivity
end
figure(2);
subplot(2,3,6); imagesc(reshape(x2,256,256)); title(['Regularised PM1, \alpha=',num2str(lambda)]);colorbar;

%% now iterate this.
nanit = 50;

derr = zeros(nanit,1);
imerr = zeros(nanit,1);
ietol = 1e-6;
tau = 0.5; % step size
for n = 1:nanit
    
    [aLap,gim,kap2d,kap1x,kap1y] = AnisoLaplacian(reshape(x2,256,256),kapPM1,eps,false);
    figure(1);clf;
    subplot(2,2,1); imagesc(im);title('Original Image');
    subplot(2,2,2); imagesc(reconFBP3);title('Filtered Back Projection')
    subplot(2,2,3); imagesc(gim);title(['Gradient at iteration ',num2str(n)]);
    subplot(2,2,4); imagesc(kap2d);title(['\kappa at iteration ',num2str(n)]);
    
    gnew = radon(reshape(x2,256,256),theta3);
    r = ATg - reshape(iradon(gnew,theta3,'linear','none',1,256),[],1);
    r = r + lambda*aLap*x2;
    derr(n) = norm(g3-gnew)^2;
    dx =  pcg(@(x) ATARadonPM1(x,theta3,256,lambda,aLap),r,tol,niter );
    x2 = x2 + tau*dx;
    if (Positivity)
        x2 = proxBox(x2); % positivity
    end
    imerr(n) = norm(tau*dx)^2;
    figure(2);
    subplot(2,3,6); imagesc(reshape(x2,256,256)); title(['Regularised PM1, \alpha=',num2str(lambda)]);colorbar;
    
    drawnow;
    % stop with discrepancy or change in image
    if(derr(n)/ndat< noiselevel || imerr(n)/imerr(1) < ietol)
        disp(['converged at iteration ',num2str(n), ' DP ',num2str(derr(n)/ndat-noiselevel)]);
        break;
    end
    eps = max(5e-2,0.8*eps);
end
figure(3); clf;
subplot(1,2,1);plot(derr(1:n)); title('data error');
subplot(1,2,2);plot(imerr(1:n)); title('image update norm');

x0im = reshape(x0,256,256);
x1im = reshape(x1,256,256);
x2im = reshape(x2,256,256);

figure(4);clf;hold on;
plot(im(100,:),'k');
plot(x0im(100,:),'b');
plot(x1im(100,:),'r');
plot(x2im(100,:),'g');
legend('true','TK0','TK1','PM1');

%%
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 z = ATARadonTK0(x,theta,N,lambda)

y = radon( reshape(x,N,N),theta);
z = iradon(y,theta,'linear','none',1,256);
z = reshape(z,[],1) + lambda*x;
end
%%
function z = ATARadonTK1(x,theta,N,lambda)

y = radon( reshape(x,N,N),theta);
z = iradon(y,theta,'linear','none',1,256);
[gx,gy]   = gradient(reshape(x,N,N));
[gxx,gxy] = gradient(gx);
[gyx,gyy] = gradient(gy);
z = reshape(z - lambda*(gxx+gyy),[],1);
end
%%
function z = ATARadonPM1(x,theta,N,lambda,aLap)

y = radon( reshape(x,N,N),theta);
z = iradon(y,theta,'linear','none',1,256);
z = reshape(z,[],1) - lambda*aLap*x;
end