%% InverseProblems/rundiff2d;

% load images.
impath = '../../StandardTestImages/';
im = double(imread([impath 'Cameraman256.png'],'PNG'));
%im = double(imread([impath 'ot.pgm'],'pgm'));
%im = double(imread([impath 'Lena512.png'],'PNG'));

%%
[n1,n2] = size(im);
N = n1; % assume square.
oneN = ones(N+1,1);
%Dx = sparse(N-1,N);
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);

% this form of the Laplacian (D2x) implicitly assumes Dirichlet b.c.s of 
% zero on the domain outside the interval [1-N]. I.e. f(0) = f(N+1) = 0

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


Lapl = -(D1x2d'*D1x2d + D1y2d'*D1y2d);

% run diffusion
imlap = reshape( Lapl*reshape(im,[],1),n1,n2);
imx = reshape( D1x2d*reshape(im,[],1),n1+1,n2);
imy = reshape( D1y2d*reshape(im,[],1),n1,n2+1);

figure(1);clf; hold on;
subplot(2,3,1); imagesc(im); colormap(gray);title('original f')
subplot(2,3,2); imagesc(imx); colormap(gray);title('f_x');
subplot(2,3,3); imagesc(imy); colormap(gray);title('f_y');
subplot(2,3,5); imagesc(imlap); colormap(gray);title('Lapl f');


%% Fully Explicit method

h = reshape(im,[],1);
dt = 0.25;  % dt<= 0.25  for stability 
kap = 1;
A = speye(n1*n2) + dt/kap .* Lapl;
tic;
for k = 1:100
    h = A*h;
    figure(2); clf;
    imagesc(reshape(h,n1,n2));colormap(gray);title(['it ',num2str(k)]);
    pause(0.1);
end
toc;
%% Fully implicit method
h = reshape(im,[],1);
dt = 0.25;%2.5;
kap = 1;
A = speye(n1*n2) - dt/kap .* Lapl;
tic;
for k = 1:100
    h = A\h;
    figure(3); clf;
    imagesc(reshape(h,n1,n2));colormap(gray);title(['it ',num2str(k)]);
    pause(0.1);
end
toc;

%% Semi-Implicit !
Lx = -D1x2d'*D1x2d;                ;%Dx2'*Dx2;
Ly = -D1y2d'*D1y2d;
h = reshape(im,[],1);
dt = 0.25;  % dt<= 0.25  for stability 
kap = 1;
A1 = speye(n1*n2) - 0.5*dt/kap .* Lx;
A2 = speye(n1*n2) + 0.5*dt/kap .* Ly;
tic;
for k = 1:100
    h = A1\(A2*h);
    h = reshape(reshape(h,n1,n2)',[],1);
    h =  A1\(A2*h);
    h = reshape(reshape(h,n1,n2)',[],1);
    
    figure(4); clf;
    imagesc(reshape(h,n1,n2));colormap(gray);title(['it ',num2str(k)]);
    pause(0.1);
end
toc;

%% include the 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));
%% make an anisotropic version
kapfunc = @(s,T) kapPM1(s,T);
[gx,gy] = gradient(im);
gim = sqrt(gx.^2 + gy.^2);
gmax = max(gim(:));
T = 0.4*gmax;
[aLap,gim2,kapim,k1x,k1y] = AnisoLaplacian(im,kapfunc,T,true);
figure(1); subplot(2,3,4); imagesc(gim2);title('gradient');
figure(1); subplot(2,3,6); imagesc(kapim);title('\kappa');

%% anisotropic diffusion (implicit);
h = reshape(im,[],1);
dt = 0.25;
kap = 1;
AA = speye(n1*n2) - dt/kap .* aLap;
tic;
for k = 1:100
    h = AA\h;
    figure(5); clf;
    imagesc(reshape(h,n1,n2));colormap(gray);title(['it ',num2str(k)]);colorbar;
    pause(0.1);
end
toc;
%% make an nonlinear version
% anisotropic diffusion
im = im.*(1+0.05*randn(size(im)));
h = reshape(im,[],1);
dt = 2;
kap = 1;
nt = 100;
hss = zeros(n1,n2,nt+1);
hss(:,:,1) = im;
figure(6); clf;
SSrate = 0.025; % make this smaller to diffuse slower.
    
for k = 1:nt
    [gx,gy] = gradient(reshape(h,n1,n2));
    gim = sqrt(gx.^2 + gy.^2);
    gmax = max(max(gim));
    T = SSrate*gmax;
    [aLap,gim2,kappa] = AnisoLaplacian(hss(:,:,k),kapfunc,T,false);
    figure(6); subplot(2,2,3); imagesc(gim2);title('\nabla f'); colorbar;
    figure(6); subplot(2,2,4);  imagesc(kappa);title('\kappa'); colorbar;
 
    AA = speye(n1*n2) - dt/kap .* aLap;    
    h = AA\h;
    hss(:,:,k+1) =reshape(h,n1,n2); 
 
    figure(6);subplot(2,2,2); 
    imagesc(hss(:,:,k+1));colormap(gray);title(['it ',num2str(k)]);colorbar;
    pause(0.1);
end
figure(6); subplot(2,2,1);
ylin = 128;
imagesc(squeeze(hss(:,128,:)));title(['Scale space, row ',num2str(ylin)]);

%%

figure(7); clf;

for k = 1:nt+1
    figure(7); clf;
    imagesc(hss(:,:,k));colormap(gray);title(['it ',num2str(k)]);
pause(0.4);
end

%%
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);

%gmax = max(gim(:));
%T = gmax/4;

% define kappa
%kappa = 0.5 + 0.5* sin(6*pi/256*[1:256]') * sin( 4*pi/256*[1:256]);
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