%
% Example of 3D curvature in a voxel array
%
USEMRI = true; % false;
if(USEMRI)
%addpath('C:/Users/Simon Arridge/Desktop/TOAST/MGH_AtlasProject');
%pathname='C:/Users/Simon Arridge/Desktop/TOAST/MGH_AtlasProject/';
addpath('S:/TOAST/MGH_AtlasProject');
pathname='S:/TOAST/MGH_AtlasProject/';

hseg=load_bin([pathname 'Subject/subject_hseg.bin'], [256 256 256]);
end

GRADMASK = true; % set to 1 to mask out gradient zero
beta = 5e-2;     % Threshold for gradient mask.
%
nx = 128;
ny = 128;
nz = 128;

% add slow background to prevent divide by zero
G1 = (1 - exp( -(nx/2- [1:nx]).^2./(nx^2/2)));


% build a synthetic 3D image
back = zeros(nx,ny,nz);
noise = 5*rand(nx,ny,nz);
im3d = back;
zstart = floor(nx/4); ystart = floor(ny/4); xstart =floor(nz/4); 
zw = floor(nz/4); yw = floor(ny/4); xw = floor(nx/4); 
zcen1 = floor(2 *nz/4); ycen1 = floor(3 *ny/4); xcen1 = floor(3*nx/4);
zcen2 = floor(3 *nz/4); ycen2 = floor(2 *ny/4); xcen2 = floor(nx/4);


for p = 1 :nz
    for m = 1:ny
        for n = 1:nx
            back(n,m,p) = 30*G1(n)*G1(m)*G1(p); % slow Gaussian
            if(m - ycen2)^2/2 + (n - xcen2)^2 + (p - zcen2)^2 < 0.8 *xw^2
                im3d(n,m,p) = im3d(n,m,p) + 150;
            end
            if(m - ycen1)^2 + (n - xcen1)^2 + (p - zcen1)^2/2 < 0.9 *xw^2
                im3d(n,m,p) = im3d(n,m,p) + 200;
            end

        end
    end
end
for p = zstart :zstart + zw
    for m = ystart :ystart + yw
        for n = xstart :xstart + xw
            im3d(n,m,p) = im3d(n,m,p) + 100;
        end
    end
end

if(USEMRI)
im3d = noise + back+double(64*hseg(1:2:256,1:2:256,1:2:256));
else
    im3d = back + im3d +noise;
end

% smooth it
% use Fourier multiplication method
sig = 1.5/128; % units of pixels
fim3 = fftshift(fftn(im3d));
kx = 2*pi*[1-nx/2:nx/2];
ky = 2*pi*[1-ny/2:ny/2];
kz = 2*pi*[1-nz/2:nz/2];
gf2 = exp( - kx'.^2 *sig^2/2)* exp( - ky.^2 *sig^2/2);
gfz = exp( - kz.^2 *sig^2/2); % z direction
for z = 1:nz
    fim3(:,:,z) = gfz(z)* (fim3(:,:,z).*gf2);
end
im3d = real(ifftn( ifftshift(fim3)));

%H2 = fspecial('gaussian',[6 6 ],sig);
%H3 = reshape([H2/4; H2/2; H2/4],6 ,6,3);
%im3d = imfilter(im3d,H3);

%%
figure(1);clf
for p = 1:nz
    imagesc(im3d(:,:,p),[0,255]);colorbar;title(['original image ',num2str(p)]);
    pause(0.1);
end

%% image gradient

[gx,gy,gz] = gradient(im3d);
g = sqrt(gx.^2 + gy.^2 + gz.^2);
g1d = reshape(g,1,nx*ny*nz);
g2 = sqrt(gx.^2 + gy.^2);
g21d = reshape(g2,1,nx*ny*nz);
gmax = max(g1d);

zind = find(g < beta * gmax);        % edges less than a threshold;
nzind = setdiff([1:nx*ny*nz],zind);
zmask = ones(nx*ny*nz,1);
zmask(zind) = 0;
zmask = reshape(zmask,nx,ny,nz);

% dilate the mask 

[mx,my,mz] = gradient(zmask); zind1 = find(mx ~= 0 | my ~=0 | mz ~= 0);
zind1 = union(zind,zind1);
mask1 = zmask; mask1(zind1) = 0;
% and again...
[mx,my,mz] = gradient(mask1); zind2 = find(mx ~= 0 | my ~=0 | mz ~= 0);
zind2 = union(zind1,zind2);
mask2 = mask1; mask2(zind2) = 0;

figure(1);clf
for p = 1:nz
    imagesc(g(:,:,p),[0,gmax]);colorbar;title(['gradient slice ',num2str(p)]);
    pause(0.1);
end

beta = 0.1*gmax;
if(GRADMASK)
    skap = zeros(nx*ny*nz,1);
    skap(nzind) = 1./g1d(nzind);
    skap = reshape(skap,nx,ny,nz);
else
    skap = 1./g;
end
skmax = max(skap(:));
%% display masks
for p = 1:nz
    figure(4);clf; colormap('hot');
    subplot(2,3,1)
    imagesc(im3d(:,:,p),[0,255]);colorbar;title('Image');
    subplot(2,3,2)
    imagesc(g(:,:,p),[0,gmax]);colorbar;title(['gradient ',num2str(p)]);
    subplot(2,3,3);
    imagesc(skap(:,:,p),[0,skmax]);colorbar;title(['reciprocal gradient ',num2str(p)]);
    
    subplot(2,3,4)
    imagesc(zmask(:,:,p),[0,1]);colorbar;title('gradient mask');
    subplot(2,3,5)
    imagesc(mask1(:,:,p),[0,1]);colorbar;title('mask dilated once');
    subplot(2,3,6)
    imagesc(mask2(:,:,p),[0,1]);colorbar;title('mask dilated twice');
    pause(0.1);
end
%% 

sx = skap.*gx;  sy = skap.*gy;   sz = skap.*gz; % direction cosines
[gxx,gxy,gxz] = gradient(gx);
[gyx,gyy,gyz] = gradient(gy);
[gzx,gzy,gzz] = gradient(gz);
[sxx,sxy,sxz] = gradient(sx);
[syx,syy,syz] = gradient(sy);
[szx,szy,szz] = gradient(sz);


if GRADMASK
    sx(zind1) = 0; sy(zind1) = 0;sz(zind1) = 0;

    sxx(zind2)= 0; sxy(zind2)= 0;sxz(zind2)= 0; 
    syx(zind2)= 0; syy(zind2)= 0;syz(zind2)= 0; 
    szx(zind2)= 0; szy(zind2)= 0;szz(zind2)= 0;
    
    gx(zind1) = 0; gy(zind1) = 0;gz(zind1) = 0;

    gxx(zind2)= 0; gxy(zind2)= 0;gxz(zind2)= 0; 
    gyx(zind2)= 0; gyy(zind2)= 0;gyz(zind2)= 0; 
    gzx(zind2)= 0; gzy(zind2)= 0;gzz(zind2)= 0;

end
%% divergence of gradient normal

%imkap = divergence(skap.*gx,skap.*gy,skap.*gz);
imkap = sxx + syy + szz;

kdmin = min(min(min(imkap)));
kdmax = max(max(max(imkap)));


%% create tangential derivative vectors
if(GRADMASK)
    plen = zeros(nx*ny*nz,1);
    plen(nzind) = 1./g21d(nzind);
    plen = reshape(plen,nx,ny,nz); % this is normalisation for p
else
    plen = 1./g2;
end
gpx = -gy.*plen; gpy = gx.*plen; gpz = zeros(size(gz)); % unit vector phat

qlen = plen.*skap;             % this is normalisation for q
gqx = -gx.*gz.*qlen; gqy = -gy.*gz.*qlen; gqz = (gx.^2 + gy.^2).*qlen; % unit vector qhat;

% this is the 2D code

% [g2txx g2txy] = gradient(-g2y.*n2len);
% [g2tyx g2tyy] = gradient(g2x.*n2len);
% g2ttx = (g2y.*g2txx - g2x.*g2txy);
% g2tty = (g2y.*g2tyx - g2x.*g2tyy);
% k2dc = g2x.*g2ttx + g2y.*g2tty;
% k2dc = k2dc.* n2len.^2;

% tangential 2nd derivative vectors 



[gpxx,gpxy,gpxz] = gradient(gpx);
[gpyx,gpyy,gpyz] = gradient(gpy);
[gpzx,gpzy,gpzz] = gradient(gpz);

[gqxx,gqxy,gqxz] = gradient(gqx);
[gqyx,gqyy,gqyz] = gradient(gqy);
[gqzx,gqzy,gqzz] = gradient(gqz);

if GRADMASK
    gpx(zind1) = 0; gpy(zind1) = 0;gpz(zind1) = 0;

    gpxx(zind2)= 0; gpxy(zind2)= 0;gpxz(zind2)= 0; 
    gpyx(zind2)= 0; gpyy(zind2)= 0;gpyz(zind2)= 0; 
    gpzx(zind2)= 0; gpzy(zind2)= 0;gpzz(zind2)= 0;
    
    gqx(zind1) = 0; gqx(zind1) = 0;gqx(zind1) = 0;

    gqxx(zind2)= 0; gqxy(zind2)= 0;gqxz(zind2)= 0; 
    gqyx(zind2)= 0; gqyy(zind2)= 0;gqyz(zind2)= 0; 
    gqzx(zind2)= 0; gqzy(zind2)= 0;gqzz(zind2)= 0;
    
end

gppx = gpx.*gpxx + gpy.*gpxy;
gppy = gpx.*gpyx + gpy.*gpyy;
gppz = zeros(size(gz));
gpp = (gx.*gppx + gy.*gppy).*skap;

gpqx = gpx.*gqxx +gpy.*gqxy;
gpqy = gpx.*gqyx +gpy.*gqyy;
gpqz = gpx.*gqzx +gpy.*gqzy;
gpq = (gx.*gpqx + gy.*gpqy + gz.*gpqz).*skap;

gqpx = gqx.*gpxx +gqy.*gpxy +gqz.*gpxz;
gqpy = gqx.*gpyx +gqy.*gpyy +gqz.*gpyz;
gqpz = gqx.*gpzx +gqy.*gpzy +gqz.*gpzz;
gqp = (gx.*gqpx + gy.*gqpy + gz.*gqpz).*skap;

gqqx = gqx.*gqxx + gqy.*gqxy + gqz.*gqxz;
gqqy = gqx.*gqyx + gqy.*gqyy + gqz.*gqyz;
gqqz = gqx.*gqzx + gqy.*gqzy + gqz.*gqzz;
gqq = (gx.*gqqx + gy.*gqqy + gz.*gqqz).*skap;

kgauss = zeros(nx,ny,nz);
kmean = zeros(nx,ny,nz);
kmin = zeros(nx,ny,nz);
kmax = zeros(nx,ny,nz);

% 
%% calculate curvature tensor per voxel
for p = 1 :nz
    disp(['processing slice ',num2str(p)]);
    for m = 1:ny
        for n = 1:nx
            H =  [gpp(n,m,p) ,  0.5*(gpq(n,m,p)+gqp(n,m,p)); 0.5*(gpq(n,m,p)+gqp(n,m,p)),  gqq(n,m,p)];
%            H =  [gpp(n,m,p) ,  gqp(n,m,p); gqp(n,m,p),  gqq(n,m,p)];
            dh = eig(H);
            if(abs(dh(2)) >= abs(dh(1)))
                kmax(n,m,p) = dh(2); kmin(n,m,p) = dh(1);
            else
                kmax(n,m,p) = dh(1); kmin(n,m,p) = dh(2);
            end
            kgauss(n,m,p) = det(H);%*skap(n,m,p);
            kmean(n,m,p) = 0.5*trace(H);%*skap(n,m,p);
        end
    end
end

k1min = min(min(min(kmax)));
k1max = max(max(max(kmin)));
k2min = min(min(min(kmean)));
k2max = max(max(max(kmean)));
kgmin = min(min(min(kgauss)));
kgmax = max(max(max(kgauss)));
kmmin = min(min(min(kmean)));
kmmax = max(max(max(kmean)));

%%
for p = 1:nz
    figure(2);clf; colormap('hot');
    subplot(2,3,1)
    imagesc(imkap(:,:,p),[kdmin,kdmax]);colorbar;title('divergence of normal');
    subplot(2,3,2)
    imagesc(kgauss(:,:,p),[kgmin,kgmax]);colorbar;title('gauss curvature');
    subplot(2,3,3)
    imagesc(kmean(:,:,p),[kmmin,kmmax]);colorbar;title('mean curvature');
    subplot(2,3,4)
    imagesc(kmax(:,:,p),[k1min,k1max]);colorbar;title('max curvature');
    subplot(2,3,5)
    imagesc(kmin(:,:,p),[k2min,k2max]);colorbar;title('min curvature');
    subplot(2,3,6)
    imagesc(g(:,:,p),[0,gmax]);colorbar;title(['gradient ',num2str(p)]);

%     if GRADMASK
%         %imagesc(zmask(:,:,j));title('mask');
%         
%         imagesc(zmask(:,:,p),[0,1]);colormap('gray');colorbar;title(['mask slice ',num2str(p)]);
%     else
%         imagesc(skap(:,:,p),[0,1]);colormap('gray');colorbar;title(['\kappa ',num2str(p)]);
%     end
    pause(0.1);
end


for p = 1:nz
    figure(3);clf
    subplot(2,3,1)
    imagesc(gpp(:,:,p),[kmmin,kmmax]);colormap('gray');colorbar;title('KPP');
    subplot(2,3,2)
    imagesc(gpq(:,:,p),[kmmin,kmmax]);colormap('gray');colorbar;title('KPQ');
    subplot(2,3,4)
    imagesc(gqp(:,:,p),[kmmin,kmmax]);colormap('gray');colorbar;title('KQP');
    subplot(2,3,5)
    imagesc(gqq(:,:,p),[kmmin,kmmax]);colormap('gray');colorbar;title('KQQ');
    subplot(2,3,6)
    imagesc(g(:,:,p),[0,gmax]);colormap('gray');colorbar;title(['gradient ',num2str(p)]);

    pause(0.1);
end

%% tests
ix = 120; iy = 100; iz = 37;
nhat = ([gx(ix,iy,iz),gy(ix,iy,iz),gz(ix,iy,iz)]).*skap(ix,iy,iz);
phat = ([gpx(ix,iy,iz),gpy(ix,iy,iz),gpz(ix,iy,iz)]);
qhat = ([gqx(ix,iy,iz),gqy(ix,iy,iz),gqz(ix,iy,iz)]);
disp(['normal at (',num2str(ix),',',num2str(iy),',',num2str(iz),'): ',num2str(nhat)]);
disp(['phat at (',num2str(ix),',',num2str(iy),',',num2str(iz),'): ',num2str(phat)]);
disp(['qhat at (',num2str(ix),',',num2str(iy),',',num2str(iz),'): ',num2str(qhat)]);