function [B,F,M] = gridsyncbp(K,U,A,M,C)
%    GRIDSYNCBP Loopy BP for a grid using a synchronous schedule
%
% [B,F,M] = GRIDSYNCBP(K,U) where K is the nodenum grid and U are the
%   model energies.
%
% [B,F,M] = GRIDSYNCBP(K,U,A) where A==0 (default) indicates use of
% the sum-product algorithm for marginals. A==1 indicates use of
% max-product to find a maximizing configuration, which will be
% returned in B.
%
% [B,F,M] = GRIDSYNCBP(K,U,A,M) where M are the set of messages to
% begin with.
%  
% [B,F,M] = GRIDSYNCBP(K,U,A,M,C) where C==1 will return beliefs
% in logspace.

    
% Jerod Weinman
% jerod@acm.org
  
%	$Id: gridsyncbp.m 1 2009-05-07 13:41:00Z weinman $
 
% Jerod Weinman
% jerod@acm.org
% Copyright 2004, 2008, 2009
 
%    This file is part of Matlab GridCRF.
%
%    Matlab GridCRF 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.
%
%    Matlab GridCRF 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 Matlab GridCRF.  If not, see <http://www.gnu.org/licenses/>.

  % TODO: Convex message updates currently happen in logspace, ie
  % a*log(Mnew)+(1-a)*log(Mold) rather than properly in expspace.
  % Still, this mode damping manages to stop the oscillations that
  % often result from a synchronous schedule.
    
  % Damping coefficient for message updates
  alpha = 0.75;
  
  N = size(U.node,2);
  
  if nargin<3
	A = 0;
  end;
  
  if nargin<5
	C = 0;
  end;
  
  if nargin<6
	T= {};
  end;
  
  L = size(U.node,1);

  [E0,E1,nH] = gedges(K);
  
  
  nEdg = length(E0);

  % Product of all messages into node n
  msgProd = zeros(N,L);

  if nargin<5 || isempty(M)
	% Initialize Messages
    M = zeros(nEdg,L,2); 

  else
	
    for n=1:N
      msgProd(n,:) = sum(M(find(E0==n),:,2),1) + sum(M(find(E1==n),:,1),1);
    end;
	
  end;

  MAX_ITER = min(5*N,5000);
  MSG_TOL = 1e-3;

  % E0->E1: 1
  % E1->E0: 2
  Mnew = zeros(size(E0,1),L,2); 
  Eh = 1:nH;
  Ev = (nH+1):length(E0);

  Un = U.node';
  Uh0 = reshape(U.hEdge,[size(U.hEdge,1) L L]); % (E0,E1)
  Uv0 = reshape(U.vEdge,[size(U.vEdge,1) L L]);
  
  Uh1 = permute(Uh0,[1 3 2]); % (E1,E0)
  Uv1 = permute(Uv0,[1 3 2]);
  
  for iter=1:MAX_ITER

    %tic
    
	converging = 1;
	l2diff = 0;

    % E0 -> E1
    
    % Message Product * Local Term
    mph = Un(E0(Eh),:) + msgProd(E0(Eh),:) - M(Eh,:,2);
    mpv = Un(E0(Ev),:) + msgProd(E0(Ev),:) - M(Ev,:,2);

    % Message "Summands"
    msh = Uh0 + mph(:,:,ones(1,L));
    msv = Uv0 + mpv(:,:,ones(1,L));

    % Calculate max
    mxh = max(msh,[],2);
    mxv = max(msv,[],2);

    if A==0 % SUM-PRODUCT      
      Mh = mxh + log(sum( exp(msh - mxh(:,ones(1,L),:)), 2));
      Mv = mxv + log(sum( exp(msv - mxv(:,ones(1,L),:)), 2));
    else    % MAX-PRODUCT
      Mh = mxh;
      Mv = mxv;
    end;
    
    Mnew(Eh,:,1) = Mh;
    Mnew(Ev,:,1) = Mv;
    
    %%
    %% E1 -> E0
    %%

    % Message Product * Local Term
    mph = Un(E1(Eh),:) + msgProd(E1(Eh),:) - M(Eh,:,1);
    mpv = Un(E1(Ev),:) + msgProd(E1(Ev),:) - M(Ev,:,1);

    % Message "Summands"
    msh = Uh1 + mph(:,:,ones(1,L));
    msv = Uv1 + mpv(:,:,ones(1,L));

    % Calculate max
    mxh = max(msh,[],2);
    mxv = max(msv,[],2);
    
    if A==0 % SUM-PRODUCT      
      Mh = mxh + log(sum( exp(msh - mxh(:,ones(1,L),:)), 2));
      Mv = mxv + log(sum( exp(msv - mxv(:,ones(1,L),:)), 2));
    else    % MAX-PRODUCT
      Mh = mxh;
      Mv = mxv;
    end;
    
    Mnew(Eh,:,2) = Mh;
    Mnew(Ev,:,2) = Mv;
    
    % Normalize messages
    mx = max(Mnew,[],2);    
    Z = mx + log(sum(exp(Mnew-mx(:,ones(1,L),:)),2));
    Mnew = Mnew - Z(:,ones(1,L),:);

    % Convergence checks/progress
    l2diff = l2diff + sumv( (M-Mnew).^2 );
    
    if converging && anyv(abs(M - Mnew)>MSG_TOL)
      converging = 0;
    end;

    % Assign new messages
    M = M*(1-alpha) + alpha*Mnew;
    

    msgProd = zeros(N,L);

    % Update message products
    
    if 1
      msgProd = calcMsgProd(E0,E1,M,N);      
    elseif 0
    % Partially vectorized:
    for c=1:L
      
      % We have to chunk the edges
      di = 2*1e4;
      for d0=1:di:nEdg
        d1 = min(d0+di-1,nEdg);
        dd = d1-d0+1;

        ndx = sub2ind([N dd],E0(d0:d1),(1:dd)');
        M0 = sparse(N,dd);
        M0(ndx) = M(d0:d1,c,2);
        
        msgProd(:,c) = msgProd(:,c) + full(sum(M0,2));
      
        ndx = sub2ind([N dd],E1(d0:d1),(1:dd)');
        M0 = sparse(N,dd);
        M0(ndx) = M(d0:d1,c,1);
        
        msgProd(:,c) = msgProd(:,c) + full(sum(M0,2));
      end;
    end;
    else
    % Non-vectorized:
    for n=1:N
      msgProd(n,:) = sum(M(find(E0==n),:,2),1) + sum(M(find(E1==n),:,1),1);
    end;
	end;
	%fprintf('%d\t%f\n',iter,l2diff);
	
	if converging
	  disp(['gridsyncbp converged in ' num2str(iter) ' iterations']);
	  F = 1;
	  break;	  
	end;

    %toc
  end;
  
  %keyboard
  
  if iter==MAX_ITER
	disp(sprintf('gridsyncbp(%d) failed to converge',A));
	F = -1;
  end;	


  if A==0
    M = permute(M,[2 1 3]);
    B = gridbeliefs(K,U,M,A,C,msgProd');
  else
    B = argmax(msgProd,2);
  end;
