%% Code summary
% Fluence Parameter Estimation (GPU version) 
% To investigate computation time
% Input files: 'multiWavParam.mat','multiWavData.mat'
%   - 'param' contains all parameters for system, data and recon
%   - 'totalImg' contains beamformed image(data)
%   - image size: 1024 x 128 x 20 x 1 x 10
%   - (axial samples x channals x sources x frames x wavelengths)
% MinWoo Kim 
% 04/02/2020

clear all
disp('** Start !! **');

%% load data
disp('-Load data');
load('multiWavParam.mat'); % parameters
load('multiWavData.mat'); % test data

%% local parameters
disp('-- parameter setting');
sourcInx = param.Core.sourcInx;    % Full source indices
ax = param.Core.laser_wav; % Spectrum wavelengh [nm]
fiberAmp = param.Core.fiberAmp; % laser source power 
lb = param.Core.lb; % searching boundary
ub = param.Core.ub; % searching boundary
ROIz = param.Core.ROIz; % region of interest
ROIx = param.Core.ROIx; % region of interest
min_thres = param.Core.min_thres; % threshold to pick up pixels
max_thres = param.Core.max_thres; % threshold to pick up pixels
[zPNum,xPNum,scNum,FrNum,WvNum] = size(totalImg); % data size
laSourPo = param.Core.laSourPo; % laser source position
alpha = param.Core.alpha; % laser source angle [radian]
xOrg = param.PData(2).Origin(1); % image origin [lambda]
zOrg = param.PData(2).Origin(3); % image origin [lambda]
pdeltaX = param.PData(2).PDelta(1); % image resolution [lambda]
pdeltaZ = param.PData(2).PDelta(3); % image resolution [lambda]
lambda = param.Core.lambda; % wavelength
ninx = param.Core.sourcEffInx; % effective source indices
scNum2 = length(ninx); %effective source number
pL = 1; % axial average number
uasBinNum = 100; % searching bin
ueffBinNum = 100; % searching bin
totBinNum = uasBinNum*ueffBinNum; % searching bin

%% Noise correction
disp('-- Noise correction');
refImg = totalImg(:,:,:,:,1);
weiImg = ones(zPNum,xPNum);
windMat = ones(100,30);
refImg = convn(refImg,windMat,'same');
weiImg = repmat(convn(weiImg,windMat,'same'),[1,1,scNum]);
refImg = refImg./weiImg;
totalImg = totalImg(:,:,:,:,2:end)-repmat(refImg,[1,1,1,1,WvNum-1]);
WvNum = WvNum-1;

%% Selection of pixels for estimation
disp('-- Pixel selection & duplication & alignment');
% selection
aveImg = squeeze(mean(mean(mean(totalImg,5),4),3));
tempImg = zeros(size(aveImg));
tempImg(ROIz,ROIx) = aveImg(ROIz,ROIx);
aveImg = tempImg;

roiInx = find(aveImg>min_thres & aveImg<max_thres);
numRoiInx = length(roiInx);
BWExt = zeros(size(aveImg));
BWExt(roiInx) = 1;
numPoints = numRoiInx*scNum2;

tempSigSet = zeros(numRoiInx, scNum,WvNum);
for ri = 1:numRoiInx
    pixPo = roiInx(ri);
    [zi,xi] = ind2sub([zPNum,xPNum],pixPo);
    tempSigSet(ri,:,:) = double(squeeze(mean(mean(totalImg(zi+(-pL:pL),xi+(-pL:pL),:,:,:),2),1)));
end

% laser source power compensation 
tempSigSet = tempSigSet./repmat(reshape(fiberAmp,[1,scNum,1]),[numRoiInx,1,WvNum]);
    
% duplication & alignment  
tempSig = tempSigSet(:,ninx,:);
meanSig = mean(tempSig,2);
tempSig = tempSig./meanSig; %normalization
tempSig = reshape(tempSig,[numPoints,WvNum]);
tempSig2 = repmat(reshape(tempSig,[1,numPoints,WvNum]),[totBinNum,1,1]);
pixPow = repmat(meanSig.^2,[1,scNum2,1]);
pixPow = repmat(reshape(pixPow,[1,numPoints,WvNum]),[totBinNum,1,1]);

%% Tables for computing fluence distributions
disp('-- Table computation');
% positions of pixels
[X,Z] = meshgrid(1:xPNum,1:zPNum);
xPos = (xOrg+pdeltaX*(X-1))*lambda*1e3;
yPos = zeros(zPNum,xPNum);
zPos = (zOrg+pdeltaZ*(Z-1))*lambda*1e3;
xyzPos = [xPos(:),yPos(:),zPos(:)];
xyzPos = xyzPos(roiInx,:);   
xyzPosMat = repmat(reshape(xyzPos,[1,numRoiInx,1,3]),[totBinNum,1,scNum,1]);

% Search range
uas_ax = linspace(lb(3),ub(3),uasBinNum);
ueff_ax = linspace(lb(2), ub(2), ueffBinNum);
[us,ueff] = meshgrid(uas_ax,ueff_ax);
us = reshape(us,[totBinNum,1]);
ueff = reshape(ueff,[totBinNum,1]);

% Diffusion parameters
ua = (-us+sqrt(us.^2+4/3*ueff.^2))/2;
D = (3*(ua+us)).^(-1); %[mm]
freePath = 1./(ua+us); %[mm]
CR = 1.03;
zb = 2*CR*D;

% Imaginary sources
hscNum = floor(scNum/2); %10
pathPoX = zeros(totBinNum,scNum);
pathPoY = [-sin(alpha)*freePath*ones(hscNum,1)' , sin(alpha)*freePath*ones(hscNum,1)'];
pathPoZ0 = zeros(totBinNum,scNum);
pathPoZ1 = [(freePath*cos(alpha))*ones(hscNum,1)' , (freePath*cos(alpha))*ones(hscNum,1)'];
pathPoZ2 = [-(freePath*cos(alpha)+2*zb)*ones(hscNum,1)', -(freePath*cos(alpha)+2*zb)*ones(hscNum,1)'];
pathPo0 = reshape([pathPoX,pathPoY,pathPoZ0],[totBinNum,scNum,3]);
pathPo1 = reshape([pathPoX,pathPoY,pathPoZ1],[totBinNum,scNum,3]);
pathPo2 = reshape([pathPoX,pathPoY,pathPoZ2],[totBinNum,scNum,3]);
sourPo0 = repmat(reshape(laSourPo,[1,scNum,3]),[totBinNum,1,1])+pathPo0;
sourPo1 = repmat(reshape(laSourPo,[1,scNum,3]),[totBinNum,1,1])+pathPo1;
sourPo2 = repmat(reshape(laSourPo,[1,scNum,3]),[totBinNum,1,1])+pathPo2;
lasPos0 = repmat(reshape(sourPo0,[totBinNum,1,scNum,3]),[1,numRoiInx,1,1]);
lasPos1 = repmat(reshape(sourPo1,[totBinNum,1,scNum,3]),[1,numRoiInx,1,1]); %+imaginary position
lasPos2 = repmat(reshape(sourPo2,[totBinNum,1,scNum,3]),[1,numRoiInx,1,1]); %-imaginary position

% Distance between sources and pixels
disMat0 = sqrt(sum((xyzPosMat-lasPos0).^2,4));
disMat1 = sqrt(sum((xyzPosMat-lasPos1).^2,4));
disMat2 = sqrt(sum((xyzPosMat-lasPos2).^2,4));

tempDis0 = reshape(disMat0(:,:,ninx),[totBinNum,numPoints]);
tempDis1 = reshape(disMat1(:,:,ninx),[totBinNum,numPoints]);
tempDis2 = reshape(disMat2(:,:,ninx),[totBinNum,numPoints]);

% Duplication for fast computation
ueff = repmat(ueff,[1,numPoints,WvNum]);
D = repmat(D,[1,numPoints,WvNum]);

% GPU Memory
ueff = gpuArray(ueff);
D = gpuArray(D);
tempDis1 = gpuArray(tempDis1);
tempDis2 = gpuArray(tempDis2);

%% Optical Parameter Estimation using GPU
disp('-- Optical parameter estimation using GPU');
tic;
phi = (1./(4*pi*D.*tempDis1).*exp(-ueff.*tempDis1)-1./(4*pi*D.*tempDis2).*exp(-ueff.*tempDis2));
phi = phi./reshape(repmat(mean(reshape(phi,[totBinNum,numRoiInx,scNum2,WvNum]),3),[1,1,1,scNum2]),[totBinNum,numPoints, WvNum]); %normalization

tempEr = (tempSig2 - 1.*phi).^2;
[minv,minInx] = min(squeeze(mean(pixPow.*tempEr,2)),[],1);
[ueInx,uasInx] = ind2sub([ueffBinNum,uasBinNum],minInx);
toc;    

uefSet = ueff_ax(ueInx);
uasSet = uas_ax(uasInx);


%% plot
disp('-- Plot');
tempUs = uasSet;
tempUeff = uefSet;
tempUa = (-tempUs+sqrt(tempUs.^2+4/3*tempUeff.^2))/2;
tempFreePath = 1./(tempUa+tempUs);

figure(107); 
subplot(2,2,1);hold on; plot(ax,tempUeff); ylim([0.02,0.5]); grid on; title('Ueff'); set(gca,'fontsize', 12);
subplot(2,2,2);hold on; plot(ax,tempUa); ylim([0.0,0.07]); grid on; title('Ua'); set(gca,'fontsize', 12);
subplot(2,2,3);hold on; plot(ax,tempUs); ylim([0,4.5]); grid on; title('reduUs'); set(gca,'fontsize', 12);
subplot(2,2,4);hold on; plot(ax,tempFreePath); ylim([0,10]); grid on; title('FreePath'); set(gca,'fontsize', 12);   
        
disp('** End !! **');