function [nodes, loglik] = EM_step(bnet, nodes, vals, obs_nodes, soft_vals, soft_nodes, cliques)
% EM_STEP Do an Expectation and a Maximization step for the specified nodes.
% [nodes, loglik] = EM_step(bnet, nodes, vals, obs_nodes, soft_vals, soft_nodes, cliques)
%
% cliques is an optional argument used by learn_dbn_structure.
% c=cliques(i) means the marginal for nodes{i} can be found in clique c.
% (We need this because these virtual nodes don't yet exist in the graph.)

if ~exist('soft_vals'), soft_vals = []; end
if ~exist('soft_nodes'), soft_nodes = []; end
if ~exist('cliques'),  cliques = []; end

verbose = 0;
loglik = 0;

num_nodes = length(nodes);
for i=1:num_nodes
  if nodes{i}.adjustable
    nodes{i} = reset_ess_bnode(nodes{i});
  end
end

num_cases = length(vals);
for l=1:num_cases
  if isempty(soft_vals)
    [bnet, ll] = enter_evidence(bnet, vals{l}, obs_nodes{l});
  else
    [bnet, ll] = enter_evidence(bnet, vals{l}, obs_nodes{l}, soft_vals{l}, soft_nodes{l});
  end
  loglik = loglik + ll;
  if isempty(cliques)
    nodes = update_ess(bnet, nodes);
  else
    nodes = update_ess_sem_dbn(bnet, nodes, cliques);
  end
  if verbose, fprintf(1, 'EM case %d, ll %4.3f\n', l, ll); end
end


for i=1:num_nodes
  if nodes{i}.adjustable
    nodes{i} = maximize_params_bnode(nodes{i});
  end
end


%%%%%%%%%%%%

function nodes = update_ess(bnet, nodes)
% UPDATE_ESS Updated the expected sufficient statistics for each adjustable node X.
% nodes = update_ess(bnet, nodes)

% Compute the marginal probability Pr(Y, Pa | E),
% where E is the evidence and Pa are parents of Y,
% for each Y in the equivalence class of X.

% To do online EM, we should multiply the marginals by 0 < gamma < 1.

if bnet.is_dbn
  ss = bnet.slice_size;
  T = bnet.num_slices;
  t = 1;
  for i=1:ss
    e = bnet.equiv_class(i);
    if nodes{e}.adjustable
      m = marginal_family(bnet, i, t);
      nodes{e} = update_ess_bnode(nodes{e}, m);
    end
  end
  for i=1:ss
    e = bnet.equiv_class(i+ss);
    if nodes{e}.adjustable
      for t=2:T
	m = marginal_family(bnet, i, t);
	nodes{e} = update_ess_bnode(nodes{e}, m);
      end
    end
  end
else
  for e=1:length(nodes)
    if nodes{e}.adjustable
      equiv = bnet.members_of_equiv_class{e};
      for j=1:length(equiv)
	i = equiv(j);
	m = marginal_family(bnet, i, 1);
	nodes{e} = update_ess_bnode(nodes{e}, m);
      end
    end
  end
end


%%%%%%%%

function nodes = update_ess_sem_dbn(bnet, nodes, cliques)
% UPDATE_ESS_SEM_DBN Updated the expected sufficient statistics for each virtual node.
% nodes = update_ess_sem_dbn(bnet, nodes, cliques)

ss = bnet.slice_size;
for i=1:length(nodes)
  c = cliques(i);
  for t=2:bnet.num_slices
    fam = [nodes{i}.parents nodes{i}.self] + (t-2)*ss;
    m = marginal_nodes(bnet, fam, t, c);
    nodes{i} = update_ess_bnode(nodes{i}, m);
  end
end

