function varargout = asyncbp(G,R,S,alpha,order)
% ASYNCBP Asynchronous loopy belief propagation on a factor graph
%
% [B,M,flag,S] = ASYNCBP(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] = ASYNCBP(G,1) indicates to run MAX-PRODUCT, yielding
% predictions P and flag indicating convergence.
%
% ASYNCBP(G,R,S) where R indicates SUM or MAX PRODUCT, and S are
% the initial variable to function messages to use.
%
% ASYNCBP(...,alpha) uses 0 <= alpha < 1 as a damping factor on
% message updates. alpha=0 corresponds to no damping. Default is 0.2.
%
% ASYNCBP(...,order) specifies the ordering over factors used to
% pass message from factors to variables asynchronously. If order
% is 'rand', a new ordering is chosen randomly at each iteration
% (thd default), otherwise order may be a vector of integers that
% is a permutation of 1:numFactors(G). If a given ordering is used,
% it is reversed in alternate iterations.


  
  
% Jerod Weinman
% jerod@acm.org
  
% $Id: asyncbp.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/>.
  
  if nargin<2 || isempty(R)
    R = 0; % sum-product
  end;
  
  E = adjmat(G);
  L = card(G);
  
  [V,F] = find(E);
  
  if nargin<3
    % 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.2;
  end
  
  if nargin<5 || strcmp(order,'rand')
    order = [];
  elseif ~(length(order)==length(G.factors) && ... % Permutation?
           isempty(setdiff(order,1:length(G.factors)))) 
    error(['order must be ''rand'' or  a permutation of the factor ' ...
           'indices']);
  end
  
  % Initialize message products (over function to variable messages)
  for v=1:G.numNodes
    msgProd{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;
  
  if ~isempty(order)
    pFac = fliplr(order);
  end
  
  for iter = 1:MAX_ITER
    
    if isempty(order)
      pFac = randperm(length(G.factors));
    else
      pFac = fliplr(pFac); % Reverse prior ordering
    end
    
    converging = 1;
    l2diff = 0;    
    
    % Calculate function to variable messages
    for f=pFac(:)'
      
      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;

        % Remove message from product: DOES NOT HANDLE INFS
        msgProd{A(a)} = msgProd{A(a)} - Mv{A(a),f};
        
        % Update function to variable message
        Mv{A(a),f} = alpha*Mv{A(a),f} + (1-alpha)*MvNew{a};
        
        % Add new message back to product
        msgProd{A(a)} = msgProd{A(a)} + Mv{A(a),f};
        
        
      end; % Argument loop (given factor)


      % Update variable to function messages
      % for any factors neighboring the arguments
      
      for a=1:length(A)

        Fa = find(V==A(a));
        
        for fa=Fa(:)'
          
          mf = msgProd{A(a)} - Mv{A(a),F(fa)};
          % Normalize message
          mx = max(mf);
          mf = mf - mx - log(sum(exp(mf-mx)));

          Mf{A(a),F(fa)} = mf;
        end;
      end;
      
    end; % Factor loop

	  

    %fprintf('L2 = %f\n',l2diff);
    
    
    if converging
      flag = 1;
      break;
    end;
    
  end; % Main loop
  
  %fprintf('iter = %d\n',iter);
  
  if ~flag
    warning(sprintf('asyncbp(%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;
