function [bnet, ll] = propagate_pearl(bnet)
% PROPOGATE_PEARL Apply Pearl's polytree algorithm for belief updating.
% [bnet, ll] = propogate_pearl(bnet)


% NODE TYPES
% Currently, this implementation assumes all the nodes are discrete.
% Noisy-or and multiplexer nodes are treated in an efficient manner.
% For Gaussian nodes, the equations are in Pearl's book; the vector versions can be found in
% "Inference Using Message Propogation and Topology Transformation in Vector {G}aussian
% Continuous Networks", S. Alag and A. Agogino, UAI 96.
% These Gaussian equations haven't been implemented yet.

% CENTRALIZED vs DISTRIBUTED
% If the graph has no undirected cycles (loops),
% we use the centralized message passing scheme discussed in
% "Fusion and propogation with multiple observations in belief networks",
%   Peot and Shachter, AI 48 (1991) p. 299-318.
% This is guaranteed to touch each node exactly twice.
% Otherwise we use the parallel update method discussed in
% "Probabilistic Reasoning in Intelligent Systems", Judea Pearl, 1988, 2nd ed.
% The centralized algorithm fails if the graph is not a single connected component

% LOOPS
% If the graph has no loops, W parallel updates are sufficient to converge, where W is the width
% of the graph. If the graph is loopy, we iterate until we exceed bnet.pearl.max_iter iterations,
% or the change in the beliefs (posteriors) at each node is less than bnet.pearl.threshold.
% Loopy belief prapogation was first suggested as an approximation in Pearl p235 exercise 4.7.
% For an analysis of the single loop case, see
% "Correctness of local probability propagation in graphical models with loops",
%   Yair Weiss, Neural Computation, 1999.
% For an empirical analysis of the general case, see
% "Loopy belief propagation for approximate inference: an empirical study",
%    Murphy, Weiss and Jordan, UAI 99.

% SCALING
% The only difference between the defns in Pearl and Peot and Shachter
% is that Pearl defines pi_X(x) = Pr(X=x | ev above X)
% and Peot defines pi_X(x) = Pr(X=x, ev above X).
% This difference only affects the computation of pi msgs.
% The advantages of normalizing are 
% (1) prevents underflow (esp. important when iterating many times)
% (2) it ensures the pi msgs sum to 1, which is needed for multiplexer nodes.
%
% In addition to scaling the pi's, we must scale the lambda's to prevent underflow.
% Hence we define lambda_X(x) = c Pr(ev below X | X=x) 
% where c is the normalizing constant.
% Note that the lambda's are not conditional probabilities, and do not automatically sum to 1.
% Again, this just affects the computation of the lambda msgs.
% Multiplying them by an arbitrary constant is valid: see Pearl p183 eq 4.52.
%
% The disadvantage of scaling is that we lose information, making it hard (impossible?)
% to compute the likelihood. Also, scaling makes cutset conditioning harder.
% The problem is more complicated than in the HMM case, where c_t = normalizing constant for
% alpha_t (i.e., pi_t), scaled_beta_t = c_t beta_t (i.e. pi_t), and likelihood = prod_t c_t.
% One difficulty is that we might compute lambda for a node before we compute pi,
% so we can't scale by c_t.

% VARIABLE NAMES
% msg{n}.pi(i) = Pr(Xn = i, ev above Xn)
% msg{n}.lambda(i) = Pr(ev below and including Xn | Xn = i)
% marginal{n}.T(i) = Pr(Xn = i | all evidence)
% msg{n}.pi_from_parent{p}(:) is the pi msg from n's p'th parent (size ns(p))
%  (We must use cell arrays because the pi_from_parent messages can be of different sizes.)
% msg{n}.lambda_from_child{c}(:) is the lambda msg from n's c'th child (size ns(n))
% See initialize_pearl for an explanation of child_index and parent_index.
%


scaled = bnet.pearl.scaled;
verbose = 0;
%centralized = ~bnet.pearl.loopy;
centralized = 0;
bnet = initialize(bnet, centralized);
bnet.pearl.maximize = bnet.maximize;

% We write the marginals to a file after each iteration so we can judge convergence
% while the process is running
if ~isempty(bnet.name)
  bnet.fid = fopen(bnet.name, 'w');
else
  bnet.fid = 0;
end

if centralized
  bnet = centralized_msg_passing(bnet, verbose);
else
  bnet = parallel_msg_passing(bnet, verbose, bnet.pearl.max_iter);
end
if ~isempty(bnet.name)
  fclose(bnet.fid);
end

if scaled
  ll = inf;
else
  n = 1; % arbitrary
  lik = bnet.pearl.msg{n}.pi(:)' * bnet.pearl.msg{n}.lambda(:);
  ll = log(lik+eps);
end
N = length(bnet.adj_mat);

for n=1:N
  bnet.pearl.marginal{n}.T = compute_bel_bnode(bnet.bnodes{n}, bnet);
end


%%%%%%%%

function bnet = initialize(bnet, central)

% In the centralized case, we check that pi msgs have been recvd from all parents except i
% before sending a lambda msg to i. Similarly, we check that lambda
% msgs have been recvd from all children except i before sending a pi msg to i.
% (This is not necessary, since the algorithm is designed such that we never
% read a message before it is written, but is a good sanity check.)
% In the non-centralized case, we pretend we have always received valid incoming msgs.

% We hard-code the assumption that all the nodes are discrete

N = length(bnet.adj_mat);
ns = bnet.node_sizes;
bnet.pearl.msg = cell(1,N);
bnet.pearl.marginal = cell(1,N);

if bnet.pearl.rand_init
  rand('seed', bnet.pearl.seed);
end

for n=1:N
  ps = parents(bnet.adj_mat, n);
  if central
    bnet.pearl.msg{n}.recvd_pi_from_parent = zeros(1, length(ps));
  else
    bnet.pearl.msg{n}.recvd_pi_from_parent = ones(1, length(ps));
  end
  bnet.pearl.msg{n}.pi_from_parent = cell(1, length(ps));
  for i=1:length(ps)
    p = ps(i);
    if bnet.pearl.rand_init
      bnet.pearl.msg{n}.pi_from_parent{i} = rand(ns(p), 1);
    else
      bnet.pearl.msg{n}.pi_from_parent{i} = ones(ns(p), 1);
    end
  end
  
  cs = children(bnet.adj_mat, n);
  bnet.pearl.msg{n}.lambda_from_child = cell(1, length(cs));
  if central
    bnet.pearl.msg{n}.recvd_lambda_from_child = zeros(1, length(cs));
  else
    bnet.pearl.msg{n}.recvd_lambda_from_child = ones(1, length(cs));
  end
  for i=1:length(cs)
    c = cs(i);
    if bnet.pearl.rand_init
      bnet.pearl.msg{n}.lambda_from_child{i} = rand(ns(n), 1);
    else
      bnet.pearl.msg{n}.lambda_from_child{i} = ones(ns(n), 1);
    end
  end

  if bnet.pearl.rand_init
    bnet.pearl.msg{n}.lambda = rand(ns(n), 1);
    bnet.pearl.msg{n}.pi = rand(ns(n), 1);
  else
    bnet.pearl.msg{n}.lambda = ones(ns(n), 1);
    bnet.pearl.msg{n}.pi = ones(ns(n), 1);
  end
  
  % Initialize the lambdas with any evidence
  if ismember(n, bnet.observed_nodes)
    v = bnet.values(n);
    bnet.pearl.msg{n}.lambda_from_self = zeros(ns(n), 1);
    bnet.pearl.msg{n}.lambda_from_self(v) = 1;
  else
    bnet.pearl.msg{n}.lambda_from_self = ones(ns(n), 1);
  end

  % We must wrap the bel inside a table for marginal_nodes.
  bnet.pearl.marginal{n} = mk_table(n, ns(n));
end



%%%%%%%%%

function bnet = parallel_msg_passing(bnet, verbose, max_iter)

% We repeat the following for T steps (or until convergence):
% For each node in parallel (i.e., simultaneously)
%  - update your state (lambda/pi values) 
%  - send messages to your neighbors
% This is guaranteed to converge to the right answer if
% (1) the graph has no undirected cycles, and
% (2) T >= width of graph

N = length(bnet.adj_mat);
ns = bnet.node_sizes;
% W is the size of the wrap-around buffer
if bnet.pearl.momentum == 0
  W = 2; % check bel(t) == bel(t-1)
else
  W = 3; % check bel(t) == bel(t-1) == bel(t-2)
end
history = cell(1,N);
for n=1:N
  history{n} = zeros(W, ns(n));
end

converged = 0;
iter = 1;
while ~converged & (iter <= max_iter)
  if verbose, fprintf('parallel iteration %d\n', iter); end
  for n=1:N
    bnet.pearl.msg{n}.lambda = compute_lambda_bnode(bnet.bnodes{n}, bnet);
    bnet.pearl.msg{n}.pi = compute_pi_bnode(bnet.bnodes{n}, bnet);
  end
 
  % Convergence criterion is that the last W bel's are approximately the same.
  % (We assume all nodes are discrete.)
  if iter >= W
    changed = 0;
    for n=1:N
      if ismember(n, bnet.hidden_nodes)
	bnet.pearl.marginal{n}.T = compute_bel_bnode(bnet.bnodes{n}, bnet);
	new_bel =  bnet.pearl.marginal{n}.T;
	fprintf(bnet.fid, '%9.7f ', new_bel(2));
	history{n}(wrap(iter, W), :) = new_bel;
	for i=2:W
	  if ~approxeq(history{n}(1,:), history{n}(i,:), 1e-5)
	    nonconv = n;
	    changed = 1;
	  end
	end
      end
    end
    fprintf(bnet.fid, '\n\n');
    if ~changed, converged = 1; end
  end
  
  if ~converged
    %if verbose, fprintf('updating neighbors\n'); end
    for n=1:N
      bnet = send_msgs_to_all_neighbors(bnet, n, 0);
    end
    iter = iter + 1;
  end
end

bnet.pearl.niterations = iter;


%%%%%%%%%

function bnet = centralized_msg_passing(bnet, verbose)

% We first send messages up to the root (pivot node), and then back towards the leaves.
% If the bnet is a singly connected graph (no loops), choosing a root induces a directed tree.
% Peot and Shachter discuss ways to pick the root so as to minimize the work,
% taking into account which nodes have changed.
% For simplicity, we always pick the root to be the last node in the graph.
% This means the first pass is equivalent to going forward in time in a DBN.

N = length(bnet.adj_mat);

% Send messages from leaves to root
if verbose, fprintf('collect to root\n'); end
for i=1:N-1
  n = bnet.pearl.postorder(i);
  above = parents(bnet.pearl.adj_mat, n);
  bnet = send_msgs_to_some_neighbors(bnet, n, verbose, above);
end

% Process root
n = bnet.pearl.root;
bnet.pearl.msg{n}.lambda = compute_lambda_bnode(bnet.bnodes{n}, bnet);
bnet.pearl.msg{n}.pi = compute_pi_bnode(bnet.bnodes{n}, bnet);
%[bnet, lik] = compute_bel(bnet, n);
%ll = log(lik);

% Send messages from root to leaves
if verbose, fprintf('distribute from root\n'); end
for i=1:N
  n = bnet.pearl.preorder(i);
  below = children(bnet.pearl.adj_mat, n);
  bnet = send_msgs_to_some_neighbors(bnet, n, verbose, below);
  % If a node has no parents, it will not send any lambda msgs, and hence not compute its lambda.
  % If a node has no children, it will not send any pi msgs, and hence not compute its pi.
  % So we compute them both here, possibly redundantly, for good measure.
  bnet.pearl.msg{n}.lambda = compute_lambda_bnode(bnet.bnodes{n}, bnet);
  bnet.pearl.msg{n}.pi = compute_pi_bnode(bnet.bnodes{n}, bnet);
  %bnet = compute_bel(bnet, n);
end


  
%%%%%%%%%%

function bnet = send_msgs_to_some_neighbors(bnet, n, verbose, valid)

ps = intersect(parents(bnet.adj_mat, n), valid);
if ~isempty(ps)
  bnet.pearl.msg{n}.lambda = compute_lambda_bnode(bnet.bnodes{n}, bnet);
  bnet = send_lambda_msgs(bnet.bnodes{n}, bnet, ps);
end
cs = intersect(children(bnet.adj_mat, n), valid);
if ~isempty(cs)
  bnet.pearl.msg{n}.pi = compute_pi_bnode(bnet.bnodes{n}, bnet);
  bnet = send_pi_msgs(bnet.bnodes{n}, bnet, cs);
end

  
%%%%%%%%%%

function bnet = send_msgs_to_all_neighbors(bnet, n, verbose)

ps = parents(bnet.adj_mat, n);
if ~isempty(ps)
  bnet = send_lambda_msgs(bnet.bnodes{n}, bnet, ps);
end
cs = children(bnet.adj_mat, n);
if ~isempty(cs)
  bnet = send_pi_msgs(bnet.bnodes{n}, bnet, cs);
end

%%%%%%%%%


function bnet = send_pi_msgs(bnode, bnet, S)

verbose = 0;
global scaled
n = bnode.self;
for i=1:length(S)
  c = S(i);
  pi = compute_pi_msg_bnode(bnode, bnet, c);
  if scaled, pi = normalise(pi); end
  if verbose
    fprintf('sending pi  msg from %d to %d\n', n, c);
    pi(:)'
  end
  j = bnet.pearl.parent_index{c}(n); % I am c's j'th parent
  if 0
    bnet.pearl.msg{c}.pi_from_parent{j} = pi;
  else
    old_pi = bnet.pearl.msg{c}.pi_from_parent{j};
    alpha = bnet.pearl.momentum;
    bnet.pearl.msg{c}.pi_from_parent{j} = alpha*old_pi + (1-alpha)*pi;
  end
  bnet.pearl.msg{c}.recvd_pi_from_parent(j) = 1;
end


%%%%%%%%%%

function bnet = send_lambda_msgs(bnode, bnet, S)

n = bnode.self;
global scaled
verbose = 0;
ns = bnet.node_sizes;
for i=1:length(S)
  if verbose, fprintf('sending lam msg from %d to %d\n', n, p); end
  p = S(i);
  lambda = compute_lambda_msg_bnode(bnode, bnet, p);
  if scaled, lambda = normalise(lambda); end
  j = bnet.pearl.child_index{p}(n); % I am p's j'th child
  if 0
    bnet.pearl.msg{p}.lambda_from_child{j} = lambda(:);
  else
    old_lam =  bnet.pearl.msg{p}.lambda_from_child{j}(:);
    alpha = bnet.pearl.momentum;
    bnet.pearl.msg{p}.lambda_from_child{j} = alpha*old_lam + (1-alpha)*lambda(:);
  end
  bnet.pearl.msg{p}.recvd_lambda_from_child(j) = 1;
end




