% Check that softmax works with a simple classification demo.
% Based on netlab's demglm2
% X -> Q where X is an input node, and Q is a softmax

  
% Check inference

input_dim = 2;
num_classes = 3;
IRLS_iter = 3;

net = glm(input_dim, num_classes, 'softmax');
x = rand(1, input_dim);
q = glmfwd(net, x);

dag = zeros(2);
dag(1,2) = 1;
discrete_nodes = [2];
bnet = mk_bnet(dag, [input_dim num_classes], discrete_nodes);
bnet.bnodes{1} = mk_input_node(bnet, 1);
bnet.bnodes{2} = mk_softmax_node(bnet, 2);
bnet.bnodes{2}.max_iter = IRLS_iter;
bnet.bnodes{2}.glim = net; % make them have the same params
bnet = compile_params(bnet);


obs_nodes = 1;
bnet = enter_evidence(bnet, x, obs_nodes);
m = marginal_nodes(bnet, 2);
assert(approxeq(m.T(:), q(:)));


% Check learning
% We use EM, but in fact there is no hidden data.
% The M step will call IRLS on the softmax node.

% Generate data from three classes in 2d
input_dim = 2;
num_classes = 3;

% Fix seeds for reproducible results
randn('state', 42);
rand('state', 42);

ndata = 10;
% Generate mixture of three Gaussians in two dimensional space
data = randn(ndata, input_dim);
targets = zeros(ndata, 3);

% Priors for the clusters
prior(1) = 0.4;
prior(2) = 0.3;
prior(3) = 0.3;

% Cluster centres
c = [2.0, 2.0; 0.0, 0.0; 1, -1];

ndata1 = prior(1)*ndata;
ndata2 = (prior(1) + prior(2))*ndata;
% Put first cluster at (2, 2)
data(1:ndata1, 1) = data(1:ndata1, 1) * 0.5 + c(1,1);
data(1:ndata1, 2) = data(1:ndata1, 2) * 0.5 + c(1,2);
targets(1:ndata1, 1) = 1;

% Leave second cluster at (0,0)
data((ndata1 + 1):ndata2, :) = ...
  data((ndata1 + 1):ndata2, :);
targets((ndata1+1):ndata2, 2) = 1;

data((ndata2+1):ndata, 1) = data((ndata2+1):ndata,1) *0.6 + c(3, 1);
data((ndata2+1):ndata, 2) = data((ndata2+1):ndata,2) *0.6 + c(3, 2);
targets((ndata2+1):ndata, 3) = 1;



%ndata = 1;
%data = x;
%targets = [1 0 0];

options = foptions;
options(1) = 0;
options(14) = IRLS_iter;
net_hat = glmtrain(net, options, data, targets);

observed_values = cell(ndata,1);
for l=1:ndata
  q = find(targets(l,:)==1);
  x = data(l,:);
  observed_values{l} = [x q];
end

obs_nodes = [1 2];
max_iter = 2; % we have complete observability
bnet_hat = learn_params(bnet, observed_values, obs_nodes, max_iter);

assert(approxeq(net_hat.w1, bnet_hat.bnodes{2}.glim.w1));
assert(approxeq(net_hat.b1, bnet_hat.bnodes{2}.glim.b1));


