function [m_up_neg, m_up_pos, m_down_neg, m_down_pos, m_left_neg, m_left_pos, m_right_neg, m_right_pos] = ...
    comb_sum_product(phi_neg, phi_pos, theta)

% Perform sum-product in a comb-structured graph. More specifically,
% subgraph A, as pictured in figure 8.2(b) of the homework. The arguments
% are the unary log-potentials phi_neg and phi_pos, where phi_neg(i, j) is
% the log value of the potential if x(i, j) is -1, and phi_pos(i, j) is the
% log value if x(i, j) is 1. Theta is the coupling parameter defined in the
% problem. All unary potentials which lie outside the comb structure are
% ignored.
%
% The return values are the sum-product messages. m_left_neg(i, j) is the
% log value of the message from left neighbor x(i, j-1) to x(i, j) when
% x(i, j) takes the value -1. The other return values are defined
% analogously. All messages which do not correspond to edges in the comb
% graph are set to zero.



[M, N] = size(phi_neg);
m_up_neg = zeros(M, N);
m_up_pos = zeros(M, N);
m_down_neg = zeros(M, N);
m_down_pos = zeros(M, N);
m_left_neg = zeros(M, N);
m_left_pos = zeros(M, N);
m_right_neg = zeros(M, N);
m_right_pos = zeros(M, N);

% up the columns
for j = 1:2:N
  [m_neg, m_pos] = forward_messages_chain(phi_neg(M-1:-1:1, j)', phi_pos(M-1:-1:1, j)', theta);
  m_up_neg(M-1:-1:1, j) = m_neg;
  m_up_pos(M-1:-1:1, j) = m_pos;
end

% upwards messages become unary potentials in the chain model
phi_neg_top = phi_neg(1, :);
phi_pos_top = phi_pos(1, :);
for j = 1:2:N
  phi_neg_top(j) = phi_neg_top(j) + m_up_neg(1, j);
  phi_pos_top(j) = phi_pos_top(j) + m_up_pos(1, j);
end

% rightwards
[m_neg, m_pos] = forward_messages_chain(phi_neg_top, phi_pos_top, theta);
m_right_neg(1, :) = m_neg;
m_right_pos(1, :) = m_pos;

% leftwards
[m_neg, m_pos] = forward_messages_chain(phi_neg_top(end:-1:1), phi_pos_top(end:-1:1), theta);
m_left_neg(1, end:-1:1) = m_neg;
m_left_pos(1, end:-1:1) = m_pos;

% downwards
for j = 1:2:N
  curr_phi_neg = phi_neg(1:M-1, j);
  curr_phi_pos = phi_pos(1:M-1, j);
  
  % leftwards and rightward messages become unary potentials in the chain model
  curr_phi_neg(1) = curr_phi_neg(1) + m_right_neg(1, j);
  curr_phi_pos(1) = curr_phi_pos(1) + m_right_pos(1, j);
  curr_phi_neg(1) = curr_phi_neg(1) + m_left_neg(1, j);
  curr_phi_pos(1) = curr_phi_pos(1) + m_left_pos(1, j);
  
  [m_neg, m_pos] = forward_messages_chain(curr_phi_neg', curr_phi_pos', theta);
  m_down_neg(1:M-1, j) = m_neg;
  m_down_pos(1:M-1, j) = m_pos;
end





function [m_right_neg, m_right_pos] = forward_messages_chain(phi_neg, phi_pos, theta)

N = size(phi_pos, 2);
m_right_neg = zeros(1, N);
m_right_pos = zeros(1, N);

for i = 2:N
  m_right_neg(i) = logaddexp(m_right_neg(i-1) + phi_neg(i-1) + theta, m_right_pos(i-1) + phi_pos(i-1) - theta);
  m_right_pos(i) = logaddexp(m_right_neg(i-1) + phi_neg(i-1) - theta, m_right_pos(i-1) + phi_pos(i-1) + theta);
end





function result = logaddexp(a, b)

mx = max(a, b);
result = log(exp(a - mx) + exp(b - mx)) + mx;
