% Factor analysis
% Z -> X,  Z in R^k, X in R^D, k << D (high dimensional observations explained by small source)
% Z ~ N(0,I),   X|Z ~ N(L z, Psi), where Psi is diagonal.

% Mixtures of FA
% Now X|Z,W=i ~ N(mu(i) + L(i) Z, Psi(i))

% Special purpose matlab code for these models is available Zoubin Ghahramani's web page
% http://www.gatsby.ucl.ac.uk/~zoubin/software.html
% Here we show how to use the toolbox to solve the same problem.

% FA

k = 2;
D = 4;
N = 20;
X = randn(N, D);

% Initialize as in Zoubin's ffa
randn('seed', 0);
X=X-ones(N,1)*mean(X);
XX=X'*X/N;
diagXX=diag(XX);
cX=cov(X);
scale=det(cX)^(1/D);
L=randn(D,k)*sqrt(scale/k);
Psi=diag(cX);

ns = [k D];
dag = zeros(2,2);
dag(1,2) = 1;
dnodes = [];
onodes = 2;
bnet = mk_bnet(dag, ns, dnodes);
bnet.bnodes{1} = mk_gaussian_node(bnet, 1);
bnet.bnodes{1}.mean = zeros(k,1);
bnet.bnodes{1}.cov = eye(k);
bnet.bnodes{1}.adjustable = 0;
bnet.bnodes{2} = mk_gaussian_node(bnet, 2);
bnet.bnodes{2}.mean = zeros(D,1);
bnet.bnodes{2}.adjustable_mean = 0;
bnet.bnodes{2}.cov = diag(Psi);
bnet.bnodes{2}.cov_type = 'diag';
bnet.bnodes{2}.weights = L;
bnet = compile_params(bnet);

max_iter = 3;
[bnet, LL1] =  learn_params(bnet, X, onodes, max_iter);

L1 = bnet.bnodes{2}.weights;
Psi1 = bnet.bnodes{2}.cov;

% Compare to Zoubin's code
[L2,Psi2,LL2]=ffa(X,k,max_iter);
assert(approxeq(LL1,LL2));
assert(approxeq(Psi1,diag(Psi2)));
assert(approxeq(L1,L2));

% Mixtures of FA
% Now X|Z,W=i ~ N(mu(i) + L(i) Z, Psi(i))

M = 2;
k = 2;
D = 5;
K = k;

X1 = randn(4, D);
X2 = randn(4, D) + 2; % move the mean to (2,2,2...)
X = [X1; X2];

% initialise as in mfa
tiny=exp(-700);
mX = mean(X);
cX=cov(X);
scale=det(cX)^(1/D);
randn('state',0);
%L=randn(D,M,K)*sqrt(scale/K);
L=randn(D*M,K)*sqrt(scale/K);
Psi=diag(cX)+tiny;
Pi=ones(M,1)/M;
Mu=randn(M,D)*sqrtm(cX)+ones(M,1)*mX;


ns = [M K D];
dag = zeros(3);
dag(1,3) = 1;
dag(2,3) = 1;
dnodes = 1;
onodes = 3;
bnet = mk_bnet(dag, ns, dnodes);
bnet.bnodes{1} = mk_tabular_node(bnet, 1);
bnet.bnodes{1}.CPT = Pi;
bnet.bnodes{2} = mk_gaussian_node(bnet, 2);
bnet.bnodes{2}.mean = zeros(k,1);
bnet.bnodes{2}.cov = eye(k);
bnet.bnodes{2}.adjustable = 0;
bnet.bnodes{3} = mk_gaussian_node(bnet, 3);
bnet.bnodes{3}.mean = Mu';
%bnet.bnodes{3}.adjustable_mean = 0;
bnet.bnodes{3}.cov = repmat(diag(Psi), [1 1 M]);
bnet.bnodes{3}.cov_type = 'diag';
bnet.bnodes{3}.tied_cov = 1;
bnet.bnodes{3}.weights = permute(reshape(L, [D M K]), [1 3 2]); % use D,K,M 
bnet = compile_params(bnet);


max_iter = 3;
[bnet1, LL1] =  learn_params(bnet, X, onodes, max_iter);

% Compare to Zoubin's code
% mfa2 is like mfa, except we pass in the initial conditions
%[Lh2, Ph2,  Mu2, Pi2, LL2] = mfa(X,M,K,max_iter);
[Lh2, Ph2,  Mu2, Pi2, LL2] = mfa2(X,M,K,max_iter,L,Mu,Psi,Pi);
Lh3 = permute(reshape(Lh2, [D M K]), [1 3 2]); % use D,K,M 
assert(approxeq(LL1,LL2));
for i=1:M
  assert(approxeq(bnet1.bnodes{3}.weights(:,:,i), Lh3(:,:,i)));
  assert(approxeq(bnet1.bnodes{3}.cov(:,:,i), diag(Ph2)));
  assert(approxeq(bnet1.bnodes{3}.mean(:,i), Mu2(i,:)));
  assert(approxeq(bnet1.bnodes{1}.CPT(:), Pi2(:)));
end


