function bnet = mk_jtree_dbn(bnet)
% MK_JTREE_DBN Make a join tree for a 1 and 2 slice DBN.
% bnet = mk_jtree_dbn(bnet)
%
% bnet.jtree contains the 2-slice join tree (see mk_jtree for details).
% bnet.jtree1 contains the 1-slice join tree.
% In addition, bnet.dbn contains the following fields, which are necessary for the
% Boyen-Koller algorithm.
%
% For the 1 slice jtree (used only in init_belief_state):
% sl1_nodes_ass_to_cl_jt1{i} = the nodes in slice 1 that are assigned to clique i
% cl_con_sl1_cluster_jt1(k) = the lightest clique that contains cluster k in slice 1
%
% For the 2 slice jtree:
% sl1_nodes_ass_to_cl{i} = the nodes in slice 1 that are assigned to clique i
% s12_nodes_ass_to_cl{i} = the nodes in slice 2 that are assigned to clique i
% cl_con_sl1_cluster(k) = the lightest clique that contains cluster k in slice 1
% cl_con_sl2_cluster(k) = the lightest clique that contains cluster k in slice 2
%
% Call this routine whenever you change bnet.dbn.cluster (which determines the approximation)
% or bnet.marginal (which determines which non-family marginals can be computed),
% since these change the jtree.

assert(bnet.is_dbn);
assert(strcmp(bnet.algo,'bk'));

% Make a join tree for 1 slice
[bnet.jtree1, bnet.dbn.sl1_nodes_ass_to_cl_jt1, bnet.dbn.cl_con_sl1_cluster_jt1] = ...
    mk_jtree_dbn_helper(unroll_dbn(bnet,1));

% Make a join tree for 2 slices
bnet = unroll_dbn(bnet, 2);
% make a 2 slice network (learn_dbn_structure may have changed inter)
[bnet.jtree, bnet.dbn.sl1_nodes_ass_to_cl, bnet.dbn.cl_con_sl1_cluster, ...
      bnet.dbn.sl2_nodes_ass_to_cl, bnet.dbn.cl_con_sl2_cluster, ...
      bnet.dbn.nodes_ass_to_cl, bnet.dbn.sl2_pnodes_ass_to_cl] = mk_jtree_dbn_helper(bnet);


%%%%%%%
function [jtree, sl1_nodes_ass_to_cl, cl_con_sl1_cluster, ...
      sl2_nodes_ass_to_cl, cl_con_sl2_cluster, nodes_ass_to_cl, sl2_pnodes_ass_to_cl] = mk_jtree_dbn_helper(bnet)

jtree = mk_jtree(bnet, bnet.dbn.cluster);
ss = bnet.slice_size;
num_slices = bnet.num_slices;
cluster = bnet.dbn.cluster;
num_clusters = length(cluster);

% Find out where the nodes in slice 1  have been assigned
num_cliques = length(jtree.cliques);
sl1_nodes_ass_to_cl = cell(1,num_cliques);
nodes_ass_to_cl = cell(1,num_cliques);
for i=1:ss
  c = jtree.cl_containing_family(i); % i in slice 1 is assigned to c
  sl1_nodes_ass_to_cl{c} = [sl1_nodes_ass_to_cl{c}, i];
  nodes_ass_to_cl{c} = [nodes_ass_to_cl{c}, i];
end
% ... and slice 2
pnodes = intersect(persistent_nodes(bnet.inter), bnet.hidden_slice);
if num_slices > 1
  sl2_nodes_ass_to_cl = cell(1,num_cliques);
  sl2_pnodes_ass_to_cl = cell(1,num_cliques);
  for i=1:ss
    c = jtree.cl_containing_family(i+ss); % i in slice 2 is assigned to c
    sl2_nodes_ass_to_cl{c} = [sl2_nodes_ass_to_cl{c}, i+ss];
    if ismember(i, pnodes)
      sl2_pnodes_ass_to_cl{c} = [sl2_pnodes_ass_to_cl{c}, i+ss];
    end
    nodes_ass_to_cl{c} = [nodes_ass_to_cl{c}, i+ss];
  end
end

% Find out where the clusters are for slice 1
cl_con_sl1_cluster = zeros(1, num_clusters);
for i=1:num_clusters
  [c, cls] = find_clique(jtree, cluster{i});
  cl_con_sl1_cluster(i) = c;
end  
% ... and slice 2
if num_slices > 1
  cl_con_sl2_cluster = zeros(1, num_clusters);
  for i=1:num_clusters
    [c, cls] = find_clique(jtree, cluster{i}+ss);
    cl_con_sl2_cluster(i) = c;
  end 
end


