% WATERSNAKES Watersnakes as in 
% Nguyen HT, Worring M, van den Boomgaard R. Watersnakes:
% Energy-Driven Watershed Segmentation. IEEE Trans on PAMI.
% 2003;25(3):330{342. Available from: citeseer.ist.psu.edu/article/
% nguyen03segmsurfwatersnakes.html.
% WAT = WATERSNAKES(IM,MINIMA) computes the segmsurfwatersnakes in image IM around
% markers MINIMA. Returning the segmsurfwatersnakes segmentation WAT.
%
% WAT = WATERSNAKES(IM,MINIMA,'TOPO',TOPO) computes the segmsurfwatersnakes with
% the given topographical distance function TOPO
% 
function [wat] = segmsurfwatersnakes(varargin)


u = varargin{1};
minima = varargin{2};
% smoothness parameter
prm.beta = 0.1;
for i = 3 : 2 : nargin
    namevar = varargin{i};
    varhere = varargin{i+1};
    switch(namevar)
        case('topo')
            c = varhere;
        case('beta')
            prm.beta = 0.1;
        otherwise
            error('Wrong input');
    end;
end;

% scale image to get same behavious
u = scale(u);

% only make topo if not exist from before
if ~exist('c','var')
    [faser,L] = bwlabeln(minima,4);
    c = Inf(size(u));
    for i = 1 : L
        % this minima
        reghere = eq(faser,i);
    
        % this topo function
        c(:,:,i) = ift(u,reghere);    
    
    end;
    clear reghere;
end;

% show(u,3)
% showall(c)

% make watershed image
[minval,wat] = min(c,[],3);
watini = wat;

[M N P] = size(c);

scalex = round(linspace(50,M,5)');
scaley = round(linspace(50,N,5)');

% different scales
wathere = wat;
for i = 1 : length(scalex)
    
    disp(sprintf('Scale %i out of %i',i,length(scalex)));
    
    scalehere = [scalex(i) scaley(i)];
    
    % resize watershed and topographical distance function
    wathere = imresize(wathere,scalehere,'nearest');
    
    [m n] = size(wathere);
    chere = zeros(m,n,P);
    for j = 1 : P
        chere(:,:,j) = imresize(c(:,:,j),[m n],'bilinear');
    end;
    c = chere;         
    
    % smooth the watershed lines
    wathereold = wathere;
    wathere = smoothwat(wathere,c,prm);
    
    
    figure(1);
    subplot(2,2,1);imagesc(u);axis off;colormap(gray);drawnow;title('Image')
    subplot(2,2,2);imagesc(watini);axis off;colormap(gray);drawnow;title('Wat initiell')
    subplot(2,2,3);imagesc(wathereold);axis off;colormap(gray);drawnow;title('Old wat')
    subplot(2,2,4);imagesc(wathere);axis off;colormap(gray);drawnow;title('New wat')
% pause
end;

% return variable
wat = wathere;

% show(u,2)
% show(wat,2)

%------------------------------------------------

function [wat] = smoothwat(wat,c,prm)


%
% Get stability for all border pixels
%

outside = zeros(size(wat));
ind = [];
de = [];
labnei = [];
labcur =[];
count = 0;
% this is done for all again, not so fast but easy!!
while 1

    % get the stability for all pixels in border
    [de,labnei,labcur,ind] = getstabborder(wat,c,prm,outside,ind,de,labnei,labcur);

    % find the maximum energy change
    [minval,indmin] = min(de);
    
    % if the minimum change is positive then stop
    if minval > 0
        break;
    end;

    % count pixels changing phase
    count = count + 1;
    
    % assign the pixel to the neighbour
    watold = wat;
    indpix = ind(indmin);
    wat(indpix) = labnei(indmin);
   
    % dont recalcualte the outside values
    outside = zeros(size(wat));
    outside(indpix) = 1;
    load ball2;se = getball(ball,2,1);
    outside = imcomplement(imdilate(outside,se));
    
% show(outside,3)
% pause
%     show(watold,1)
%     show(wat,2);
%     pause
    
end;

disp(sprintf('Number of pixels changing phase: %i',count))

%--------------------------------------------------

function [de,labnei,labcur,ind] = getstabborder(wat,c,prm,outside,indold,deold,labneiold,labcurold)

% the perimeter of watershed regions to get the stability for
ind = perimwat(wat);

% the coordinates to use in neighbour considerations
[kperim(:,1) kperim(:,2)] = ind2sub(size(wat),ind);

% the indices outside the neighbourhood to test whether we need to update
% here or not
indoutside = find(outside);

% get stability value for every pixel in the boundaries
numperim = length(ind);    
de = zeros(numperim,1);
labcur = zeros(numperim,1);
labnei = zeros(numperim,1);
for i = 1 : numperim
    
    % this coordinate
    k = kperim(i,:);
    
    % this indices to compare if we have been here before, and if we are
    % outside the neighbourhood, then we dont waste time here and continue!!
    indhere = ind(i);

    % the intersection to the old perimeter pixels outsid the neighboruhood
    inters = intersect(indoutside,indhere);

    % get stability value if we are not outside the neighbourhood
    if isempty(inters)            
        [de(i) labnei(i) labcur(i)] = getstabpxl(k,c,wat,prm);
    else
        % just keep the old values since we are outside the neighboorhood,
        % no need to update here
        indlocal = find(indold == inters);
%         indold
%         inters
%         indlocal
        de(i) = deold(indlocal);
        labnei(i) = labneiold(indlocal);
        labcur(i) = labcurold(indlocal);
        ind(i) = indold(indlocal);
    end;
        
    
end;


%------------------------------------------------

function [de,labnei,labcur] = getstabpxl(k,c,wat,prm)
k1 = k(1);
k2 = k(2);

% find the neighbour regions
dil = zeros(size(wat));
dil(k1-2:k1+2,k2-2:k2+2) = 1;
load ball1;se = getball(ball,1,1);
labcur = wat(k1,k2);
while 1

    labnei = unique(wat(dil == 1));
    labnei(labnei == labcur) = [];
    % we have a neighbour
    if length(labnei >=1)
        labnei = labnei(1);
        break;
    end;
    dil = imdilatefast(dil,se);
end;
    

% get the stability
regcur = eq(wat,labcur);
[n4c,n8c,nkc] = getfilter(regcur,k);

regnei = eq(wat,labnei);
[n4n,n8n,nkn] = getfilter(regnei,k);


% curvature term
dtopo = c(k(1),k(2),labnei) - c(k(1),k(2),labcur);

% energy term
w4 = 0.26;
w8 = 0.19;
wk = 0.06;
curv = prm.beta*(w4*(n4n-n4c) + w8*(n8n-n8c) + wk*(nkn-nkc));

% update
de =  dtopo + curv;
% dtopo
% curv
% de
% pause

% regcur = double(regcur);
% regcur(k1,k2) = 2;
% show(regcur,2)
% k
% labnei
% labcur
% [n4c n4n]
% [n8c n8n]
% [nkc nkn]
% dtopo
% curv
% de
% pause
 %--------------------------------------------------

function [n4,n8,nk] = getfilter(reghere,k)

k1 = k(1);
k2 = k(2);
reghere(k1,k2) = 1;
outside = imcomplement(reghere);
k = [k1-1 k2;   ...
     k1+1 k2;   ...
     k1   k2-1; ...
     k1   k2+1];
ind4 = sub2ind(size(reghere),k(:,1),k(:,2));
n4 = sum(outside(ind4));

k = [k1-1 k2-1; ...
     k1+1 k2-1; ...
     k1-1 k2+1; ...
     k1+1 k2+1];
ind8 = sub2ind(size(reghere),k(:,1),k(:,2));
n8 = sum(outside(ind8));

k = [k1-1 k2-2; ...
     k1-2 k2-1; ...
     k1-2 k2+1; ...
     k1-1 k2+2; ...
     k1+1 k2+2; ...
     k1+2 k2+1; ...
     k1+2 k2-1; ...
     k1+1 k2-2];
indk = sub2ind(size(reghere),k(:,1),k(:,2));
nk = sum(outside(indk));

function [ind] = perimwat(wat)
    
% the perimeter of watershed regions
valwat = unique(wat(wat > 0));
numwat = length(valwat);
perimall = zeros(size(wat));
for i = 1 : numwat
    reghere = eq(wat,valwat(i));
    perim = bwperim(reghere);
    perimall(perim) = 1;
end;
perim = perimall;
perim = constborder(perim,1,0);
ind = find(perim);


%------------------------------------------------


% reghere = double(reghere);
% reghere(k1,k2) = 2;
% show(reghere,1)
% show(outside,2)
% k1
% k2
% n4
% n8
% nk
% pause




% 
% % iterate until no more negative change in energy
% dehere = -Inf;
% while 1
%     
%     % find the maximum energy change
%     [minval,indmin] = min(de);
%     
%     if minval > 0
%         break;
%     end;
%     
%     % assign the pixel to the neighbour
%     wat(ind(indmin)) = neigh(indmin);
%     
%     
%     % recompute the stability in the neighbourhood
%     [k1 k2] = kperim(indmin,:);
%     
%     p = 5;
%     dil = zeros(size(wat));
%     dil(k1-p:k1+p,k2-p:k2+p) = 1;
%     perim = bwperim(gt(wat,0));
%     indhere = find(perim.*dil);
%     
%     % the common ones to be replacestrd
%     inters = intersect(indhere,ind)
%     
%     % local labelling
%     indhere = ind == inters;
%     
%     de(indhere) = [];
%     neigh(indhere) = [];
%     label(i) = [];
%     numperimhere = length(indhere);
%     kperimhere = ind2sub(size(wat),inters);
%     for i = 1 : numperimhere
%         % this coordinate
%         k = kperimhere(i,:);
%     
%     
%         % get stability value
%         [dehere neighhere labelhere] = getstab(k,c,wat);
%     end;
%     de = [de ; dehere];
%     neigh = [neigh ; neighhere];
%     label = [label ; labelhere];
%     
% end;
