function varargout = spasyncbp(G,D,R,S,alpha,order)
% LOOPYBP Asynchronous loopy belief propagation on a factor graph
%
% [B,M,flag,Mf,N,Nf] = SPASYNCBP(G,D) where B are the factor marginals, M
% are the node marginals, flag indicates whether BP converged, N
% are the resulting state space sizes after pruning.
%
% [P,flag,S] = SPASYNCBP(G,D,1) indicates to run MAX-PRODUCT, yielding
% predictions P and flag indicating convergence.
%
% SPASYNCBP(G,D,R,S) where R indicates SUM or MAX PRODUCT, and S are
% the initial variable to function messages to use.
%
% SPASYNCBP(...,alpha) uses 0 <= alpha < 1 as a damping factor on
% message updates. alpha=0 corresponds to no damping. Default is 0.5.
%
% SPASYNCBP(...,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.
%
% If the length of D is greater than one, the excess dimensions are
% interpreted as limits to the state space sparsity of nodes in the
% factor graph. If the entries are positive integers, they are
% interpeted as absolute minimum cardinalities. If the entries are
% reals between 0-1, they are interpreted as fractions of the
% original node cardinality. Absolute minimum cardinalities are
% calculated by multiplying the original cardinality by the given
% fraction and taking the ceiling.
%
% IMPORTANT NOTA BENE:
%   For very loopy BP and small state spaces, this  is likely to
%   take longer than its synchronous counterpart SPLOOPYBP. This is
%   because sparsity of impacted nodes is updated after each factor
%   is processed; this means many more sorts are likely to be
%   required than would be optimal. Unless a specified, fixed
%   ordering is used, the performance of SPLOOPYBP is likely to be
%   better. 
%
% See also SPLOOPYBP, LOOPYBP, ASYNCBP
  

% Jerod Weinman
% jerod@acm.org
% (c) 2008-2011
  
% $Id: spasyncbp.m 35 2011-07-29 15:56:28Z 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<3 || isempty(R)
    R = 0; % sum-product
  end;
  
  E = adjmat(G);
  L = card(G);

  N = L;
  
  if length(D)>1
    I = D(2:end);
    
    D = D(1);
    
    if length(I)==1
      I = I(ones(1,length(L)));
    end;
    
    % Multiply fractional minimum sparsities by node cardinality
    I(I<1) = ceil(L(I<1).*I(I<1));
  else
    I = zeros(1,length(L));
  end;

  [V,F] = find(E);

  
  if nargin<4 || isempty(S)
    % INITIALIZE MESSAGES: ln 1 = 0
    
    Mf = cell(size(E)); % Variable to function
    Mv = cell(size(E)); % Function to variable
    Mvi = cell(size(E)); % Function to variable infinite locations 
    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);
      Mvi{V(e),F(e)} = false(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; % Copy messages
    clear S;
    
    % Check messages
    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;
      
      % Initialize function to variable messages	  
      Mv{V(e),F(e)} = zeros(L(V(e)),1);
    end;
  end;

  if nargin<5
    alpha = 0.5;
  end
    
  if nargin<6 || 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), tracking for infinite products, and  belief sparsity
  msgProd = cell(G.numNodes,1);
  belPrune = cell(G.numNodes,1);
  
  for v=1:G.numNodes
    belPrune{v} = false(L(v),1);
    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;

  if ~isempty(order)
    pFac = fliplr(order);
  end
  
  for iter = 1:MAX_ITER
    
    % Consistency check on messages
    for e=1:length(V)
      if 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;
    end;
    
    converging = 1;
    l2diff = 0;
    
    if isempty(order)
      pFac = randperm(length(G.factors));
    else
      pFac = fliplr(pFac); % Reverse prior ordering
    end
    
    for f=pFac(:)'

      %% Calculate function to variable messages
      MvNew = spmessages(G.factors{f},Mf(:,f),R,T{f});
      
      A = arguments(G.factors{f});
      
      % For each argument (node) of the factor
      for a=1:length(A)

        v = A(a);
        
        if any(isnan(MvNew{a}))
          error('NaN in new factor->variable message');
        end;
        
        mvf = ~Mvi{v,f};          % Finite indices of old message
        mvnf = ~isinf(MvNew{a});  % Finite indices of new message
        mvni = ~mvnf;             % Infinite indices of new message
        
        mf = mvf & mvnf;          % Finite indices of both messages
        
        l2diff = l2diff + sum((Mv{v,f}(mf) - MvNew{a}(mf)).^2);

        if converging && ...
              (any(abs(Mv{v,f}(mf) - MvNew{a}(mf))>MSG_TOL) || ...
               any(xor(mvf,mvnf)))
          converging = 0;
        end;
        
        % Handle infinite messages appropriately
        if any(Mvi{v,f}) || any(mvni)
        
          % Remove finite parts of old message from product
          msgProd{v}(mvf) = msgProd{v}(mvf) - Mv{v,f}(mvf);
          
          % Update function to variable message
          if alpha==0 % Handle undamped updates specially, since 0*Inf=NaN
            Mv{v,f} = MvNew{a};
          else
            Mv{v,f} = alpha*Mv{v,f} + (1-alpha)*MvNew{a};
          end
        
          % Update infinite indices
          Mvi{v,f} = isinf(Mv{v,f});
          mvf = ~Mvi{v,f};
          
          % Add new message back to product (where finite)
          msgProd{v}(mvf) = msgProd{v}(mvf) + Mv{v,f}(mvf);
          
        else % Update message and message product normally
          
          % Remove  message from product
          msgProd{v} = msgProd{v} - Mv{v,f};
          
          % Update function to variable message
          Mv{v,f} = alpha*Mv{v,f} + (1-alpha)*MvNew{a};
          
          % Add new message back to product
          msgProd{v} = msgProd{v} + Mv{v,f};
        end % Message (factor to variable) and product updates
          
        if any(isnan(msgProd{v}))
          error('NaN in new message product');
        end

        if any (isnan(Mv{v,f}))
          error('NaN in updated variable->factor message');
        end

        if any(isinf(msgProd{v}))
          error('Inf in new message product');
        end
        
      end; % Argument loop (given factor)

      
      % Update variable to function messages
      % for any factors neighboring the arguments

      for a=1:length(A)
        
        v = A(a);
        
        % Update sparsity
        belPrune{v} = sparsify(msgProd{v},D,I(v));
        
        N(v) = sum(belPrune{v});
        
        % Product of messages from any other factor into this variable
        mf = msgProd{v};
        
        % Remove message from factor to variable
        if any(Mvi{v,f}) % Handle infinite messages appropriately
          
          ix = Mvi{v,f}; % infinite message locations
          
          % subtract (div in logspace) any finite part
          mf(~ix) = mf(~ix) - Mv{v,f}(~ix);
          
          % Is there another message in the product that would still make the
          % product infinite?
          ix2 = ix & any(cell2mat(Mvi(v,mysetdiff(find(E(v,:)),f))'),2);
          
          if any(ix2) % If so, restore infinite part
            mf(ix2) = -Inf;
          end;
          
        else % Perform standard (non-Inf safe) subtraction (div in logspace)
          mf = mf - Mv{v,f};
        end;

        if any(size(mf) ~= [L(v) 1]) % Quick assertion
          error('Updated message V(%d)->F(%d) has incorrect size',v,f);
        end;

        % Conform to sparsity
        mf(belPrune{v}) = -Inf;
      
        % Normalize
        mx = max(mf);
        mf = mf - mx - log(sum(exp(mf-mx)));
      
        if any(isnan(mf))
          error('NaN in variable (%d) to factor (%d) message',V(e),F(e));
        end;
      
        % Assign final message to factor
        Mf{v,f} = mf;
      
      end; % Argument loop (given factor)
          
    end; % Factor loop

    
    if converging
      flag = 1;
      break;
    end;
    
  end; % Main loop
  
  %fprintf('iter = %d\n',iter);
  
  if ~flag
    warning(sprintf('spasyncbp(%d) did not converge',R));
  end;
  
  % Calculate beliefs
  if R
    for v=1:G.numNodes
      B(v) = argmax( msgProd{v} );
    end;
    
    switch nargout
      case 1
        varargout = {B};
      case 2
        varargout = {B,flag};
      otherwise
        varargout = {B,flag,Mf};
    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;
    
    if nargout>4
      varargout{5} = N;
    end;

  end;
  
function pruned = sparsify(mm,D,I)
  
  % Normalize message product (belief)
  mx = max(mm);
  bel = mm - mx - log(sum(exp(mm-mx)));
  
  % Sparsify
  [pm,px] = sort(bel,'descend'); % order probs
  
  ZI = -log(cumsum(exp(pm))); % sorted divergences
  
  prune = min(find(ZI<D)); % cut-off index 
  
  prune = max(prune,I); % respect variable maximum pruning
        
  pruned = false(length(mm),1);
  pruned(px(prune+1:end)) = 1; % prune
        
