function B = gridbeliefs(K,U,M,A,C,P)
% BELIEFS Calculate node and edge beliefs
%
%   B = BELIEFS(K,U) where K is the nodenum grid and U are the
%   model energies
%
%   B = BELIEFS(K,U,M) where M are the messages from belief propagation. 
%
%   B = BELIEFS(K,U,M,A) here A==0 indicates the usual sum-product
%   marginals will be returned. A==1 indicates to use max-product
%   messages and a maximal configuration will be returned in B.
%
%   B = BELIEFS(K,U,M,A,C) with C==1 will not exponentiate
%   beliefs, but return them in logspace.
%
%
%   B = BELIEFS(K,U,M,A,C,P) where P contains the message products


% Jerod Weinman
% jerod@acm.org
    
% $Id%
 
% 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/>.
    

  
  N = size(U.node,2); % Number of nodes
  L = size(U.node,1); % Number of labels

  if nargin<3
    M = [];
  end;
  
  if nargin<4
	A = 0;
  end;

  if nargin<5
	C = 0;
  end;

  [E0,E1,nH] = gedges(K);
  
  % Edge index
  Ei = zeros(length(E0),1);
  Ei(1:nH) = edgeind(mkvector(K(:,1:end-1)),size(K,1),size(K,2),1);
  Ei(nH+1:end) = edgeind(mkvector(K(1:end-1,:)),size(K,1),size(K,2),0);

  if nargin<6
    
	% Product of all messages into each node
	msgProd = zeros(L,N);

    if nargin>3 && ~isempty(M)
      msgProd = calcMsgProd(E0,E1,M,N);
    end;
  else
	msgProd = P;
  end;
  
	
  Bn = msgProd + U.node; % LxN

  mx = max(Bn,[],1); % 1xN
  
  if A
	% MAX PRODUCT
	[ll,nn] = find(mx(ones(1,L),:)==B);
	[nns,nni] = unique(nn);
	Bn = ll(nni)';
  else
	% SUM PRODUCT

	% To normalize, find the sum in log space and subtract it.

	Z = mx + log ( sum( exp(Bn-mx(ones(1,L),:)), 1));
	
	Bn = Bn - Z(ones(1,L),:);
	
	if ~C
	  Bn = exp(Bn);
	  
	  % doubly-normalize
	  Z = sum(Bn);
	  Bn = Bn./Z(ones(1,L),:);
	end;
	
  end;
  
  if A
    B = Bn;
	return;
  end;
  
  B.node = Bn;
  B.hEdge = zeros(size(U.hEdge));
  B.vEdge = zeros(size(U.vEdge));
  
  if ~isempty(M)
    for e=1:length(E0)
      
      n0 = E0(e);
      n1 = E1(e);
	
      % Product of all messages into n0 (sans n1)
      prod0 = msgProd(:,n0) - M(:,e,2);   % Lx1
      
      % repmat( prod0,[1 L])
      prod0 = prod0(:,ones(1,L));    % LxL -- n0 changes in row
      
      % Product of all messages into n1 (sans n0)
      prod1 = msgProd(:,n1) - M(:,e,1);
      
      prod1 = prod1'; % 1xL
      
      % repmat( prod1',[L 1])
      prod1 = prod1(ones(1,L),:); % LxL -- n1 changes in col
      
      phi0 = U.node(:,n0*ones(1,L));  % LxL -- n0 changes in row
      phi1 = U.node(:,n1*ones(1,L))'; % LxL -- n1 changes in col
      
      if e<=nH
        psi = U.hEdge(Ei(e),:);
      else
        psi = U.vEdge(Ei(e),:);
      end;
      
      psi = reshape(psi,[L L]);
      
      prod = phi0 + phi1 + psi + prod0 + prod1;
      
      mx = max(max(prod));
      Z = mx + log(sumv(exp(prod-mx)));
      
      be = prod - Z;
      
      if ~C
        be = exp(be);
        
        % doubly normalize
        be = be./sumv(be);
      end
      
      if e<=nH
        B.hEdge(Ei(e),:) = be(:);
      else
        B.vEdge(Ei(e),:) = be(:);
      end;
      
    end;
  else
    for e=1:length(E0)
      
      n0 = E0(e);
      n1 = E1(e);
	
      
      phi0 = U.node(:,n0*ones(1,L));  % LxL -- n0 changes in row
      phi1 = U.node(:,n1*ones(1,L))'; % LxL -- n1 changes in col
      
      prod = phi0 + phi1;
      
      mx = max(max(prod));
      Z = mx + log(sumv(exp(prod-mx)));
      
      be = prod - Z;
      
      if ~C
        be = exp(be);
        
        % doubly normalize
        be = be./sumv(be);
      end
      
      if e<=nH
        B.hEdge(Ei(e),:) = be(:);
      else
        B.vEdge(Ei(e),:) = be(:);
      end;
      
    end;
  end;
