function [alphas,Theta,Phi,Nkj,Nwk,fitinfo,scinfo,Theta_sc] = ...
    solve_em(solver,dofitalpha0,dofitalphas,alphas,beta0,...
    X,Xval,Theta,Phi,Nkj,Nwk,...
    dohist,reltol,first_iter,maxiter,verbose,usemex,dosc,scinfo,...
    freqval,rsh,rowidx,colidx,z,initial_history)
% Implementation of the main loop of the EM algorithm

% Inputs for N documents, W words and K topics:
%   solver      - Solver name, char vector
%   dofitalpha0 - Fit topic concentration, true or false
%   dofitalphas - Fit topic probabilities, true or false
%   alphas      - K-by-1 vector of topic concentrations, double
%   beta0       - Word concentration, double scalar
%   X           - W-by-N word-document matrix, sparse double
%   Xval        - W-by-Nval validation word-document matrix, sparse double
%   Theta       - K-by-N estimates of topic probabilities per document, double.
%   Phi         - W-by-K estimates of word probabilities per topic, double
%   Nkj         - K-by-N estimates of topic counts per document, double.
%                 If [], computed from X and Theta.
%   Nwk         - W-by-K estimates of word counts per topic, double
%                 If [], computed from X and Phi.
%   dohist      - Record fitting history, true or false
%   reltol      - Relative tolerance on log-likelihood, double scalar
%   first_iter  - Index of the first iteration (1 if fit from scratch, >1
%                 if resuming), double scalar 
%   maxiter     - max number of iterations, double scalar
%   verbose     - Verbosity level, double scalar integer
%   usemex      - Use mex, true or false
%   dosc        - Use stochastic solver, true or false
%   scinfo      - struct with stochastic fit info saved earlier, [] for batch solver
%   freqval     - Frequency of print-outs for validation data
%   rsh         - randstream handle
%   rowidx      - Row indices of nonzero elements in X. Used for collapsed
%                 Gibbs sampler only. [] for other solvers. Vector of
%                 double integers.
%   colidx      - Column indices of nonzero elements in X. Used for collapsed
%                 Gibbs sampler only. [] for other solvers. Vector of
%                 double integers.
%   z           - Initial topic assignments for collapsed Gibbs sampler.
%                 Vector of double integers of the same length as rowidx
%                 and colidx.
%   initial_history - [] or struct with prior fitinfo history.
%
% Outputs:
%   alphas      - K-by-1 vector of topic concentrations, double
%   Theta       - K-by-N estimates of topic probabilities per document, double
%   Phi         - W-by-K estimates of word probabilities per topic, double
%   Nkj         - K-by-N estimates of topic counts per document, double
%   Nwk         - W-by-K estimates of word counts per topic, double
%   fitinfo     - struct with fitting info common to all solvers.
%   scinfo      - struct with fitting info specific to stochastic solvers.
%                 [] for batch solvers.
%   Theta_sc    - K-by-N estimates of topic probabilities per document
%                 obtained by stochastic solver, double. [] for batch
%                 solvers.

%   Copyright 2017-2022 The MathWorks, Inc.

% Theta = P(z|d)  K-by-N
% Phi   = P(w|z)  W-by-K

% NOTE: Nwk and Phi below are passed to C++ via shared copy. Do not save
%       Nwk and Phi in other variables.

tstart = tic;

nj = full(sum(X,1));
num_words = sum(nj);

N = numel(nj);

doval = ~isempty(Xval);
num_words_validation = full(sum(Xval(:)));

curridx = [];
if dosc
    batchsize  = scinfo.MiniBatchSize;
    decay      = scinfo.LearnRateDecay;
    scidx      = scinfo.PermutedIndices;
    
    nend          = scinfo.LastProcessedDocumentIndex;
    [nbegin,nend] = scrange(nend,batchsize,N);
    
    reltol = 0;
    curridx = scidx(nbegin:nend);
end

reltol_sc = 1e-3; % loglik tolerance for learning theta at fixed phi (final pass)
maxiter_sc = 2; % max num iterations for learning theta at fixed phi (final pass)

if dosc
    reltol_val  = reltol_sc; % tolerance for validation data
    maxiter_val = 10; % max num iterations for validation data
else
    reltol_val = reltol;
    maxiter_val = 100;
end

if isempty(curridx)
    if usemex
        logprob = clustering.ldautils.docprob_atmaxlkh(X,Phi,Theta);
    else
        logprob = docprob_atmaxlkh(X,Phi,Theta);
    end
    
    neglogL = -nansum(logprob);
else
    % Stochastic solver requires mex
    logprob = clustering.ldautils.docprob_atmaxlkh(X(:,curridx),Phi,Theta(:,curridx));

    curr_num_words = sum(nj(curridx));
    neglogL = -nansum(logprob)*num_words/curr_num_words;
end

ppl = exp( neglogL/num_words );

ppl_validation = [];
if doval && mod(first_iter,freqval)==0
    [~,ppl_validation] = evaluate_model(usemex,alphas,Phi,Xval,...
        reltol_val,maxiter_val,num_words_validation);
end

% Initial estimate of number of topics in each document
if isempty(Nkj)
    Nkj = nj.*Theta;
end

% Initial estimate of number of words assigned to each topic across the
% entire document collection
if isempty(Nwk)
    Nwk = Phi./sum(Phi,2).*full(sum(X,2));
end

if dohist
    neglogL_vals            = zeros(maxiter,1,'like',Theta);
    ppl_vals                = zeros(maxiter,1,'like',Theta);
    alpha_vals              = zeros(maxiter,1,'like',alphas);
    nalphaiter_vals         = zeros(maxiter,1,'like',alphas);
    telapsed_vals           = zeros(maxiter,1,'like',Theta);
    
    neglogL_validation_vals = NaN(maxiter,1,'like',Theta);
    ppl_validation_vals     = NaN(maxiter,1,'like',Theta);
end

alpha0 = sum(alphas);

iter = 0;

relDeltaObjective = Inf;

telapsed = toc(tstart);

if verbose>0
    fitinfolen = clustering.ldautils.showfitinfo(...
        true,doval,dosc,dofitalpha0,dofitalphas,...
        iter+first_iter,telapsed,...
        [],ppl,ppl_validation,alpha0,0);
end

oldNeglogL = neglogL;

while relDeltaObjective>reltol && iter<maxiter
    tstart = tic;
    
    iter = iter + 1;

    if     dosc
        eta = 1/(1 + (first_iter + iter)).^decay;

        [Nwk,Phi] = clustering.ldautils.learn_phi_at_fixed_theta(...
            solver,alphas,beta0,Phi,Nwk,Theta,Nkj,X,curridx,eta);

        [nbegin,nend] = scrange(nend,batchsize,N);
        curridx = scidx(nbegin:nend);
        
        curr_num_words = sum(nj(curridx));
        if curr_num_words==0
            continue;
        end

        % capturing all output arguments from the built-in function to avoid memory leak issues
        [Theta(:,curridx), ~] = ...
            clustering.ldautils.learn_theta_at_fixed_phi(alphas,Phi,X(:,curridx),reltol,2);
        Nkj(:,curridx) = Theta(:,curridx).*(nj(curridx) + alpha0) - alphas;
        Nkj(Nkj<0) = 0;
        
    elseif strcmp(solver,'cgs')
        [Nwk,Nkj,Phi,Theta,z] = clustering.ldautils.sampling_step(...
            solver,alphas,beta0,Nwk,Nkj,rowidx,colidx,z,20,rsh);
        
    elseif usemex
        [Nwk,Nkj,Phi,Theta] = clustering.ldautils.em_step(...
            solver,alphas,beta0,Phi,Nwk,Theta,Nkj,X,16);
        
    else
        [Nwk,Nkj,Phi,Theta] = one_em_step(solver,alphas,beta0,Phi,Nwk,Theta,Nkj,X);
    
    end
    
    % Compute likelihood
    if isempty(curridx)
        if usemex
            logprob = clustering.ldautils.docprob_atmaxlkh(X,Phi,Theta);
        else
            logprob = docprob_atmaxlkh(X,Phi,Theta);
        end
        
        neglogL = -nansum(logprob);
        
    else % usemex is always true
        logprob = clustering.ldautils.docprob_atmaxlkh(...
            X(:,curridx),Phi,Theta(:,curridx));
        
        neglogL = -nansum(logprob)*num_words/curr_num_words;
    end        
    
    deltaObj = neglogL - oldNeglogL;
    oldNeglogL = neglogL;
    relDeltaObjective = abs(deltaObj/neglogL);
    
    ppl = exp( neglogL/num_words );
    
    neglogL_validation = [];
    ppl_validation = [];
    if doval && mod(iter+first_iter,freqval)==0
        [neglogL_validation,ppl_validation] = ...
            evaluate_model(usemex,alphas,Phi,Xval,...
            reltol_val,maxiter_val,num_words_validation);
    end
        
    % Fit alphas
    nalpha_iter = 0;
    if dofitalphas && iter+first_iter>10
        old_alpha0 = alpha0;
        old_alphas = alphas;        
        
        relAlphas = Inf;
        
        if ~dofitalpha0
            denom = sum( psi(nj+alpha0) - psi(alpha0) );
        end
        
        while relAlphas > 1e-3
            bad = sum(Nkj,2)==0 | alphas==0;
            if dofitalpha0
                denom = sum( psi(nj+alpha0) - psi(alpha0) );
            end
            alphas = alphas.*sum( psi(Nkj+alphas) - psi(alphas), 2 )./denom;
            alphas(bad) = 0;
            
            if dofitalpha0
                alpha0 = sum(alphas);
            else
                alphas = alphas/sum(alphas)*alpha0;
            end
            
            relAlphas = sum(abs(alphas - old_alphas))/old_alpha0;
            old_alphas = alphas;
            old_alpha0 = alpha0;
            
            nalpha_iter = nalpha_iter + 1;
        end
    end        
    
    telapsed = toc(tstart);
    
    exitflag = 0; % number of iterations exceeded
    if relDeltaObjective <= reltol
        exitflag = 1; % converged due to tolerance on loglik
    end
    % exitflag = 2 reserved for tensor decomposition
    
    numiter = iter + first_iter;

    if dohist
        neglogL_vals(iter)    = neglogL;
        ppl_vals(iter)        = ppl;
        alpha_vals(iter)      = alpha0;
        nalphaiter_vals(iter) = nalpha_iter;
        telapsed_vals(iter)   = telapsed;
        
        if doval && mod(iter+first_iter,freqval)==0
            neglogL_validation_vals(iter) = neglogL_validation;
            ppl_validation_vals(iter)     = ppl_validation;
        end
    end
    
    if verbose>0
        doheader = iter==0 || (iter>1 && mod(iter,20)==1);
        fitinfolen = clustering.ldautils.showfitinfo(...
            doheader,doval,dosc,dofitalpha0,dofitalphas,...
            iter+first_iter,telapsed,...
            relDeltaObjective,ppl,ppl_validation,alpha0,nalpha_iter);
    end
end

if verbose>0
    clustering.ldautils.showseparator(fitinfolen);
end

Theta_sc = [];
if dosc
    % Theta's computed for the entire training set.
    % This includes recomputing Theta's for documents
    % that have been processed stochastically.
    
    Theta_sc = Theta;
    
    if verbose > 0
        [Theta,logprob] = ...
            sc_final_pass_verbose(alphas,Phi,X,reltol_sc,maxiter_sc);
    else
        [Theta,logprob] = ...
            clustering.ldautils.learn_theta_at_fixed_phi(...
            alphas,Phi,X,reltol_sc,maxiter_sc);
    end

    neglogL = -nansum(logprob);
end

if dosc
    solver = ['s' solver];
end

fitinfo = struct;
fitinfo.TerminationCode        = exitflag;
fitinfo.TerminationStatus      = clustering.ldautils.code2str(exitflag);
fitinfo.NumIterations          = numiter;
fitinfo.NegativeLogLikelihood  = neglogL;
fitinfo.Perplexity             = exp( neglogL/full(sum(X(:))) );
fitinfo.Solver                 = string(solver);

history = [];
if dohist
    history = struct;
    history.NegativeLogLikelihood              = neglogL_vals(1:iter);
    history.Perplexity                         = ppl_vals(1:iter);
    history.TopicConcentration                 = alpha_vals(1:iter);
    history.TopicConcentrationIterations       = nalphaiter_vals(1:iter);
    history.TimeSinceStart                     = cumsum(telapsed_vals(1:iter));
    
    history.ValidationNegativeLogLikelihood    = neglogL_validation_vals(1:iter);
    history.ValidationPerplexity               = ppl_validation_vals(1:iter);
end
fitinfo.History                = history;

if ~isempty(initial_history)
    old = initial_history;

    history.NegativeLogLikelihood = ...
        [old.NegativeLogLikelihood;
        history.NegativeLogLikelihood];
    history.Perplexity = ...
        [old.Perplexity; history.Perplexity];
    history.TopicConcentration = ...
        [old.TopicConcentration; history.TopicConcentration];
    history.TopicConcentrationIterations = ...
        [old.TopicConcentrationIterations; history.TopicConcentrationIterations];
    history.TimeSinceStart = ...
        [old.TimeSinceStart; old.TimeSinceStart(end)+history.TimeSinceStart];
    
    history.ValidationNegativeLogLikelihood = ...
        [old.ValidationNegativeLogLikelihood;
        history.ValidationNegativeLogLikelihood];
    history.ValidationPerplexity = ...
        [old.ValidationPerplexity; history.ValidationPerplexity];
    
    fitinfo.History = history;
end

if dosc
    scinfo.LastProcessedDocumentIndex = nbegin - 1;
end

end


function logprob = docprob_atmaxlkh(X,Phi,Theta)

N = size(X,2);

logprob = NaN(N,1,'like',Phi);

parfor n=1:N
    x = X(:,n);
    if all(x==0)
        continue;
    end
    pw_given_d = Phi*Theta(:,n);
    xlogp = x.*log(pw_given_d);
    pwd0 = pw_given_d==0;
    xeq0 = x == 0;
    xlogp(xeq0 & pwd0) = 0;
    logprob(n) = sum(xlogp);
end

end


function [Nwk,Nkj,Phi,Theta] = one_em_step(solver,alphas,beta0,Phi,Nwk,Theta,Nkj,X)

N = size(X,2);
K = size(Phi,2);

%alpha0 = sum(alphas);

spmd
    idx = chunk(spmdSize,spmdIndex,N);
    
    local_Nwk = zeros(size(Nwk),'like',Nwk);
    local_Nkj = zeros(K,numel(idx),'like',Nkj);
    
    for j=1:numel(idx)
        n = idx(j);
        x = X(:,n);
        
        nj = full(sum(x));
        if nj==0
            continue;
        end
        
        % E-step: T-by-K array of P(z|d,w) for given document
        gamma_wjk = estep(solver,alphas,beta0,Phi,Nwk,Theta(:,n),Nkj(:,n),x);
        
        % M-step
        gamma_wjk = gamma_wjk.*x;
        
        sumGamma_wjk = sum(gamma_wjk(:));
        if sumGamma_wjk==0
            continue;
        end
        
        local_Nwk      = local_Nwk + gamma_wjk;
        local_Nkj(:,j) = sum(gamma_wjk,1)';
    end
    
    % Sum over workers
    local_Nwk = spmdPlus(local_Nwk,1);
end

%
% Collect from spmd
%

Nwk = local_Nwk{1};

B = numel(local_Nkj);
for b=1:B
    idx = chunk(B,b,N);
    Nkj(:,idx) = local_Nkj{b};
end
%end

% Compute probabilities

Phi = clustering.ldautils.nwk2phi(Nwk,beta0);

% nk = sum(Nwk,1);
% Phi = (Nwk + beta0/T)./(nk + beta0);
% if beta0==0
%     bad = nk==0;
%     Phi(:,bad) = 0;
% end

Theta = clustering.ldautils.nkj2theta(Nkj,alphas);

% nj = sum(Nkj,1);
% Theta = (Nkj + alphas)./(nj + alpha0);
% if alpha0==0
%     bad = nj==0;
%     Theta(:,bad) = 0;
% end

end


function idx = chunk(B,b,N)
n = floor(N/B);
ibegin = (b-1)*n+1;
if b < B
    iend = ibegin + n-1;
else
    iend = N;
end
idx = ibegin:iend;
end


function gamma_wjk = estep(type,alphas,beta0,Phi,Nwk,theta,nkj,x)
% Phi - W-by-K array of P(w|z)
% Nwk - W-by-K array of counts
% theta - K-by-1 array of P(z|d)
% nkj - K-by-1 array of counts
% gamma_wjk = W-by-K array of P(z|d,w), i.e., sum(gamma_wjk,2)==ones(T,1)

W = size(Phi,1);
nj = full(sum(x)); % cannot be 0; checked in one_em_step

switch type
    case 'plsa' % Hoffman's plsa
        pwd = Phi*theta;
        Pzdw = (Phi.*theta') ./ pwd;
        bad = pwd == 0;
        Pzdw(bad,:) = 0;
        gamma_wjk = Pzdw;
        
%     case 'evb' % exact variational Bayes
%         gamma_wjk = exp(psi(Nwk + beta0/W) - psi(sum(Nwk,1) + beta0) + psi(nkj' + alphas'));
%         gamma_wjk(gamma_wjk<0) = 0;
%         gamma_wjk2 = sum(gamma_wjk,2);
%         gamma_wjk = gamma_wjk./gamma_wjk2;
%         bad = gamma_wjk2==0;
%         gamma_wjk(bad,:) = 0;
        
    case 'avb' % approximate variational Bayes
        gamma_wjk = (Nwk + beta0/W - 0.5)./(sum(Nwk,1) + beta0 - 0.5).*(nkj' + alphas' - 0.5);
        gamma_wjk(gamma_wjk<0) = 0;
        gamma_wjk2 = sum(gamma_wjk,2);
        gamma_wjk = gamma_wjk./gamma_wjk2;
        bad = gamma_wjk2==0;
        gamma_wjk(bad,:) = 0;
        
    case 'cvb0' % collapsed variational Bayes
        pwd = Phi*theta;
        Pzdw = (Phi.*theta') ./ pwd;
        bad = pwd == 0;
        Pzdw(bad,:) = 0;
        
        Nkj_ij = nkj' - nj/full(sum(x>0))*Pzdw;
        Nkj_ij(Nkj_ij<0) = 0;
        
        Nwk_ij = Nwk - nj*(Phi.*theta');
        Nwk_ij(Nwk_ij<0) = 0;
        nk_ij = sum(Nwk_ij,1);
        
        gamma_wjk = ( Nwk_ij + beta0/W )./( nk_ij + beta0 ).*( Nkj_ij + alphas' );
        gamma_wjk = gamma_wjk./sum(gamma_wjk,2);
        
    otherwise
        error(message('stats:clustering:ldautils:UnknownSolver'));
end

end


function [nbegin,nend] = scrange(nend,batchsize,N)
% Range of indices for stochastic learning

if nend==0
    nbegin = 1;
    nend   = min(batchsize,N);
else
    nbegin = nend + 1;
    if nbegin >= N
        nbegin = 1;
    end
    nend = min( N, nbegin + batchsize - 1 );
end

end


function [Theta,logprob] = sc_final_pass_verbose(alphas,Phi,X,reltol,maxiter)
% Perform final pass (fit Theta keeping Phi fixed) after stochastic solver
% chunk by chunk to monitor progress

K = size(Phi,2);
N = size(X,2);

fprintf('%s\n',...
    message('stats:clustering:ldautils:solve_em:EstimatingTopicProbabilitiesPerDocument').getString());

Theta = NaN(K,N);
logprob = NaN(N,1);

chunksize = 10000;
ibegin = 1;
iend   = min(N,chunksize);

fprintf('==========================================\n');
fprintf('|   Number of processed   | Elapsed time |\n');
fprintf('| documents (%% completed) |  (HH:MM:SS)  |\n'); % 1 off due to %%
fprintf('==========================================\n');

tstart = tic;

while ibegin <= N
    idx = ibegin:iend;
    
    [Theta(:,idx),logprob(idx)] = ...
        clustering.ldautils.learn_theta_at_fixed_phi(...
        alphas,Phi,X(:,idx),reltol,maxiter);
    telapsed = toc(tstart);
    
    percComplete = iend/N*100;
    stelapsed = datestr(telapsed/(24*3600),'HH:MM:SS');
    
    fprintf('| %13i (%6.2f%%) | %12s |\n',iend,percComplete,stelapsed);
    
    ibegin = iend + 1;
    iend   = min(N,iend + chunksize);
end

fprintf('==========================================\n');
fprintf('%s\n',...
    message('stats:clustering:ldautils:solve_em:DoneEstimatingTopicProbabilitiesPerDocument',...
    sprintf('%10.2f',telapsed)).getString());
end


function [neglogL,ppl] = ...
    evaluate_model(usemex,alphas,Phi,Xval,reltol,maxiter,num_words)
if usemex
    [~,logprob] = clustering.ldautils.learn_theta_at_fixed_phi(...
        alphas,Phi,Xval,reltol,maxiter);
else
    [~,logprob] = clustering.ldautils.learnThetaAtFixedPhi(...
        alphas,Xval,Phi,reltol,maxiter);
end

neglogL = -nansum(logprob);
ppl = exp( neglogL/num_words );
end

% LocalWords:  dofitalpha dofitalphas Xval Nval Nkj Nwk dohist dosc scinfo freqval rowidx colidx sc
% LocalWords:  nk nj nkj wjk evb
