%******************************************************
% Joint Gaussian Mixture Model
% Simon Rogers, 8/01/07
%******************************************************
%
% function out = jointmix(data,varargin)
%
% Performs a joint clustering on the two datasets data{1} and data{2}.
% Datasets should have the same number of data points, N (rows)
% Options should be passed in pairs, the following options can be set
% (defaults in brackets)
% covopt (diag) -> restrictions on covariance matrices {'full','diag'}
% no_its (100) -> maximum number of iterations
% verbose (1) -> 0 = nothing, 1 = text, 2 = plots and text
% tol (1e-6) -> Stopping tolerance in relative likelihood change
% k (2) -> Number of components in first level cluster
% j (3) -> Number of cimponents in second level cluster
% outx (data) -> Second data file to calculate the held-out likelihood
% condprioropt (free) -> sets restrictions on p(j|k) and hence chooses the
%   type of algorithm.  Options are 
%       'free' - p(j|k) is inferred
%       'dependent' - p(j|k) = 1 iff j==k (note J must equal K)
%       'independent' - p(j|k) = p(j), i.e. there is no dependence on k 
%           and we have two uncoupled mixtures.
% startind (0) -> Only relevant if 'free' model is used - does the
%       first 20 iterations with independent and then links.
%dirpriors (1) -> parameter of Dirichlet priors on conditional multinomials
%(pseudo-counts).  Setting this value lower will attempt to enforce
%sparsity but may lead to numerical instability
%covconst ([]) -> We did some experiments with a covariance function where
%each diagonal element was a constant multiplied by a single variance
%value.  If you set covconst to a vector of 1s, it is the same as using
%spherical covariance (variance is the same in each dimension/at each
%timepoint.  Note that this is only used if covopt = 'timedep'
%
% out has various fields...
% out.means1 -> K*M1 matrix of first level means
% out.vars1 -> K dimensional cell array of level one M1*M1 covariances
% out.condpriors -> K*J matrix of conditional priors
% out.means2 -> K*M2 matrix of second level means
% out.vars2 -> K dimensional cell array of level two M2*M2 covariances
% out.priors -> 1*K vector p(k)
% out.like -> likelihood values on training data
% out.likeout -> likelihood values on held-out data
% out.diff -> relative likelihood changes
function out = jointmix(data,varargin)

% Set some default values
K = 2;
J = 3;

[N M1] = size(data{1});
[N M2] = size(data{2});
covopt = 'diag';
no_its = 100;
tol = 1e-6;
outx = data;
verbose = 1;
condprioropt = 'free';
startind = 0;
dirpriors = 1;
for i = 1:2:length(varargin)-1
    switch lower(varargin{i})
        case 'covopt'
            covopt = varargin{i+1};
        case 'no_its'
            no_its = varargin{i+1};            
        case 'verbose'
            verbose = varargin{i+1}; %0 nothing, 1 text, 2 plots+text
        case 'tol'
            tol = varargin{i+1};
        case 'k'
            K = varargin{i+1};
        case 'outx'
            outx = varargin{i+1};
        case 'j'
            J = varargin{i+1};
        case 'condprioropt'
            condprioropt = varargin{i+1};
        case 'covconst'
            covconst = varargin{i+1};
        case 'dirpriors'
            dirpriors = varargin{i+1};
    end
end
% Check the length of dirpriors - if it is a scalar, replicate it to make
% the correct length vector
if length(dirpriors)~=J
    dirpriors = repmat(dirpriors(1),1,J);
end

if verbose
    fprintf('\n\n%s\nJoint Mixture Model',repmat('*',1,30));
    fprintf('\nSimon Rogers, 7/01/06');
    fprintf('\n%s\n',repmat('*',1,30));
    fprintf('\nModel is ',condprioropt);
    if strcmp(condprioropt,'free')
        fprintf('Unconstrained, p(j,k) = p(j|k)p(k)');
    else
        fprintf('Constrained such that ');
        if strcmp(condprioropt,'independent')
            fprintf(' p(j,k) = p(j)p(k)');
        else
            fprintf(' p(j,k) = (j==k)');
        end
    end
end
if verbose, fprintf('\n\n'),end;
if strcmp(condprioropt,'dependent') & J~=K
    fprintf('\n\nERROR: dependent conitional prior and J~=K, setting J=K\n\n');
    J = K;
end

Nout = size(outx{1},1);
% If there is missing data, initially fill it with random values that are
% in some way characteristic
for i = 1:2
    misv{i} = (isnan(data{i}));
    Nmisv{i} = sum(sum(misv{i}));
    misvout{i} = (isnan(outx{i}));
    Nmisvout{i} = sum(sum(misvout{i}));
    pos = find(~isnan(data{i}));
    data{i}(misv{i}) = randn(Nmisv{i},1)*std(data{i}(pos)) + ...
        mean(data{i}(pos));
    pos = find(~isnan(outx{i}));
    outx{i}(misvout{i}) = randn(Nmisvout{i},1)*std(outx{i}(pos)) + ...
        mean(outx{i}(pos));
end

% Initialise p(k) and p(j|k)
priors = repmat(1/K,1,K);
switch condprioropt
    case 'free'
        condpriors = repmat(1/J,K,J);
    case 'independent'
        condpriors = repmat(1/J,K,J);
    case 'dependent'
        condpriors = eye(K);
end
% Initialise component means
s = find(sum(isnan(data{1}),2)==0);
order = randperm(length(s));
for k = 1:K
    means1(k,:) = data{1}(s(order(k)),:);
end
s = find(sum(isnan(data{2}),2)==0);
order = randperm(length(s));
for j = 1:J
    means2(j,:) = data{2}(s(order(j)),:);
end

% Initialise component covariances
for k = 1:K
    switch covopt
        case 'diag'
            vars1{k} = diag(var(data{1}));
        case 'full'
            vars1{k} = cov(data{1});   
        case 'timedep'
            vars1{k} = diag(covconst{1}*var(data{1}(:)));
    end
end
for j = 1:J
    switch covopt
        case 'diag'
            vars2{j} = diag(var(data{2}));
        case 'full'
            vars2{j} = cov(data{2});            
        case 'timedep'
            vars2{j} = diag(covconst{2}*var(data{2}(:)));

    end
end

gam = rand(N,K,J);
gam = gam./repmat(sum(gam,3),[1 1 J]);
it = 0;
like = zeros(no_its,1);
likeout = like;
diff = zeros(no_its,2);
diff(1) = Inf;

condprioropt_temp = condprioropt;
if startind
  switch condprioropt
   case 'free'
    condprioropt_temp = 'independent';
  end
end
while it <= no_its
    it = it + 1;
    if strcmp(condprioropt,'free') & strcmp(condprioropt_temp,'independent') ...
          & it > 20 & startind
      condprioropt_temp = 'free';
    end
    for j = 1:J
        pdfvals2(:,j) = mynormpdf(data{2},means2(j,:),vars2{j});
        pdfvals2out(:,j) = mynormpdf(outx{2},means2(j,:),vars2{j});
    end
    
    % Update second level assignments
    for j = 1:J
        gam(:,:,j) = repmat(pdfvals2(:,j),1,K) .* ...
            repmat(condpriors(:,j)',N,1);
    end   
    
    
    gam = gam./repmat(sum(gam,3),[1 1 J]);
    pos = find(gam<1e-30);
    gam(pos) = 1e-30;
    gam = gam./repmat(sum(gam,3),[1 1 J]);

    % Compute the first level assignments    
    for k = 1:K
        pdfvals1(:,k) = mynormpdf(data{1},means1(k,:),vars1{k});
        pdfvals1out(:,k) = mynormpdf(outx{1},means1(k,:),vars1{k});
        assignments(:,k) = priors(k)*pdfvals1(:,k).*...
            sum(pdfvals2.*reshape(gam(:,k,:),N,J),2);
    end
    % Calculate the in and out likelihoods
    for k = 1:K
        Lsecond(:,k) = sum(pdfvals2.*repmat(condpriors(k,:),N,1),2);
        Lsecondout(:,k) = sum(pdfvals2out.*repmat(condpriors(k,:),Nout,1),2);
    end
    like(it) = sum(safelog(sum(pdfvals1.*repmat(priors,N,1).*Lsecond,2)))/N;
    likeout(it) = sum(safelog(sum(pdfvals1out.*...
        repmat(priors,Nout,1).*Lsecondout,2)))/Nout;
    if it>1
        diff(it) = abs(like(it)-like(it-1));
        diff(it) = diff(it)/abs(like(it));
    end
    % Add a small val for numerical stability
    assignments = assignments./repmat(sum(assignments,2),1,K);
    pos = find(assignments<1e-10);
    assignments(pos) = 1e-10;
    assignments = assignments./repmat(sum(assignments,2),1,K);
    
    
    priors = sum(assignments,1)./N;

    switch condprioropt_temp
        case 'free'
            for j = 1:J
                condpriors(:,j) = sum(assignments.*gam(:,:,j),1)' + ...
                    repmat(dirpriors(j)-1,K,1);
                condpriods(:,j) = condpriors(:,j)./...
                    (sum(assignments,1)'+sum(dirpriors)-J);
                if any(condpriors(:)<0)
                    condpriors(find(condpriors<0)) = 0;
                end
            end
            condpriors = condpriors./repmat(sum(condpriors,2),1,J);
        case 'independent'
            condpriors = repmat(sum(reshape(gam(:,1,:),N,J),1)./N,K,1);
    end
            
    % Update the parameters
    for k = 1:K       
        means1(k,:) = sum(data{1}.*repmat(assignments(:,k),1,M1))/...
            sum(assignments(:,k));
        mk = data{1} - repmat(means1(k,:),N,1);
        if ~strcmp(covopt,'timedep')
            vars1{k} = (mk.*repmat(assignments(:,k),1,M1))'*mk./...
                        sum(assignments(:,k));
            pos = find(vars1{k})<1e-10;
            vars1{k}(pos) = 1e-10;
            switch covopt
                case 'diag'
                    vars1{k} = diag(diag(vars1{k}));
            end   
        else
            % Time dependent covariates
            sigk = sum(sum(repmat(1./covconst{1},N,1).*(mk.^2),2).*...
                assignments(:,k))./(M1*sum(assignments(:,k)));
            vars1{k} = diag(covconst{1}*sigk);
        end
    end
    
    for j = 1:J
        means2(j,:) = sum(data{2}.*repmat(...
            sum(assignments.*gam(:,:,j),2),1,M2),1)./...
            sum(sum(assignments.*gam(:,:,j)));
        mj = data{2} - repmat(means2(j,:),N,1);
        if ~strcmp(covopt,'timedep')
            vars2{j} = (mj.*repmat(sum(assignments.*gam(:,:,j),2),1,M2))'*mj./...
                sum(sum(assignments.*gam(:,:,j)));
            pos = find(vars2{j})<1e-10;
            vars2{j}(pos) = 1e-10;
            switch covopt
                case 'diag'
                    vars2{j} = diag(diag(vars2{j}));
            end
        else
            % Time dependent covariance
            sigJ = sum(sum(repmat(1./covconst{2},N,1).*(mj.^2),2).*...
                sum(assignments.*gam(:,:,j),2))/...
                (M2*sum(sum(assignments.*gam(:,:,j))));
            vars2{j} = diag(covconst{2}*sigJ);
            
        end
    end 
    % Compute second level assignments
    
    ass2 = reshape(sum(gam.*...
            repmat(assignments,[1 1 J]),2),N,J);
    % Update any missing values
    if Nmisv{1}>0
        wmean = assignments * means1;
        data{1}(misv{1}) = wmean(misv{1});
    end
    if Nmisvout{1}>0
        wmeanout = repmat(priors*means1,size(outx{1},1),1);
        outx{1}(misvout{1}) = wmeanout(misvout{1});
    end
    if Nmisv{2}>0
        wmean = ass2 * means2;
        data{2}(misv{2}) = wmean(misv{2});
    end
    if Nmisvout{2}>0
        wmeanout = repmat(priors*condpriors*means2,size(outx{2},1),1);
        outx{2}(misvout{2}) = wmeanout(misvout{2});
    end
    if verbose
        fprintf('\nIt: %g, Like: %g, Likeout: %g, diff: %g',...
            it,like(it),likeout(it),diff(it));
    end
    if verbose>1
        % Do some plots     
        figure(1)
        subplot(2,3,1),plot([like(1:it) likeout(1:it)]);
        title('Log Likelihood')
        subplot(2,3,2),imagesc(assignments);
        title('First Level Assignments');
        subplot(2,3,3),imagesc(ass2);
        title('Second Level Assignments');
        subplot(2,3,5),imagesc(condpriors);
        title('Conditional Priors');
        subplot(2,3,4),imagesc(priors);
        title('Priors');
        subplot(2,3,6),hold off;
        semilogy(diff(1:it));
        hold on
        plot(1:it,repmat(tol,1,it),'r--');
        ylim([tol/10 max(max(diff))]);
        title('Relative Likelihood Change');
        drawnow
    end
    if diff(it)<tol
        if verbose
            fprintf('\nFinished - relative likelihood change (%g) < tolerance (%g)',diff(it),tol);
        end
        break
    end
end
if verbose
    if it>no_its
        fprintf('\nFinished - reached max iterations');
    end
    fprintf('\n\n');
end
if it<no_its
    like(it+1:no_its) = [];
    likeout(it+1:no_its) = [];
end
out.means1 = means1;
out.vars1 = vars1;
out.condpriors = condpriors;
out.means2 = means2;
out.vars2 = vars2;
out.priors = priors;
out.like = like;
out.likeout = likeout;
out.diff = diff;
out.assigments = assignments;
out.ass2 = ass2;
out.gam = gam;
function o = mynormpdf(data,mean,covar)

covinv = inv(covar);% 0.1 added for numerics
d = length(mean);
dm = data - repmat(mean,size(data,1),1);

o = exp(-0.5*diag(dm*covinv*dm'));
o = o./((2*pi)^(d/2));
o = o./sqrt(det(covar));

function o = safelog(x)
pos = find(x<1e-300);
x(pos) = 1e-300;
o = log(x);