% Waring! this script involves running 72 MC simulations, this can take
% 5-10 min on an integrated GPU.

clear all;
close all;
clc;

% Comment the below line if you have an NVIDIA GPU
USE_MCXCL=1;

%% define medium

% sphere
[yy, xx, zz] = meshgrid(1:40,1:40,1:20);
S1 = sqrt((xx-22).^2+(yy-10).^2+(zz-10).^2)<=6;
S2 = sqrt((xx-20).^2+(yy-30).^2+(zz-10).^2)<=6;
homo = S1+1;
homo(17:26,25:34,7:14) = 2;

cfg.vol=uint8(homo);
cfg.issrcfrom0 = 1;
cfg.isreflect = 1;

%% define time-gates

cfg.tstart=0;
cfg.tend=5e-9;
cfg.tstep=5e-9;

%% define source/detector positions

[x,y] = meshgrid([7.5:5:32.5],[7.5:5:32.5]);  % 6*6 point sources & detectors
posx = x(:);    posy = y(:);
srcz = 0;   % source on z=0 mm
detz = 20;  % detector on z=20 mm

%% other information

cfg.nphoton = 1e6;
cfg.seed = -1;
cfg.gpuid = 1;
cfg.autopilot = 1;
nSrc = 36; 
nDet = 36;

%% simulation for source
% we get both measurements and source Green's function

Green_s = zeros(nSrc,40*40*20);
mea0 = zeros(nSrc*nDet,1);
mea1 = zeros(nSrc*nDet,1);
Green_sd = zeros(nSrc*nDet,1);

srccfg = cfg;

% background optical properties
srccfg.prop=[0 0 1 1;
        0.01 1 0.01 1.37
        0.01 1 0.01 1.37];

% media with mua contrast    
prop2 = [0 0 1 1;
        0.01 1 0.01 1.37
        0.05 1 0.01 1.37];
    
srccfg.srcdir = [0 0 1];
srccfg.detpos = [posx,posy];
srccfg.detpos(:,3) = detz;
srccfg.detpos(:,4) = 2;

    
for i=1:nSrc
    srccfg.srcpos = [posx(i),posy(i),srcz];
    tic
    [flux,detp] = mcxlab(srccfg);
    t_s = toc
    Green_s(i,:) = flux.data(:)' * srccfg.tstep;
    % trajectory and partial path length (ppl) do not change if only mu_a is
    % changed, so we can use the same ppl to compute the remaining weight
    % for homogeneous background (w0) and turbid volume (w1). This is
    % correct only when you want to reconstruct delta mu_a.
    w0 = mcxdetweight(detp.data,srccfg.prop);
    w1 = mcxdetweight(detp.data,prop2);
    for j=1:nDet
        row = (i-1)*nDet+j;
        % calculate measurement for this sd pair (both background and with perturbation)
        mea0(row) = sum(w0(detp.data(1,:)==j));
        mea1(row) = sum(w1(detp.data(1,:)==j));
        % get the detector position and source Green's function at the
        % detector (Grs,rd)
        detx = ceil(srccfg.detpos(j,1));
        dety = ceil(srccfg.detpos(j,2));
        Green_sd(row) = flux.data(detx,dety,detz) * srccfg.tstep;
    end
end

figure(1); hold on;
plot(mea0);
plot(mea1,'r');

%% simulation for detector
% we get detector Green's function

Green_d = zeros(nDet,40*40*20);

detcfg = cfg;
detcfg.prop=[0 0 1 1;
        0.01 1 0.01 1.37
        0.01 1 0.01 1.37];

detcfg.srcdir = [0 0 -1];

for i=1:nDet
    detcfg.srcpos = [posx(i),posy(i),detz];
    flux = mcxlab(detcfg);
    Green_d(i,:) = flux.data(:)' * detcfg.tstep;
end

%% Jacobian by adjoint method for absorption

J = zeros(nSrc*nDet,40*40*20);
for i=1:nSrc
    for j=1:nDet
        row = (i-1)*nDet+j;
        J(row,:) = Green_s(i,:).*Green_d(j,:)/Green_sd(row);
    end
end

%% Reconstruction
% ground truth of x is det_mua = 0.04;

x0 = (homo-1)*(0.05-0.01);

% Born Formulation
x1 = lsqr(J,(mea0-mea1)./mea0,1e-4,30);
x1_3d = reshape(x1,[40 40 20]);

% Rytov Formulation, preferred
x2 = lsqr(J,log(mea0./mea1),1e-4,30);
x2_3d = reshape(x2,[40 40 20]);

% Tikhnov Regularization
x3 = J'*((J*J'+100*eye(size(J,1)))\((mea0-mea1)./mea0));
x3_3d = reshape(x3,[40 40 20]);

figure(2);

caxis = [0 0.02];
subplot(2,2,1);
imagesc(x0(:,:,10)); axis equal;
xlim([0.5 40.5]); colorbar;
title('Groundtruth')

subplot(2,2,2);
imagesc(x1_3d(:,:,10)); axis equal;
xlim([0.5 40.5]); colorbar;
title('Born')

subplot(2,2,3);
imagesc(x2_3d(:,:,10)); axis equal;
xlim([0.5 40.5]); colorbar;
title('Rytov')

subplot(2,2,4);
imagesc(x3_3d(:,:,10)); axis equal;
xlim([0.5 40.5]); colorbar;
title('Tikhnov')