function varargout = sploopybp(G,D,R,S,alpha)
% LOOPYBP Synchronous loopy belief propagation on a factor graph
%
% [B,M,flag,Mf,N] = SPLOOPYBP(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] = SPLOOPYBP(G,D,1) indicates to run MAX-PRODUCT, yielding
% predictions P and flag indicating convergence.
%
% SPLOOPYBP(G,D,R,S) where R indicates SUM or MAX PRODUCT, and S are
% the initial variable to function messages to use.
%
% SPLOOPYBP(...,alpha) uses 0 <= alpha < 1 as a damping factor on
% message updates. alpha=0 corresponds to no damping. Default is 0.5.
%
% 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.

% Jerod Weinman
% jerod@acm.org
  
% $Id: sploopybp.m 34 2011-07-28 17:46:52Z 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
    
    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; % 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
    
  % Initialize message products (over function to variable messages) and
  % belief sparsity
  msgProd0 = cell(G.numNodes,1);
  belPrune = cell(G.numNodes,1);
  
  for v=1:G.numNodes
    belPrune{v} = false(L(v),1);
    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;
  
  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;
    
    %% Calculate function to variable messages

    l2diff = 0;
    
    for f=1:length(G.factors)
      
      MvNew = spmessages(G.factors{f},Mf(:,f),R,T{f});
      
      A = arguments(G.factors{f});
      
      for a=1:length(A)
        
        if any(isnan(MvNew{a}))
          error('NaN in new factor->variable message');
        end;
        
        mvi = ~isinf(Mv{A(a),f});
        mvni = ~isinf(MvNew{a});
        l2diff = l2diff + sum((Mv{A(a),f}(mvi&mvni) - MvNew{a}(mvi&mvni)).^2);
        
        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);
    
    msgProd = msgProd0;
    
    for e=1:length(V)
      msgProd{V(e)} = msgProd{V(e)} + Mv{V(e),F(e)};
    end;
    
    
    %% Calculate node beliefs and sparsity
    if ~converging % & iter==1 %& 0
      for v=1:G.numNodes
        mm = msgProd{v};
        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(v)); % respect variable maximum pruning
        
        bp = false(L(v),1);
        bp(px(prune+1:end)) = 1; % prune
        
        if 0 && any(bp & ~belPrune{v})
          fprintf('Pruned V(%d) by %d to %d/%d\n', v, ...
                  sum(bp & ~belPrune{v}),...
                  sum(~(bp | belPrune{v})),...
                  L(v));
        end;
        
        N(v) = sum(~(bp | belPrune{v}));

        % The following line incrementally prunes, disallowing
        % state resuscitation.
        %belPrune{v} = belPrune{v} | bp;
        % Use the following line to dynamically prune based on
        % current information:
        belPrune{v} = bp;
        
      end; % Node loop
    end; % Not converging check
    
    %% Calculate variable to function messages
		
    msgProd = msgProd0;
    infProd = msgProd0;
    
    % First Pass: Product of messages from *any* factor into a variable
    for e=1:length(V)
      mm = Mv{V(e),F(e)}; % var to factor msg
      ix = isinf(mm);    % "illegal" (infinite) entries
      
      if any(ix) % Only allow finite entries to contribute to sum
        msgProd{V(e)}(~ix) = msgProd{V(e)}(~ix) + mm(~ix);
        infProd{V(e)}(ix) = infProd{V(e)}(ix) + 1;
      else
        msgProd{V(e)} = msgProd{V(e)} + mm;
      end;
    end;

    % Second pass: Remove reverse message from total products 
    for e=1:length(V)

      % Product of messages from any other factor into this variable
      mf = msgProd{V(e)};

      % Remove message from factor to variable
      
      if any(infProd{V(e)}) % Handle infinite messages appropriately

        ix = isinf(Mv{V(e),F(e)}); % infinite msg

        % subtract (div in logspace) any finite part
        if any(ix) 
          mf(~ix) = mf(~ix) - Mv{V(e),F(e)}(~ix);
        end;

        % is there another message that would make it infinite?
        ix2 = ix & infProd{V(e)}>1;
        
        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(e),F(e)};
      end;

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

      if ~converging
        % Conform to sparsity
        mf(belPrune{V(e)}) = -Inf;
      end;
      
      % 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(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 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;
