% Check correctness of an HMM with Mixture of Gaussians output.

%
%   q1 -> q2 -> ...
%  / |   / |
% m1 |  m2 |
%  \ |   \ |
%   y1    y2
%
% where Pr(m=j|q=i) is a multinomial and Pr(y|m,q) is a Gaussian

intra = zeros(3);
intra(1,2) = 1; intra(1,3)=1;
intra(2,3) = 1;
inter = zeros(3);
inter(1,1) = 1;
ss = length(intra);
Q = 3;
M = 2;
O = 3;
cov_type = 'diag';
node_sizes = [Q M O];
obs_nodes = [3];
discrete_nodes = [1 2];
equiv_class = [1 2 3   4 2 3];
bnet = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'bk', 'exact');
bnet.bnodes{1} = mk_tabular_node(bnet, 1);
bnet.bnodes{2} = mk_tabular_node(bnet, 2);
bnet.bnodes{3} = mk_gaussian_node(bnet, 3, cov_type);
bnet.bnodes{4} = mk_tabular_node(bnet, 4);
bnet = compile_params(bnet);

T = 4;
num_training = 10;
max_iter = 2;
[init_state_prob, transmat, mixmat, mu, Sigma] = extract_mhmm_from_dbn(bnet);
data = sample_mhmm(T, num_training, init_state_prob, transmat, mu, Sigma, mixmat);
observed_values = num2cell(data, [1 2]); % put dims 1 and 2 in the same cell

% Start both models off from the same point in parameter space, bnet0
%bnet0 = randomize_params(bnet);
%[init_state_prob0, transmat0, mixmat0, mu0, Sigma0] = can_dbn_to_mhmm(bnet0);
bnet0 = bnet;
[init_state_prob0, transmat0, mixmat0, mu0, Sigma0] = init_mhmm(data, Q, M, cov_type, 0);
bnet0.bnodes{1}.CPT = init_state_prob0;
bnet0.bnodes{2}.CPT = mixmat0;
bnet0.bnodes{3}.mean = mu0;
bnet0.bnodes{3}.cov = Sigma0;
bnet0.bnodes{4}.CPT = transmat0;
bnet0 = compile_params(bnet0);
verbose = 1;

[ll, init_state_prob, transmat, mu, Sigma, mixmat] = ...
    learn_mhmm(data, init_state_prob0, transmat0, mu0, Sigma0, mixmat0, ...
    max_iter, 1e-3, verbose, cov_type);

[bnet_hat, ll1] = learn_params(bnet0, data, obs_nodes, max_iter);

[init_state_prob1, transmat1, mixmat1, mu1, Sigma1] = extract_mhmm_from_dbn(bnet_hat);

assert(approxeq(ll, ll1));
assert(approxeq(init_state_prob, init_state_prob1));
assert(approxeq(transmat, transmat1));
assert(approxeq(mixmat, mixmat1));
assert(approxeq(mu, mu1));
assert(approxeq(Sigma, Sigma1));

