function [pi,alpha,theta,Lik] = ccg( counts, E, W, options )
% [pi,pl,Lq,loglikelihood_samples] =  ccg( counts, E, W, options )
% 
% Componential Counting Grids - NIPS 2013
%
% ---- INPUT ----
% counts:            data, features x samples
% E:                 counting grid size  (e.g., [30,30])
% W:                 window size (e.g., [4 4])
% options (Struct) - All fields are optional. 
%   min_change:      convergence criterion. Min relative change in 
%                    loglikelihood (defalut 1e-5)
%   max_iter:        max number of EM iteration (default 100)
%   learn_pi:        learn counting grid \pi (default 1=true)
%   learn_alpha:     learn dirichlet prior on location admixtures
%                    (default 0=alse - May give instabilities if learned)
%   alpha_value:     initialization for the parameters of the dirichlet 
%                    distribution alpha (default = \alpha_i = 1.01)
%   eta_value:       initialization for the parameters of the dirichlet 
%                    prior on pi, eta (default = \eta_i = 1.0001)
%   pi:              Learned (Componential) Counting Grid (for inference, 
%                    or to pass an initialization as input)
%   alpha:           Learned Dirichlet Prior (for inference, or to pass an
%                    initialization)
%
% ---- OUTPUT ----
% pi:                counting grid pi(\bf{k},z) = \pi_{\bf{k},z}
% alpha:             dirichlet prior on theta \alpha_\bf{k}
% theta:             Admixture of locations for each sample \theta_\bf{k}^t
% Lik:               (final) loglikelihood of each sample
% 
%
% Written by Alessandro Perina, alessandro.perina@gmail.com /
% alperina@microsoft.com
dbstop if error

if ~exist( 'options', 'var'); options = []; end
if ~isfield( options,'min_change'); options.min_change = 1e-4; end
if ~isfield( options,'max_iter'); options.max_iter = 100; end
if ~isfield( options,'learn_pi'); options.learn_pi = 1; end
if ~isfield( options,'learn_alpha'); options.learn_alpha = 0; end
if ~isfield( options,'alpha_value'); options.alpha_value = 1.01; end
if ~isfield( options,'eta_value'); options.eta_value = 1.0001; end

L = prod(E);
[Z,T] = size(counts);
total = sum(counts(:));

if isfield( options,'pi')
    pi = options.pi;
    if prod(W) ~= 1
        extPi = padarray(pi,W,'circular','post' );
        H = permute(cumsum( permute( cumsum(extPi),[2 1 3]) ),[2 1 3] );
        Hz = padarray( H,[1 1],0,'pre');
        tmp = compute_h_noLoopFull( Hz, W(2)+1, W(1)+1);
        h = bsxfun(@rdivide, tmp(1:end-1,1:end-1,:), sum( tmp(1:end-1,1:end-1,:),3 ));
    else
        % W = [1 1] - CCG is LDA ("LDA is a special case of CCG :)") 
        h = pi;
    end
else
    tmp = 3+4*rand([E,Z]);
    pi = bsxfun(@rdivide, tmp, sum(tmp,3));
    if prod(W) ~= 1
        extPi = padarray(pi,W,'circular','post' );
        H = permute(cumsum( permute( cumsum(extPi),[2 1 3]) ),[2 1 3] );
        Hz = padarray( H,[1 1],0,'pre');
        tmp = compute_h_noLoopFull( Hz, W(2)+1, W(1)+1);
        h = bsxfun(@rdivide, tmp(1:end-1,1:end-1,:), sum( tmp(1:end-1,1:end-1,:),3 ));
    else
        h = pi;
    end
end


alpha = options.alpha_value *ones(L,1);
eta = options.eta_value*ones(Z,1);

tmp  = gamrnd(repmat(alpha',T,1),1,T,L)'; tmp = bsxfun( @rdivide, tmp, sum(tmp) ); tmp2 = tmp*sum(counts(:))/T;
theta = bsxfun(@rdivide, tmp2, sum(tmp2) );
max_iter = options.max_iter;
U = padarray( ones(W),E-W,0,'post') ./ prod(W); % U^W in the paper.
free_energy = zeros(1,max_iter);
iter = 0;
done = 0;

while ~done; iter= iter+1;
  
    % E-Step
    sum_posts = theta .* (  reshape( h, [L,Z])* (counts./( reshape( h, [L,Z])'*theta )) );
    theta  = bsxfun(@plus, alpha-1, sum_posts);
    theta( theta< 0 ) = 0;
    theta = bsxfun( @rdivide, theta, sum(theta) );
    
    % M-Step
    if options.learn_pi
        
        if prod(W) == 1
            exp_theta   = theta;
        else
            exp_theta = zeros(L,T);
            for d=1:T
                tmp = real( ifft2( fft2( reshape( theta(:,d),E )).*fft2( U ) ) );
                tmp(tmp<0)=0;
                exp_theta(:,d) = tmp(:);
            end
        end
       sum_posts   = reshape( pi, [L,Z])'.* ( (counts ./ ( reshape( pi, [L,Z])'*exp_theta +eps ) )* exp_theta' );
       
       pi = reshape( bsxfun( @rdivide, bsxfun(@plus, sum_posts, eta-1), sum( bsxfun(@plus, sum_posts, eta-1) ))',[E,Z]);
       if prod(W)~=1
           extPi = padarray(pi,W,'circular','post' );
           H = permute(cumsum( permute( cumsum(extPi),[2 1 3]) ),[2 1 3] );
           Hz = padarray( H,[1 1],0,'pre');
           tmp = compute_h_noLoopFull( Hz, W(2)+1, W(1)+1);
           h = bsxfun( @rdivide, tmp(1:end-1,1:end-1,:), sum(tmp(1:end-1,1:end-1,:),3));
       else
           h = pi;
       end
       
    end
   
    Lik = sum( counts.*log( reshape(h,[L,Z])' * theta ),1) ;
    if options.learn_alpha
        % Use a dirichlet distribution estimator. Like the one in the
        % fastfit toolbox by T.Minka 
        % alpha = dirichlet_fit( theta )
        Lik =  Lik + dirichlet_logProb( alpha, theta );
    end    
    
    % Check for convergence
    free_energy(iter) = -sum( Lik );
    done = iter >= max_iter;
    if iter > 20
        F1 = -free_energy(iter)/total;
        F2 = -free_energy(iter-1)/total;
        rel_ch = (F1-F2) / abs(mean([F1,F2]));
        if rel_ch < options.min_change
            done = 1;
        end
    end

end;  
end


function h = compute_h_noLoopFull( H, xW, yW )
        h = H(yW:end,xW:end,:,:,:) - H(1:end-yW+1,xW:end,:,:,:) ...
            - H(yW:end,1:end-xW+1,:,:,:) + H(1:end-yW+1,1:end-xW+1,:,:,:);    
end
