function [bnet, changelog] = search_sem_dbn(bnet, ev, em_options, sem_options, must_connect, must_not_connect)
% SEM_DBN Use Structural EM to learn the inter-slice connectivity matrix of a DBN.
% [bnet, changelog] = sem_dbn(bnet, ev, em_options, sem_options, must_connect, must_not_connect)
%
% ev and options are as in learn_params.
% Optional arguments:
% must_connect(p,i) = 1 if we must have an arc from p to i.
% must_not_connect(p,i) = 1 if we must not have an arc from p to i.
%
% See "Learning the structure of DBNs", Friedman, Murphy and Russell, UAI 1998.

ss = bnet.slice_size;
if nargin < 3, em_options = []; end
if nargin < 4, sem_options = []; end
if nargin < 5, must_connect = zeros(ss); end
if nargin < 6, must_not_connect = zeros(ss); end

if isfield(sem_options, 'verbose')
  verbose = sem_options.verbose;
else
  verbose = 1;
end
if isfield(sem_options, 'max_iter')
  max_iter = sem_options.max_iter;
else
  max_iter = 4;
end

converged = 0;
num_iter = 1;
changelog = {};

while ~converged  & (num_iter <= max_iter) %{
  
  if verbose, fprintf(1, '\nSEM iteration %d\n', num_iter); end
  num_iter = num_iter + 1;
  
  % Get a rough guess of parameters for current model, before we compute ESS for future families
  if ~isempty(em_options)
    [bnet, LL] = learn_params(bnet, ev, em_options);
  else
    [bnet, LL] = learn_params(bnet, ev);
  end
  
  if approxeq(LL(end), 0) % fits the data perfectly
    return;
  end
  
  % pp_add(p,i) = 1 means we can add p->i
  % pp_rem(p,i) = 1 means we can remove p->i
  % pp = potential parent
  pp_add = ones(ss) - bnet.inter; % all parents in prev slice except those we have already
  pp_add(logical(must_not_connect))=0;
  pp_rem = bnet.inter; % all parents in prev slice
  pp_rem(logical(must_connect))=0;

  % vnodes{j}, where j = index_add(p,i), means virtual node j is the same as node i except
  % it has an extra parent p, etc.
  vnodes = {};
  families = {};
  [vnodes, index_add, families] = mk_vnodes(pp_add, bnet, vnodes, families, 1);
  [vnodes, index_rem, families] = mk_vnodes(pp_rem, bnet, vnodes, families, 0);
  [bnet, cliques] = schedule_marginals(bnet, families);
  vnodes = EM_step(bnet, vnodes, ev, cliques);

  % For each node, find its best family.
  nsamples = bnet.num_slices;
  changed = 0;
  for i=1:ss
    bnode = bnet.bnodes{i+ss};
    old_score = log_likelihood_bnode(bnode) - bnode.nparams*0.5*log(nsamples);

    score_add = ones(1,ss)*-inf;
    pp = find(pp_add(:,i));
    for pi=1:length(pp)
      p = pp(pi);
      bnode = vnodes{index_add(p,i)};
      score_add(p) = log_likelihood_bnode(bnode) - bnode.nparams*0.5*log(nsamples);
    end

    score_rem = ones(1,ss)*-inf;
    pp = find(pp_rem(:,i));
    for pi=1:length(pp)
      p = pp(pi);
      bnode = vnodes{index_rem(p,i)};
      score_rem(p) = log_likelihood_bnode(bnode) - bnode.nparams*0.5*log(nsamples);
    end

    score = max(score_add, score_rem);
    [best_score, best_parent] = max(score - old_score);
    threshold = 1;
    if best_score < threshold
      s = sprintf('not changing parents of node %d\n', i);
    else
      if pp_add(best_parent,i)
	bnode = vnodes{index_add(best_parent,i)};
	bnet.inter(best_parent,i) = 1;
	s = sprintf('adding parent %d to node %d\n', best_parent, i);
      else
	bnode = vnodes{index_rem(best_parent,i)};
	bnet.inter(best_parent,i) = 0;
	s = sprintf('removing parent %d from node %d\n', best_parent, i);
      end
      bnet.bnodes{i+ss} = bnode;
      changed = 1;
    end
    if verbose, fprintf(s); end
    changelog{end+1} = s;
  end
  
  if ~changed
    SEM_converged = 1;
  else
    bnet = mk_jtree_dbn(bnet);
  end
  
end %}




%%%%%%%%%%
function [vnodes, index, families] = mk_vnodes(pp, bnet, vnodes, families, add)

ss = length(pp);
index = zeros(ss);
[ps is] = find(pp);
parent_offset = 0; % since in slice 2
child_offset = ss; % since in slice 1
for j=1:length(ps)
  p = ps(j) + parent_offset;
  i = is(j) + child_offset;
  if add
    parents = union(bnet.bnodes{i}.parents, p);
  else
    parents = setdiff(bnet.bnodes{i}.parents, p);
  end
  families{end+1} = [parents i];
  n = length(vnodes)+1;
  vnodes{n} = modify_parents_bnode(bnet.bnodes{i}, parents, bnet);
  index(p - parent_offset, i - child_offset) = n;
end

