% Do regularised inversion using Richardson-Lucy
%
%
addpath('../','../Lin1D');
n = 64;
sb = 0.03; % blurring width
%sigma = 0.1; % std.dev. of added white noise
f = zeros(n,1);
f(12) = 1; f(20:28) = 1.5; f(32:36) = 2; % arbitrary function. Must be non-negative
[y,K] = linblur(f,sb);

LowCount = false; % set to zero for 'High Count' simultation
if LowCount
    yp = 1e11*imnoise(1e-11*y,'poisson');
else
    yp = 1e9*imnoise(1e-9*y,'poisson');

end
figure (1);clf
hold on;
plot(y);
plot(yp,'g');
plot(f,'r');


%%
figure(3);clf;
niter = 30;
alpha = 1e-1;
plot(f,'r');
hold on;
fi = ones(size(f)); % initial guess
ks = K'*ones(size(y));
derr = zeros(niter,1);
ferr = derr;
KL = derr;
eps = 1e-9;
for k=1:niter 
    fi = (fi./ks) .* (K'*(yp./(K*fi))); 
    figure(3);title('reconstruction');
    plot(fi,'k');
    pause(1);
    figure(1);
    plot(K*fi,'k');title('data');
    z = K*fi;
    iz =find(z < eps);
    iy = find(yp <eps);
    it = find(f < eps);
    ir = find(fi < eps);
    ij = union(ir,it);
    ii = union(iz,iy);
    zz = yp./z;
    ff = f./fi;
    zz(ii) = 1;
    ff(ij) = 1;
    derr(k) = norm(ones(size(y)) - zz).^2;
    ferr(k) = norm(ones(size(f)) - ff).^2;
    KL(k) = ones(size(yp))' * ( yp.* log(zz) + (z - yp) );
    pause(0.1);
end
frl = fi;

derr1 = derr; ferr1=ferr; KL1 = KL;
%%
figure(11); clf; hold on; 
semilogy(KL1,'r');
semilogy(derr1,'k');
semilogy(ferr1,'b');
legend('Kullback-Leibler divergence','relative L2 error in data ||1-g/Af||^2','relative solution error ||1-f/f_{true}||^2'); 
%% compare to Least-Square 
D1 = lindf(f); 
D2 = D1'*D1;
alphaML = 0; %  alpha;

flsqML= (K'*K + alphaML*D2)\(K'*yp);

% and weighted Least Square : 
eps = 1e-6;
ws = diag(1./(sqrt(yp)+eps));
fwlsqML= (K'*ws*ws*K + alphaML*D2)\(K'*ws*ws*yp);

figure(4);clf;
hold on;
plot(frl,'k');
plot(f,'r');
plot(flsqML,'b');
plot(fwlsqML,'--m');
legend('Richardson-Lucy','f_{true}','Least-Squares','Weighted Least-Squares');

%% regularised RL : iterative EM-MAP
D1 = lindf(f); 
D2 = D1'*D1;
fi = ones(size(f)); % initial guess
tau = alpha;
k = 1; rr = 2; er = 1;ifnz = [1:n];
%while abs(rr - 1) > 1e-3 && k < 20
while er > 1e-4 && k < niter
    disp(['iteration ',num2str(k),' relative change ', num2str(rr),' absolute change ', num2str(er)]);
    finext = (fi./(ks+tau*D2*fi)) .* (K'*(yp./(K*fi)));
    rr = norm(finext(ifnz)./fi(ifnz),1)/n;
    er = norm(finext - fi)/n;
    flast = fi;
    fi = finext;
    figure(3);
    plot(fi,'m');
    pause(1);
    figure(1);
    plot(K*fi,'--m');
    z = K*fi;
    iz =find(z < eps);
    iy = find(yp <eps);
    it = find(f < eps);
    ir = find(fi < eps);
    ifnz = setdiff([1:n],ir); % index non-zeros in solution
    ij = union(ir,it);
    ii = union(iz,iy);
    zz = yp./z;
    ff = f./fi;
    zz(ii) = 1;
    ff(ij) = 1;
    derr(k) = norm(ones(size(y)) - zz);
    ferr(k) = norm(ones(size(f)) - ff);
    KL(k) = ones(size(yp))' * ( yp.* log(zz) + (z - yp) );
    pause(0.1);
    k = k+1;
end
frlrTK1 = fi;

% compare to Least-Squares
flsq= (K'*K + alpha*D2)\(K'*yp);

% and weighted Least Square : 
eps = 1e-6;
ws = diag(1./(sqrt(yp)+eps));
fwlsq= (K'*ws*ws*K + alpha*D2)\(K'*ws*ws*yp);


%%
figure(5);clf;
hold on;
plot(frl,'k');
plot(f,'r');
plot(flsq);
plot(fwlsq,'m');
plot(frlrTK1,'g');
legend('Richardson-Lucy','f_{true}','Least-Squares','Weighted Least-Squares','One-step-Late');

derr2 = derr; ferr2=ferr; KL2 = KL;
figure(12); clf; hold on; 
title('TK1 Regularised Richardson-Lucy (One step late)');
semilogy(KL2,'r');
semilogy(derr2,'k'); 
semilogy(ferr2,'b');
legend('Kullback-Leibler divergence','relative L2 error in data ||1-g/Af||^2','relative solution error ||1-f/f_{true}||^2'); 


%% Newton type ?
% 
% fi = ones(size(f)); % initial guess
% 
% for k=1:10
%     activeset = find(fi < 0.01);
%     b = zeros(size(f));
%     b(activeset) = 1;
% %    H = -diag(fi)*K'*diag(yp./((K*fi).^2) )*K*diag(fi) + diag(b);
%     H = -K'*diag(yp./((K*fi).^2) )*K + diag(b);
% 
% %    fi = fi + (H\(K' * ((ones(size(y))-yp./(K*fi)))) );
%     fi = fi./(H\(K'*ones(size(y)))) .* (H\(K'*(yp./(K*fi)))); 
%     figure(3);
%     plot(fi,'m');
%     pause(1);
%     figure(1);
%     plot(K*fi,'m');
%     pause(3);
% end
% fpN = fi;
% figure(4);
% plot(fpN,'m');