%%% Copyright: K. Takabe, T. Hirayama, Y. Yawata, 2020 %%%
%% Script3D

[pathname] = uigetdir('','Select a directory');
cd(pathname);
mkdir(pathname,'/Signature/');
savepath = [pathname,'/Signature/'];

CRMlist = dir(['*CRM*','.tif']);
Z = length(CRMlist);

%% Segmentaion Parameters
prompt = {'Threshold for image binarization [0-10]','Upper threshold for volume of a cell (in pixels)','Lower threshold for volume of a cell',...
    'X pixel size [?¿½?¿½m]', 'Y pixel size[?¿½?¿½m]','Z pixel size [?¿½?¿½m]','Number of detectors'};
dlg_title = 'Segmentation Parameters';
num_lines = [1 70];
defaultans = {'0.01','1000','20',...
    '0.26', '0.26', '0.42', '32'};
Parameters = inputdlg(prompt,dlg_title,num_lines,defaultans);

Thres = str2num(Parameters{1});
MAX = str2num(Parameters{2});
MIN = str2num(Parameters{3});
mpp_x = str2num(Parameters{4});
mpp_y = str2num(Parameters{5});
mpp_z = str2num(Parameters{6});
EM_ch = str2num(Parameters{7});

prompt = {'Number of excitation wavelengths used for image acquisition'};
dlg_title = 'Input';
num_lines = [1 50];
defaultans = {'6'};
NumInput = inputdlg(prompt,dlg_title,num_lines,defaultans);

NumEX = str2num(NumInput{1});

Pex = [];
Dex = [];
for i = 1:NumEX
    TempPex = ['Excitation No.',num2str(i),' (nm)'];
    TempDex = [''];
    Pex = [Pex,{TempPex}];
    Dex = [Dex,{TempDex}];
end    

prompt = Pex;
dlg_title = 'Input wavelengths (in a sequence from shortest to longest wavelength)';
num_lines = [1 130];
defaultans = Dex;
EXlist = inputdlg(prompt,dlg_title,num_lines,defaultans);

EX_wl = EXlist;

%% Image processing

[H W] = size(imread(CRMlist(1).name));
CRMstack = uint16(zeros(H,W,Z));
Imstack = uint16(zeros(H,W,Z));

for k = 1:Z
    Im =  imread(CRMlist(k).name);
    Imstack(:,:,k) = Im;
    Img = imgaussfilt(Im,2.5);
    se = strel('disk',5);
    Ims = imsubtract(imadd(Img,imtophat(Img,se)),imbothat(Img,se));
        
    CRMstack(:,:,k) = Ims;
end
CRMstack2 = medfilt3(CRMstack, [1,1,3]);
BWstack = imbinarize(CRMstack2,Thres);

CC = bwconncomp(BWstack,26);
stats = regionprops3(CC,"Volume");
removeMask = [stats.Volume]>MAX;
BWstack(cat(1,CC.PixelIdxList{removeMask})) = false;
BWstack2 = bwareaopen(BWstack,MIN,26);
BWstack3 = Watershed(BWstack2);

CC2 = bwconncomp(BWstack3,26);
stats2 = regionprops3(CC2, "Centroid");
NUM=CC2.NumObjects;

%% Extraction of innate fluorescence signature based on segmented regions
EX = length(EX_wl); %The number of excitaions
EM = max(EM_ch); %Max number of detector channels

uiwait(msgbox(('Draw a rectangle to select a background region')));
figure('Name', 'Draw a rectangle to select a background region');
MIP = max(CRMstack,[],3);
[BG, BG_rect] = imcrop(imadjust(MIP));
close
%%
FL_list = zeros(EX,EM,NUM);
for k = 1:EX
    for l = 1:EM
       %% Loading fluorescence(FL) images
        FL = dir(['*',EX_wl{k},'c',num2str(l, '%02i'),'z*','.tif']);
        FL_stack=uint16(zeros(H,W,Z));
        for z=1:Z
            FL_stack (:,:,z) = imread (FL(z).name);
        end
        
       %% Calculation of background FL intensity
        BG_mask = imcrop3(FL_stack,[BG_rect(1) BG_rect(2) 1 BG_rect(3) BG_rect(4) Z-1]);
        BG_mean = mean(BG_mask(:));
        
       %% Extruction of single-cell FL intensity
        for n = 1:NUM
            mask=zeros(H,W,Z);
            mask(CC2.PixelIdxList{1,n})=1;
            ROI = FL_stack.*uint16(mask);
            FL_cell = mean(FL_stack(mask>0)) - BG_mean;
            FL_list(k,l,n) = FL_cell;
        end
    end
end

%% Export signatures as PNG format

FL_list_norm = uint16(zeros(EX,EM,NUM));
for n=1:NUM
    int_lim = stretchlim(uint16(FL_list(:,:,n)),5e-4);
    FL_list_norm(:,:,n) = imadjust(uint16(FL_list(:,:,n)),[int_lim(1) int_lim(2)],[0 1]);
    imwrite(FL_list_norm(:,:,n),[savepath,'Signature',num2str(n,'%02i'),'.png']);
end

%% Display heatmap

Signature = uint16(zeros(EX*EM,NUM));
for n = 1:NUM
    Signature(:,n) = transpose(reshape(FL_list_norm(:,:,n).',[192,1]));
end
Heatmap(Signature,pathname);

%% Display 3D map

Contour = zeros(H,W,Z);
for n = 1:NUM
    Contour(CC2.PixelIdxList{1,n}) = 1;
end
    
figure

data = Contour;
data = smooth3(data,'box',3);
xx = 0:mpp_x:(W-1)*mpp_x;
yy = 0:mpp_y:(H-1)*mpp_y;
zz = 0:mpp_z:(Z-1)*mpp_z;
[X1,Y1,Z1] = meshgrid(xx,yy,zz);
p0 = patch(isocaps(X1,Y1,Z1,data,.3),...
    'FaceColor','none','EdgeColor','none');
p1 = patch(isosurface(X1,Y1,Z1,data,.3),...
    'FaceColor','red','EdgeColor','none');
p1.FaceAlpha = 0.6;
p1.EdgeAlpha = 0.4;
p0.FaceAlpha = 0.6;
p0.EdgeAlpha = 0.4;
isonormals(X1,Y1,Z1,data,p1);
view([0 0 45]); 

daspect([1 1 1])
camlight left
lighting gouraud
xlabel('?¿½?¿½m')
ylabel('?¿½?¿½m')
zlabel('?¿½?¿½m')
ax = gca;
ax.LineWidth = 2;
ax.FontName = 'Calibri';
ax.FontSize = 18;
ax.XLim = [0 (W-1)*mpp_x];
ax.YLim = [0 (H-1)*mpp_y];
ax.ZLim = [0 (Z-1)*mpp_z];
    
hold on

Cent = stats2.Centroid;
xc = mpp_x.*Cent(:,1);
yc = mpp_y.*Cent(:,2);
zc = mpp_z.*Cent(:,3);
        
for n=1:NUM
    X = xc(n);
    Y = yc(n);
    ZZ = zc(n);
    TEX = text(X,Y,ZZ,[' ',num2str(n)],'FontSize', 20);
    TEX.Color = 'black';
    TEX.FontName = 'Calibri';
    TEX.VerticalAlignment = 'middle';
end

im = image(MIP,'CDataMapping','scaled');
im.AlphaData = 0.8;
im.XData = [0 (W-1)*mpp_x];
im.YData = [0 (H-1)*mpp_y];
colormap(gray);

hold off
set(gcf,'PaperUnits','inches','PaperPosition',[0 0 15 8])
print -dtiff 3D.tif -r100

%% Function for display of heatmap
function Heatmap(A,path)
    MaxV = max(max(A));
    figure
    h = heatmap(A,'Colormap',parula,'ColorLimits',[0 MaxV]);
    h.XLabel = 'Cells';
    h.YLabel = 'Detector Number';
    h.FontSize = 8;
    h.colorbar;
    saveas(gcf,[path,'/Heatmap.png'])
end

function L = Watershed(d)
    BW = d;
    D = bwdist(~BW, 'cityblock');
    D = -D;
    D(~BW) = Inf;
    L = watershed(D);
    L(~BW) = 0;
end