function varargout = loopybp(G,R,S,alpha,order)
% LOOPYBP Synchronous loopy belief propagation on a factor graph
%
% [B,M,flag,S] = LOOPYBP(G) where B are the factor marginals, M
% are the node marginals, flag indicates whether BP converged. S
% are the variable to function messages at termination.
%
% [P,flag,S] = LOOPYBP(G,1) indicates to run MAX-PRODUCT, yielding
% predictions P and flag indicating convergence.
%
% LOOPYBP(G,R,S) where R indicates SUM or MAX PRODUCT, and S are
% the initial variable to function messages to use.
%
% LOOPYBP(...,alpha) uses 0 <= alpha < 1 as a damping factor on
% message updates. alpha=0 corresponds to no damping. Default is 0.5.
%
%
% See also ASYNCBP, SPLOOPYBP
  
% Jerod Weinman
% jerod@acm.org
    
% $Id: loopybp.m 33 2011-07-26 20:28:07Z weinman $
 
% Jerod Weinman
% jerod@acm.org
% Copyright 2008, 2009, 2011
 
%    This file is part of MatLibFactorGraph.
%
%    MatLibFactorGraph is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    MatLibFactorGraph is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with MatLibFactorGraph.  If not, see <http://www.gnu.org/licenses/>.
 
%    This file is part of MatLibFactorGraph.
%
%    MatLibFactorGraph is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    MatLibFactorGraph is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with MatLibFactorGraph.  If not, see <http://www.gnu.org/licenses/>.
  
  if nargin<2 || isempty(R)
    R = 0; % sum-product
  end;
  
  E = adjmat(G);
  L = card(G);

  [V,F] = find(E);

  if nargin<3 || isempty(S)
    % INITIALIZE MESSAGES: ln 1 = 0
	
    Mf = cell(size(E)); % Variable to function
    Mv = cell(size(E)); % Function to variable
    
    for e=1:length(V)
      Mf{V(e),F(e)} = zeros(L(V(e)),1);
      Mv{V(e),F(e)} = zeros(L(V(e)),1);
    end;
  else
    
    if any(size(S)<size(E))
      % Expand size of messages if any are undefined
      S{size(E,1),size(E,2)} = [];
    end;
	  
    Mf = S;
    clear S;
    
    for e=1:length(V)
      if isempty(Mf{V(e),F(e)})
        Mf{V(e),F(e)} = zeros(L(V(e)),1);
      elseif any(size(Mf{V(e),F(e)}) ~= [L(V(e)) 1])
        error('Message V(%d)->F(%d) has incorrect size',V(e),F(e));
      end;
      
      Mv{V(e),F(e)} = zeros(L(V(e)),1);
    end;
  end;
  
  if nargin<4
    alpha = 0.5;
  end
	
  % Initialize message products (over function to variable messages)
  for v=1:G.numNodes
	msgProd0{v} = zeros(L(v),1);
  end;

  T = cell(1,length(G.factors));
  
  % Precompute message expansion indices
  for f=1:length(G.factors)
	if nargs(G.factors{f})>1
	  T{f} = combindices(G.factors{f});
	end;
  end;
  
  MAX_ITER = 5*G.numNodes;
  MSG_TOL = 1e-3;
  
  nEdg = length(V);
  
  flag = 0;
  
  % Normalize all messages
  for e=1:length(V)
    mm = Mf{V(e),F(e)};
    mx = max(mm);
    Mf{V(e),F(e)} = mm-mx-log(sum(exp(mm-mx)));
    
    mm = Mv{V(e),F(e)};
    mx = max(mm);
    Mv{V(e),F(e)} = mm-mx-log(sum(exp(mm-mx)));
  end;
 
  for iter = 1:MAX_ITER
    
    converging = 1;
    
    % Calculate function to variable messages
    
    l2diff = 0;
    
    for f=1:length(G.factors)
      
      MvNew = messages(G.factors{f},Mf(:,f),R,T{f});
      
      A = arguments(G.factors{f});
      
      for a=1:length(A)
        
        mvi = ~isinf(Mv{A(a),f});
        mvni = ~isinf(MvNew{a});
        
        l2diff = l2diff + sum((Mv{A(a),f}(mvi&mvni) - MvNew{a}(mvi&mvni)).^2);
        
        %a,A(a),f
        %Mv{A(a),f}
        %MvNew{a}
		
        if converging && ...
              (any(abs(Mv{A(a),f}(mvi&mvni) - MvNew{a}(mvi&mvni))>MSG_TOL) || ...
               any(xor(mvi,mvni)))
          
          converging = 0;
        end;
        
        % Weighted (damped) update
        Mv{A(a),f} = alpha*Mv{A(a),f} + (1-alpha)*MvNew{a};
        
      end; % Argument loop (given factor)
      
    end; % Factor loop

    %fprintf('L2 = %f\n',l2diff);
    
    % Calculate variable to function messages
    msgProd = msgProd0;
    
    % Product of messages from any factor into a variable
    for e=1:length(V)
      msgProd{V(e)} = msgProd{V(e)} + Mv{V(e),F(e)};
    end;
    

    for e=1:length(V)
      
      mf = msgProd{V(e)} - Mv{V(e),F(e)};
      
      % Normalize message
      mx = max(mf);
      mf = mf - mx - log(sum(exp(mf-mx)));
      
      if converging && any(abs(Mf{V(e),F(e)} - mf)>MSG_TOL)
        converging = 0;
      end;
      
      Mf{V(e),F(e)} = mf;
    end;
    
    if converging
      flag = 1;
      break;
    end;
    
  end;
  
  %fprintf('iter = %d\n',iter);
  
  if ~flag
    warning(sprintf('loopybp(%d) did not converge',R));
  end;
  
  % Calculate beliefs
  if R
    for v=1:G.numNodes
      B(v) = argmax( msgProd{v} );
    end;
    
    switch nargout
      case 3
        varargout = {B,flag,Mf};
      case 2
        varargout = {B,flag};
      otherwise
        varargout = {B};
    end;
  else
	  
    for f=1:length(G.factors)
      B{f} = belief(G.factors{f},Mf(:,f));
    end;
    
    varargout = {B};
    
    if nargout>1
      % Normalize log marginals
      for v=1:G.numNodes
        mm = msgProd{v};
        mx = max(mm);
        M{v} = mm - mx - log(sum(exp(mm-mx)));
      end;
      
      varargout{2} = M;
    end;
    
    if nargout>2
      varargout{3} = flag;
    end;
    
    if nargout>3
      varargout{4} = Mf;
    end;
  end;
