% Check correctness of inference and learning in a discrete DBN.

N = 5;
single_obs = 0;
[intra, inter, node_sizes, obs_nodes] = mk_cancer_topology(N, single_obs);
ss = length(intra); % slice size
discrete_nodes = 1:ss;
equiv_class = mk_equiv_classes_for_dbn(inter, intra, intra);

% Make a DBN with random parameters
bnet = mk_dbn(intra, inter, node_sizes, discrete_nodes,  equiv_class, obs_nodes, 'bk');
bneth = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'hmm');
for i=1:length(bnet.bnodes)
  bnet.bnodes{i} = mk_tabular_node(bnet, i);
  bneth.bnodes{i} = mk_tabular_node(bneth, i);
end
% Make sure they have the same params
bnet = randomize_params(bnet, 0);
bneth = randomize_params(bneth, 0);

% Make a random observation sequence
T = 3; % num slices
O = length(bnet.observed_slice);
vals = sample_discrete([0.5 0.5], O, T);


% Check inference

[bnet, llj] = enter_evidence(bnet, vals);
[bneth, llh] = enter_evidence(bneth, vals);
assert(approxeq(llj, llh));

% Evaluate the marginal on node i in slice t = node j
t=2;
i=3;
mh = marginal_nodes(bneth, i, t);
mj = marginal_nodes(bnet, i, t);
assert(approxeq(mh.T, mj.T));

% Evaluate the marginal on node 3's family
t = 2;
mh = marginal_family(bneth, i, t);
mj = marginal_family(bnet, i, t);
assert(approxeq(mh.T, mj.T));




% Check learning 

ss = 2; % slice size
intra = zeros(ss, ss);
intra(1,2) = 1;
inter = zeros(ss, ss);
inter(1,1) = 1;
node_sizes = [2 4]; % 2 hidden states, 4 output symbols

dnodes = [1 2];
eclass = [1 2 3 2];
obs_nodes = [2]; % need to know what nodes are observed to convert to HMM
bnet = mk_dbn(intra, inter, node_sizes, dnodes, eclass, obs_nodes);
for i=1:length(bnet.bnodes)
  bnet.bnodes{i} = mk_tabular_node(bnet, i);
end
seed = 0;
bnet = randomize_params(bnet, seed);

% Generate training data from the true model
seed = 1;
rand('state', seed);
T = 5;
ntrain = 3;
vals = cell(1, ntrain);
for i=1:ntrain
  sample = sample_dbn(bnet, T);
  vals{i} = sample(obs_nodes, :);
end
celldisp(vals)

bnet1 = randomize_params(bnet, 2);
max_iter = 5;
tic; [bnet2, ll2] = learn_params(bnet1, vals, obs_nodes, max_iter); toc

[prior1, transmat1, obsmat1] = dbn_to_hmm(bnet1);
tic; [ll3, prior3, transmat3, obsmat3] = learn_hmm(vals, prior1, transmat1, obsmat1, max_iter); toc
[prior2, transmat2, obsmat2] = dbn_to_hmm(bnet2);

assert(approxeq(ll3, ll2));
assert(approxeq(prior3(:), prior2(:)));
assert(approxeq(transmat3, transmat2));
assert(approxeq(obsmat3, obsmat2));


