function lambdaf = rw_seg(F,labels,opt,lambda0)
% 
%RW_SEG Learn a segmentation using the Meila & Shi random walk method
%
% LAMBDAF = RW_SEG(F,LABELS,OPT,LAMBDA0)
%
% F                 Cell array with Q elements containing the similarity matrices 
%                   based on different cues
% F{q}.data  I x I  Similarity matrix for cue q
% LABELS     I x 1  Array of segmentation labels (1/0)
% OPT        3 x 1  Parameters for the gradient ascent optimization
% OPT(1) = nu       Step size for gradient ascent
% OPT(2) = maxniter Maximum number of iterations in gradient ascent
% OPT(3) = tol      Tolerance for |lambda(n+1)-lambda(n)|^2
% LAMBDA0    1 x Q  Initial guess for the lambdas (optional)
%
% Note that this implementation is only for demonstration purposes - its not very
% fast (lots of loops that could be rolled out to matrix operations).
%
% Implementation based directly on:
% Meila M & Shi J: Learning Segmentation Using Random Walks. NIPS 13 (2001).
%
% Markus Herrgard
% $ Revision: 0.5 $ $ Date: 10/11/01 15:31:00 $

% Number of ques
Q = length(F);

% Options
nu = opt(1);
maxniter = opt(2);
tol = opt(3);

% Set initial guess
if nargin < 4
    lambda0 = -rand(Q,1);
end
lambda = zeros(Q,maxniter);
lambda(:,1) = lambda0;

% Calculate target probabilities
Ps = calc_Ps(labels);

% Gradient ascent - iterate until converged
conv = 0;
n = 0;
while (conv == 0)
    n = n+1;
    % Calculate the current transition probabilities
    P = calc_P(F,lambda(:,n));
    % Update the lambdas
    for q=1:Q
        lambda(q,n+1) = lambda(q,n) + nu * gradient(Ps,P,F{q}.data);
    end
    % Calculate the cross-entropy
    J = calc_J(Ps,P);
    % Calculate the change in lambda
    dlambda = sqrt(sum((lambda(:,n+1)-lambda(:,n)).^2));
    fprintf('%3d %6.3f %6.3f\n',n,dlambda,J);
    % Check for convergence
    if (error <= tol | n == maxniter-1)
        conv = 1;
    end
end

% Take the final lambda and return it
lambdaf = lambda(:,n+1);

function J = calc_J(Ps,P)
%CALC_J Calculate the cross entropy between Ps and P

I = size(Ps,1);

J = 0;

for i = 1:I
    for j = 1:I
        J = J + Ps(i,j)*log(P(i,j));
    end
end

J = J/I;

function Ps = calc_Ps(labels)
%CALC_PS Calculate the target transition probabilities

I = length(labels);

Ps = zeros(I);

for i = 1:I
    for j = 1:I
        if (labels(i) == labels(j))
            Ps(i,j) = 1/sum(labels == labels(i));
        else
            Ps(i,j) = 0;
        end
    end
end

function P = calc_P(F,lambdavec)
%CALC_P Calculate new transition matrix P

Q = length(lambdavec);
S = zeros(size(F{1}.data));

for q = 1:Q
    S = S + lambdavec(q)*F{q}.data; 
end

S = exp(S);
D = diag(sum(S,2));
P = inv(D)*S;

function grad = gradient(Ps,P,fq)
%GRADIENT Calculate the gradient of J

I = size(Ps,1);

grad = 0;
for i=1:I
    for j=1:I
        grad = grad + (Ps(i,j)-P(i,j))*fq(i,j);
    end
end

grad = grad/I;