% Check mixture of Gaussians works by comparing with netlab
% Q -> Y

dag = zeros(2);
dag(1,2) = 1;
Q = 2; Y = 5;
node_sizes = [Q Y];
obs_nodes = [2];
discrete_nodes = 1;
bnet = mk_bnet(dag, node_sizes, discrete_nodes);
bnet.bnodes{1} = mk_tabular_node(bnet, 1, 0);
bnet.bnodes{2} = mk_gaussian_node(bnet, 2, 'diag');
bnet = compile_params(bnet);
bnet0 = bnet;

% Generate the data
randn('state', 0); rand('state', 0);
ndat1 = 20; ndat2 = 20; ndata = ndat1+ndat2;
x1 = gsamp(bnet.bnodes{2}.mean(:,1), bnet.bnodes{2}.cov(:,:,1), ndat1);
x2 = gsamp(bnet.bnodes{2}.mean(:,2), bnet.bnodes{2}.cov(:,:,2), ndat2);
x = [x1; x2];

max_iter = 5;

mix = gmm(Y, Q, 'diag');
options = foptions;
options(1) = 1; % verbose
options(14) = max_iter;
%mix = gmminit(mix, x, options); % Initialize with K-means
mu0 = mix.centres';
[mix, options] = gmmem(mix, x, options);
ll1 = options(8);
mu1 = mix.centres';

bnet.bnodes{1}.CPT = normalise(ones(Q,1));
bnet.bnodes{2}.mean = mu0;
bnet.bnodes{2}.cov = repmat(eye(Y),[1 1 Q]);
bnet  = compile_params(bnet);
[bnet_hat, ll2] = learn_params(bnet, x, obs_nodes, max_iter);
mu2 = bnet_hat.bnodes{2}.mean;
S2 = bnet_hat.bnodes{2}.cov;

% assert(approxeq(ll1,ll2)); % gmmem returns the value after the final M step, BNT before
assert(approxeq(mu1, mu2));
assert(approxeq(mix.covars(1,:), diag(S2(:,:,1))));
assert(approxeq(mix.covars(2,:), diag(S2(:,:,2))));


