% Check correctness of inference and learning in a Gaussian DBN.
% This is the Gaussian analog of dbn1.m
  
if 1
  intra = zeros(2); intra(1,2) = 1;
  inter = zeros(2); inter(1,1)=1;
  node_sizes = [2 2];
else
  N = 5;
  single_obs = 1;
  [intra, inter, node_sizes] = mk_cancer_topology(N, single_obs);
end
node_sizes(end) = sum(node_sizes(1:end-1));
O = node_sizes(end);
% The observable node encodes the state of the hidden nodes if its weight matrix = I
% (subject to white noise)

ss = length(intra); % slice size
obs_nodes = [ss]; % the last node in each slice is observed
discrete_nodes = [];
equiv_class = mk_equiv_classes_for_dbn(inter, intra, intra);

% Make a DBN with random parameters
bnetj = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'jtree');
bnetb = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'bk', 'exact');
bneth = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'lds');
for i=1:length(bnetj.bnodes)
  bnetj.bnodes{i} = mk_gaussian_node(bnetj, i, 'full');
  bnetb.bnodes{i} = mk_gaussian_node(bnetb, i, 'full');
  bneth.bnodes{i} = mk_gaussian_node(bneth, i, 'full');
end
% Make sure they have the same params
bnetj = randomize_params(bnetj, 0);
bnetb = randomize_params(bnetb, 0);
bneth = randomize_params(bneth, 0);
bnetj = mk_zero_mean(bnetj);
bnetb = mk_zero_mean(bnetb);
bneth = mk_zero_mean(bneth);


% Make a random observation sequence
A = bneth.hmm.trans_mat;
C = bneth.hmm.obs_mat;
Q = bneth.hmm.trans_cov;
R = bneth.hmm.obs_cov;
init_x = bneth.hmm.init_state;
T = 3;
[x,y] = sample_lds(A, C, Q, R, init_x, T);
% x(i,t) = i'th component of state at time t.
% y(i,t) = i'th component of observation at time t.

% Check inference

[bnetj, llj] = enter_evidence(bnetj, y);
[bnetb, llb] = enter_evidence(bnetb, y);
[bneth, llh] = enter_evidence(bneth, y);
assert(approxeq(llj, llh));
assert(approxeq(llj, llb));

% Evaluate the marginal on node i in slice t = node j (for prediction)
t=2;
i=1;
j = i + (t-1)*bnetj.slice_size;
mh = marginal_nodes(bneth, i, t);
mj = marginal_nodes(bnetj, i, t);
mb = marginal_nodes(bnetb, i, t);
[mg.mu, mg.Sigma, llg] = naive_gaussian_marginal(bnetj, j);

assert(approxeq(mg.mu, mj.mu));
assert(approxeq(mh.mu, mj.mu));
assert(approxeq(mh.mu, mb.mu));
assert(approxeq(mg.Sigma, mj.Sigma));
assert(approxeq(mh.Sigma, mj.Sigma));
assert(approxeq(mh.Sigma, mb.Sigma));





% We can compare bneth, bnetj or bnetb to learn_kalman, since they share the same M step.

intra = zeros(2); intra(1,2) = 1;
inter = zeros(2); inter(1,1)=1;
node_sizes = [2 3];
ss = length(intra); % slice size
obs_nodes = [2];
discrete_nodes = [];
equiv_class = [1 2 3 2];
bneth = mk_dbn(intra, inter, node_sizes, discrete_nodes, equiv_class, obs_nodes, 'lds');
for i=1:length(bneth.bnodes)
  bneth.bnodes{i} = mk_gaussian_node(bneth, i, 'full');
end
bneth = mk_zero_mean(bneth);
bneth = compile_params(bneth);

% Make a random observation sequence
A = bneth.hmm.trans_mat;
C = bneth.hmm.obs_mat;
Q = bneth.hmm.trans_cov;
R = bneth.hmm.obs_cov;
init_x = bneth.hmm.init_state;
init_V = bneth.hmm.init_cov;
T = 5; % num slices
num_train = 3;
[os ss] = size(C);
x = zeros(ss, T, num_train);
y = zeros(os, T, num_train);
% x(i,t,l) = i'th component of state at time t in example l
% y(i,t,l) = i'th component of observation at time t in example l
for i=1:num_train
  [x(:,:,i), y(:,:,i)] = sample_lds(A, C, Q, R, init_x, T);
end
vals = y;
%vals = num2cell(y, [1 2]); % put an os*T matrix into num_train separate cells

%[xsmooth, Vsmooth, VVsmooth, loglik] = kalman_smoother(y(:,:,1), A, C, Q, R, init_x, init_V);

bnet0 = randomize_params(bneth, 1);
[A0, Q0, C0, R0, x0, V0] = dbn_to_lds(bnet0);

max_iter = 3;
[bnet1, ll1] = learn_params(bnet0, vals, obs_nodes, max_iter);
[A1, Q1, C1, R1, x1, V1] = dbn_to_lds(bnet1);

[A, C, Q, R, x, V, ll] = learn_kalman(vals, A0, C0, Q0, R0, x0, V0, max_iter);


assert(approxeq(ll,ll1));
assert(approxeq(A,A1));
assert(approxeq(Q,Q1));
assert(approxeq(C,C1));
assert(approxeq(R,R1));
assert(approxeq(x,x1));
assert(approxeq(V,V1));

