//    Copyright 2011 Brad Block, Pawjaw, LLC. (brad@pawjaw.com)
//
//    This program is free software: you can redistribute it and/or modify
//    it under the terms of the GNU General Public License as published by
//    the Free Software Foundation, either version 3 of the License, or
//    (at your option) any later version.
//
//    This program is distributed in the hope that it will be useful,
//    but WITHOUT ANY WARRANTY; without even the implied warranty of
//    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//    GNU General Public License for more details.
//
//    You should have received a copy of the GNU General Public License
//    along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
//    This program implements concepts described in the paper entitled
//    "Collapsed Variational Inference for Hierarchical Dirichlet Processes"
//    written by Yee Whye Teh, Kenichi Kurihara, and Max Welling.

package com.pawjaw.topics;

import com.pawjaw.linalg.Matrix;
import com.pawjaw.math.SpecFun;
import com.pawjaw.utils.ProgressCallback;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.Semaphore;

public class CVHDP {
    // output
    public double E_alpha, G_alpha;
    public double E_beta, G_beta;
    public double E_gamma;
    public double[] E_pi_k, G_pi_k; // [topic]
    public double[] E_tau_w, G_tau_w; // [word]
    public long real_variational_posterior_update_millis, real_bookkeeping_millis, real_hyperparameter_update_millis;
    public String timing;
    public List<Double> avglls = new ArrayList(); // record of average log likelihood of held out training data during inference
    public List<double[]> E_n_ks = new ArrayList(); // record of expected number of tokens assigned to topics during inference
    public double[] E_n_k; // [topic] expected number of tokens assigned to topic
    public Integer[] dk_k; // [topic] map topics in descending order of E_n_k
    public double[][] theta_dk; // [document][topic] expected document-topic multinomial mixture
    public double[][] phi_kw; // [topic][word] expected topic-word multinomial mixture

    // constructors, all work is done during construction
    public CVHDP(Matrix word_count_documents_in_columns, int topics, ProgressCallback notifier) throws InterruptedException {
        this(word_count_documents_in_columns, null, topics, 1,
                DEFAULT_VARIATIONAL_ITERATIONS, DEFAULT_GROUP_TOKENS_BY_WORD, 0,
                DEFAULT_E_BETA, DEFAULT_G_BETA,
                DEFAULT_A_ALPHA, DEFAULT_B_ALPHA,
                DEFAULT_A_BETA, DEFAULT_B_BETA,
                DEFAULT_A_GAMMA, DEFAULT_B_GAMMA,
                DEFAULT_A_TAU / (double)word_count_documents_in_columns.rows, notifier);
    }

    public CVHDP(Matrix word_count_documents_in_columns_training, Matrix word_count_documents_in_columns_heldout, int topics,
            int threads, int variational_iterations, boolean group_tokens_by_word, int heldout_log_likelihood_estimates,
            ProgressCallback notifier) throws InterruptedException {
        this(word_count_documents_in_columns_training, word_count_documents_in_columns_heldout, topics,
                threads, variational_iterations, group_tokens_by_word, heldout_log_likelihood_estimates,
                DEFAULT_E_BETA, DEFAULT_G_BETA,
                DEFAULT_A_ALPHA, DEFAULT_B_ALPHA,
                DEFAULT_A_BETA, DEFAULT_B_BETA,
                DEFAULT_A_GAMMA, DEFAULT_B_GAMMA,
                DEFAULT_A_TAU / (double)word_count_documents_in_columns_training.rows, notifier);
    }

    public CVHDP(Matrix word_count_documents_in_columns_training, Matrix word_count_documents_in_columns_heldout, int topics,
            int threads, int variational_iterations, boolean group_tokens_by_word, int heldout_log_likelihood_estimates,
            double E_beta, double G_beta, double a_alpha, double b_alpha, double a_beta, double b_beta,
            double a_gamma, double b_gamma, double a_tau, ProgressCallback notifier) throws InterruptedException {
        this.K = topics > word_count_documents_in_columns_training.rows ? word_count_documents_in_columns_training.rows : topics;
        this.ho_wd = word_count_documents_in_columns_heldout;
        this.threads = threads < 1 ? 1 : threads;
        this.group_tokens_by_word = group_tokens_by_word;
        this.E_beta = E_beta;
        this.G_beta = G_beta;
        this.a_alpha = a_alpha;
        this.b_alpha = b_alpha;
        this.a_beta = a_beta;
        this.b_beta = b_beta;
        this.a_gamma = a_gamma;
        this.b_gamma = b_gamma;
        this.a_tau = a_tau;
        initialize(word_count_documents_in_columns_training);
        run(variational_iterations, heldout_log_likelihood_estimates, notifier);
    }

    // defaults
    public static final double DEFAULT_E_BETA = 100;
    public static final double DEFAULT_G_BETA = 100;
    public static final double DEFAULT_A_ALPHA = 3.0; // [2,6]
    public static final double DEFAULT_B_ALPHA = 5.0; // [2,6]
    public static final double DEFAULT_A_BETA = 3.0; // [2,6]
    public static final double DEFAULT_B_BETA = 5.0; // [2,6]
    public static final double DEFAULT_A_GAMMA = 6.0; // [5,10]
    public static final double DEFAULT_B_GAMMA = 8.0; // [5,10]
    public static final double DEFAULT_A_TAU = 40; // this default gets divided by W, e.g. [30,50]/W
    public static final boolean DEFAULT_GROUP_TOKENS_BY_WORD = true; // calculate one variational pr for each word not token
    public static final int DEFAULT_VARIATIONAL_ITERATIONS = 200; // steps equals em iterations multiplied by topics
    private static final double SMALL_VALUE_THRESHOLD = 1e-8; // to keep various operations from diverging

    /************************ private section ************************/
    // expectations of quantities
    private double[][] E_n_dk, V_n_dk, Z_n_dk; // [document][topic] sum over token
    private double[][] E_n_kw, V_n_kw, Z_n_kw; // [topic][global_word] sum over document
    private double[] V_n_k, Z_n_k; // [topic] = sum over document and token

    // temporary storage per document
    private double[][] tE_n_dk, tV_n_dk, tZ_n_dk; // [document][topic]
    private double[][][] tE_n_kw, tV_n_kw, tZ_n_kw; // [thread][topic][document_word]

    // calculating probabilities for each document token
    private boolean group_tokens_by_word; // calculate only one topic assignment probability for each unique word in document
    private double[][][] q_z_id_eq_k_dki; // [document][topic][document_token] variational topic assignment probabilities

    // auxillary variables
    private double E_s, digamma_E_beta, E_log_one_minus_pi_tilde;
    private double E_log_eta, E_log_xi, E_t;
    private double[] E_s_k;
    private double[] E_s_gt_k;
    private double[] E_t_w;
    private double[] G_alpha_x_pi_k, G_beta_x_tau_w;
    private double[] digamma_G_alpha_x_pi_k, digamma_G_beta_x_tau_w;
    private double[] E_pi_tilde_k, E_one_minus_pi_tilde_k;
    private double[] G_pi_tilde_k, G_one_minus_pi_tilde_k;
    private double[][] E_s_dk;

    // variational approximation parameters
    private int K; // topic truncation threshold (number of potential assigned topics)
    private double a_alpha, b_alpha;
    private double a_beta, b_beta;
    private double a_gamma, b_gamma;
    private double a_tau;

    // documents, words, tokens
    private int D; // number of documents in corpus
    private int W; // number of words in vocabulary
    private int N_d[]; // [document] = number of document_tokens in document d
    private Matrix ho_wd = null; // held-out words from documents for calculating average log-likelihood
    private int[][] docTokenToWord; // [document][document_token] = global word index
    private int[][][] docWordToken; // [document][document_word] = { global word, first doc_token index, doc_token count for word }

    // execution options
    private int threads;
    private Thread[] thread_pool;
    private Semaphore parent_lock, thread_lock;

    private void initialize(Matrix word_count_documents_in_columns) {
        initializeDocuments(word_count_documents_in_columns);
        initializeParameters();
        initializeTopicAssignments();
        initializeTopicAssignmentSums();
    }

    private void run(int iterations, int llestimates, ProgressCallback notifier) throws InterruptedException {
        int iteration = 0, avglliter = llestimates <= 1 ? 0 : (int)Math.max(iterations / llestimates, 1);
        startThreads();
        real_variational_posterior_update_millis = real_bookkeeping_millis = real_hyperparameter_update_millis = 0;
        while(iteration++ < iterations) {
            updateVariationalPosterior();
            updateHyperParameters();
            if(avglliter > 0 && iteration % avglliter == 0) {
                estimateMixtures();
                avglls.add(avgLogLikelihood());
                E_n_ks.add(Arrays.copyOf(E_n_k, K));
            }
            if(notifier != null)
                notifier.notify(iteration, iterations);
        }
        stopThreads();
        estimateMixtures();
        if(llestimates == 1) {
            avglls.add(avgLogLikelihood());
            E_n_ks.add(Arrays.copyOf(E_n_k, K));
        }
        describeTiming();
    }

    private void updateVariationalPosterior() throws InterruptedException {
        long start = System.currentTimeMillis();
        thread_lock.release(threads);
        parent_lock.acquire(threads);
        real_variational_posterior_update_millis += System.currentTimeMillis() - start;
        start = System.currentTimeMillis();
        updateTopicAssignmentSums();
        real_bookkeeping_millis += System.currentTimeMillis() - start;
    }

    private void updateHyperParameters() {
        long start = System.currentTimeMillis();
        updateAlpha();
        updateBeta();
        updatePi();
        updateGamma();
        updateTau();
        real_hyperparameter_update_millis += System.currentTimeMillis() - start;
    }

    private void initializeDocuments(Matrix document_in_columns) {
        int d, x, X, y, Y, c, w;
        D = document_in_columns.cols;
        W = document_in_columns.rows;
        docWordToken = new int[D][][];
        docTokenToWord = new int[D][];
        N_d = new int[D];
        for(d = 0;d < D;d++) {
            x = document_in_columns.pointr[d];
            X = document_in_columns.pointr[d+1];
            docWordToken[d] = new int[X - x][3];
            for(y = 0;x < X;x++) {
                docWordToken[d][y][0] = document_in_columns.rowind[x];
                docWordToken[d][y][1] = N_d[d];
                N_d[d] += docWordToken[d][y++][2] = (int)Math.ceil(document_in_columns.value[x]);
            }
            docTokenToWord[d] = new int[N_d[d]];
            c = 0; Y = docWordToken[d].length;
            for(y = 0;y < Y;y++) {
                w = docWordToken[d][y][0];
                X = docWordToken[d][y][2];
                for(x = 0;x < X;x++)
                    docTokenToWord[d][c++] = w;
            }
        }
    }

    private void initializeTopicAssignments() {
        int d, k, i, I;
        double sum;
        Random r = new Random();
        q_z_id_eq_k_dki = new double[D][K][];
        for(d = 0;d < D;d++) {
            I = N_d[d];
            for(k = 0;k < K;k++)
                q_z_id_eq_k_dki[d][k] = new double[I];
            for(i = 0;i < I;i++) {
                sum = 0;
                for(k = 0;k < K;k++)
                    sum += q_z_id_eq_k_dki[d][k][i] = 1.0 + r.nextDouble();
                for(k = 0;k < K;k++)
                    q_z_id_eq_k_dki[d][k][i] /= sum;
            }
        }
    }

    private void initializeTopicAssignmentSums() {
        int t, d, k, i, I, w;
        double E_q_z_idk, V_q_z_idk, Z_q_z_idk;
        E_n_dk = new double[D][K];
        V_n_dk = new double[D][K];
        Z_n_dk = new double[D][K];
        E_n_kw = new double[K][W];
        V_n_kw = new double[K][W];
        Z_n_kw = new double[K][W];
        E_n_k = new double[K];
        V_n_k = new double[K];
        Z_n_k = new double[K];
        tE_n_dk = new double[D][K];
        tV_n_dk = new double[D][K];
        tZ_n_dk = new double[D][K];
        tE_n_kw = new double[threads][][];
        tV_n_kw = new double[threads][][];
        tZ_n_kw = new double[threads][][];
        for(t = 0;t < threads;t++) {
            tE_n_kw[t] = new double[K][W];
            tV_n_kw[t] = new double[K][W];
            tZ_n_kw[t] = new double[K][W];
        }
        for(d = 0;d < D;d++) {
            for(k = 0;k < K;k++) {
                I = N_d[d];
                for(i = 0;i < I;i++) {
                    E_q_z_idk = q_z_id_eq_k_dki[d][k][i];
                    V_q_z_idk = E_q_z_idk * (1.0 - E_q_z_idk);
                    Z_q_z_idk = Math.log(1.0 - E_q_z_idk);
                    w = docTokenToWord[d][i];
                    E_n_dk[d][k] += E_q_z_idk;
                    V_n_dk[d][k] += V_q_z_idk;
                    Z_n_dk[d][k] += Z_q_z_idk;
                    E_n_kw[k][w] += E_q_z_idk;
                    V_n_kw[k][w] += V_q_z_idk;
                    Z_n_kw[k][w] += Z_q_z_idk;
                    E_n_k[k] += E_q_z_idk;
                    V_n_k[k] += V_q_z_idk;
                    Z_n_k[k] += Z_q_z_idk;
                }
            }
        }
    }

    private void initializeParameters() {
        int k;
        G_tau_w = new double[W];
        E_tau_w = new double[W];
        E_s_dk = new double[D][K];
        E_s_k = new double[K];
        E_s_gt_k = new double[K];
        Arrays.fill(G_tau_w, 1.0 / (double)W);
        E_alpha = a_alpha / b_alpha;
        G_alpha = Math.exp(SpecFun.digamma(a_alpha)) / b_alpha;
        E_gamma = a_gamma / b_gamma;
        G_pi_k = new double[K];
        Arrays.fill(G_pi_k, 1.0 / (double)K);
        E_pi_k = new double[K];
        E_pi_tilde_k = new double[K];
        E_one_minus_pi_tilde_k = new double[K];
        G_pi_tilde_k = new double[K];
        G_one_minus_pi_tilde_k = new double[K];
        E_t_w = new double[W];
        dk_k = new Integer[K];
        for(k = 0;k < K;k++)
            dk_k[k] = k;
        G_alpha_x_pi_k = new double[K];
        G_beta_x_tau_w = new double[W];
        digamma_G_alpha_x_pi_k = new double[K];
        digamma_G_beta_x_tau_w = new double[W];
    }

    private void startThreads() {
        int t, c, d = 0;
        int[] ds;
        thread_pool = new Thread[threads];
        thread_lock = new Semaphore(0, true);
        parent_lock = new Semaphore(0, true);
        for(t = 0;t < threads;t++) {
            ds = new int[D / threads + (D % threads > t ? 1 : 0)];
            for(c = 0;c < ds.length;c++)
                ds[c] = d++;
            (thread_pool[t] = new TopicAssignmentThread(t, ds)).start();
        }
    }

    private void stopThreads() {
        int t;
        for(t = 0;t < threads;t++)
            thread_pool[t].interrupt();
    }

    private class TopicAssignmentThread extends Thread {
        private int tid;
        private int[] ds;
        public TopicAssignmentThread(int tid, int[] ds) {
            this.tid = tid;
            this.ds = ds;
        }
        public void run() {
            int x, k;
            try {
                while(true) {
                    thread_lock.acquire();
                    for(k = 0;k < K;k++) {
                        Arrays.fill(tE_n_kw[tid][k], 0);
                        Arrays.fill(tV_n_kw[tid][k], 0);
                        Arrays.fill(tZ_n_kw[tid][k], 0);
                    }
                    for(x = 0;x < ds.length;x++)
                        updateTopicAssignments(tid, ds[x]);
                    parent_lock.release();
                }
            } catch (InterruptedException e) {}
        }
    }

    private void updateTopicAssignments(int t, int d) {
        int l, L, w, i, I, k;
        double sum, n, E_q_z_idk, V_q_z_idk, Z_q_z_idk, term1, term2, term3, term4, term5, term6;
        Arrays.fill(tE_n_dk[d], 0);
        Arrays.fill(tV_n_dk[d], 0);
        Arrays.fill(tZ_n_dk[d], 0);
        L = docWordToken[d].length;
        for(l = 0;l < L;l++) {
            w = docWordToken[d][l][0];
            if(group_tokens_by_word) {
                i = docWordToken[d][l][1];
                n = docWordToken[d][l][2];
                sum = 0;
                for(k = 0;k < K;k++) {
                    E_q_z_idk = q_z_id_eq_k_dki[d][k][i];
                    V_q_z_idk = E_q_z_idk * (1.0 - E_q_z_idk);
                    term1 = G_alpha * G_pi_k[k] + E_n_dk[d][k] - E_q_z_idk;
                    term2 = G_beta * G_tau_w[w] + E_n_kw[k][w] - E_q_z_idk;
                    term3 = E_beta + E_n_k[k] - E_q_z_idk;
                    term4 = (V_n_dk[d][k] - V_q_z_idk) / term1 / term1;
                    term5 = (V_n_kw[k][w] - V_q_z_idk) / term2 / term2;
                    term6 = (V_n_k[k] - V_q_z_idk) / term3 / term3;
                    sum += q_z_id_eq_k_dki[d][k][i] = term1 * term2 / term3 * Math.exp((- term4 - term5 + term6) / 2.0);
                }
                for(k = 0;k < K;k++) {
                    E_q_z_idk = (q_z_id_eq_k_dki[d][k][i] /= sum);
                    V_q_z_idk = E_q_z_idk * (1.0 - E_q_z_idk);
                    Z_q_z_idk = Math.log(1.0 - E_q_z_idk);
                    tE_n_dk[d][k] += E_q_z_idk * n;
                    tV_n_dk[d][k] += V_q_z_idk * n;
                    tZ_n_dk[d][k] += Z_q_z_idk * n;
                    tE_n_kw[t][k][w] += E_q_z_idk * n;
                    tV_n_kw[t][k][w] += V_q_z_idk * n;
                    tZ_n_kw[t][k][w] += Z_q_z_idk * n;
                }
            } else {
                I = (i = docWordToken[d][l][1]) + docWordToken[d][l][2];
                while(i < I) {
                    sum = 0;
                    for(k = 0;k < K;k++) {
                        E_q_z_idk = q_z_id_eq_k_dki[d][k][i];
                        V_q_z_idk = E_q_z_idk * (1.0 - E_q_z_idk);
                        term1 = G_alpha * G_pi_k[k] + E_n_dk[d][k] - E_q_z_idk;
                        term2 = G_beta * G_tau_w[w] + E_n_kw[k][w] - E_q_z_idk;
                        term3 = E_beta + E_n_k[k] - E_q_z_idk;
                        term4 = (V_n_dk[d][k] - V_q_z_idk) / term1 / term1;
                        term5 = (V_n_kw[k][w] - V_q_z_idk) / term2 / term2;
                        term6 = (V_n_k[k] - V_q_z_idk) / term3 / term3;
                        sum += q_z_id_eq_k_dki[d][k][i] = term1 * term2 / term3 * Math.exp((- term4 - term5 + term6) / 2.0);
                    }
                    for(k = 0;k < K;k++) {
                        E_q_z_idk = q_z_id_eq_k_dki[d][k][i] /= sum;
                        V_q_z_idk = E_q_z_idk * (1.0 - E_q_z_idk);
                        Z_q_z_idk = Math.log(1.0 - E_q_z_idk);
                        tE_n_dk[d][k] += E_q_z_idk;
                        tV_n_dk[d][k] += V_q_z_idk;
                        tZ_n_dk[d][k] += Z_q_z_idk;
                        tE_n_kw[t][k][w] += E_q_z_idk;
                        tV_n_kw[t][k][w] += V_q_z_idk;
                        tZ_n_kw[t][k][w] += Z_q_z_idk;
                    }
                    i++;
                }
            }
        }
    }

    private void updateTopicAssignmentSums() {
        int d, k, t, w;
        for(k = 0;k < K;k++) {
            Arrays.fill(E_n_kw[k], 0);
            Arrays.fill(V_n_kw[k], 0);
            Arrays.fill(Z_n_kw[k], 0);
        }
        Arrays.fill(E_n_k, 0);
        Arrays.fill(V_n_k, 0);
        Arrays.fill(Z_n_k, 0);
        for(d = 0;d < D;d++) {
            for(k = 0;k < K;k++) {
                E_n_dk[d][k] = tE_n_dk[d][k];
                V_n_dk[d][k] = tV_n_dk[d][k];
                Z_n_dk[d][k] = tZ_n_dk[d][k];
                E_n_k[k] += tE_n_dk[d][k];
                V_n_k[k] += tV_n_dk[d][k];
                Z_n_k[k] += tZ_n_dk[d][k];
            }
        }
        for(t = 0;t < threads;t++) {
            for(k = 0;k < K;k++) {
                for(w = 0;w < W;w++) {
                    E_n_kw[k][w] += tE_n_kw[t][k][w];
                    V_n_kw[k][w] += tV_n_kw[t][k][w];
                    Z_n_kw[k][w] += tZ_n_kw[t][k][w];
                }
            }
        }
    }

    private double E_s_dk(int d, int k) { // d = document, k = topic
        double Pplus_n_dk, Eplus_n_dk, Vplus_n_dk, psi1, psi3;
        if(E_n_dk[d][k] < SMALL_VALUE_THRESHOLD)
            return 0;
        Pplus_n_dk = 1.0 - Math.exp(Z_n_dk[d][k]);
        Eplus_n_dk = E_n_dk[d][k] / Pplus_n_dk;
        Vplus_n_dk = V_n_dk[d][k] / Pplus_n_dk - Math.exp(Z_n_dk[d][k]) * Eplus_n_dk;
        psi1 = SpecFun.digamma(G_alpha_x_pi_k[k] + Eplus_n_dk);
        psi3 = SpecFun.tetragamma(G_alpha_x_pi_k[k] + Eplus_n_dk);
        return G_alpha_x_pi_k[k] * Pplus_n_dk * (psi1 - digamma_G_alpha_x_pi_k[k] + Vplus_n_dk * psi3 / 2.0);
    }

    private double E_t_kw(int k, int w) { // k = topic, w = word
        double Pplus_n_kw, Eplus_n_kw, Vplus_n_kw, psi1, psi3;
        if(E_n_kw[k][w] < SMALL_VALUE_THRESHOLD)
            return 0;
        Pplus_n_kw = 1.0 - Math.exp(Z_n_kw[k][w]);
        Eplus_n_kw = E_n_kw[k][w] / Pplus_n_kw;
        Vplus_n_kw = V_n_kw[k][w] / Pplus_n_kw - Math.exp(Z_n_kw[k][w]) * Eplus_n_kw;
        psi1 = SpecFun.digamma(G_beta_x_tau_w[w] + Eplus_n_kw);
        psi3 = SpecFun.tetragamma(G_beta_x_tau_w[w] + Eplus_n_kw);
        return G_beta_x_tau_w[w] * Pplus_n_kw * (psi1 - digamma_G_beta_x_tau_w[w] + Vplus_n_kw * psi3 / 2.0);
    }

    private double E_log_xi_k(int k) { // k = topic
        double Pplus_n_k, Eplus_n_k, Vplus_n_k, psi2, psi3;
        if(E_n_k[k] < SMALL_VALUE_THRESHOLD)
            return 0;
        Pplus_n_k = 1.0 - Math.exp(Z_n_k[k]);
        Eplus_n_k = E_n_k[k] / Pplus_n_k;
        Vplus_n_k = V_n_k[k] / Pplus_n_k - Math.exp(Z_n_k[k]) * Eplus_n_k;
        psi2 = SpecFun.digamma(E_beta + Eplus_n_k);
        psi3 = SpecFun.tetragamma(E_beta + Eplus_n_k);
        return Pplus_n_k * (digamma_E_beta - psi2 - Vplus_n_k * psi3 / 2.0);
    }

    private double E_log_eta_d(int d) {
        return SpecFun.digamma(E_alpha) - SpecFun.digamma(E_alpha + N_d[d]);
    }

    // update order 1
    private void updateAlpha() {
        int d, k;
        double a, b;
        E_s = 0;
        E_log_eta = 0;
        for(k = 0;k < K;k++)
            digamma_G_alpha_x_pi_k[k] = SpecFun.digamma(G_alpha_x_pi_k[k] = G_alpha * G_pi_k[k]);
        for(d = 0;d < D;d++) {
            for(k = 0;k < K;k++)
                E_s += E_s_dk[d][k] = E_s_dk(d, k);
            E_log_eta += E_log_eta_d(d);
        }
        a = a_alpha + E_s; b = b_alpha - E_log_eta;
        E_alpha = a / b;
        G_alpha = Math.exp(SpecFun.digamma(a)) / b;
    }

    // update order 2
    private void updateBeta() {
        int k, w;
        double a, b, E_t_kw;
        E_t = 0;
        E_log_xi = 0;
        Arrays.fill(E_t_w, 0);
        for(w = 0;w < W;w++)
            digamma_G_beta_x_tau_w[w] = SpecFun.digamma(G_beta_x_tau_w[w] = G_beta * G_tau_w[w]);
        digamma_E_beta = SpecFun.digamma(E_beta);
        for(k = 0;k < K;k++) {
            for(w = 0;w < W;w++) {
                E_t_w[w] += E_t_kw = E_t_kw(k, w);
                E_t += E_t_kw;
            }
            E_log_xi += E_log_xi_k(k);
        }
        a = a_beta + E_t; b = b_beta - E_log_xi;
        E_beta = a / b;
        G_beta = Math.exp(SpecFun.digamma(a)) / b;
    }

    private Comparator topicsorter = new Comparator() {
        public int compare(Object o1, Object o2) {
            double E_n_k1 = E_n_k[(Integer)o1], E_n_k2 = E_n_k[(Integer)o2];
            return E_n_k1 < E_n_k2 ? 1 : E_n_k1 > E_n_k2 ? -1 : 0;
        }};

    // update order 3
    private void updatePi() {
        int k, l, d, dk;
        double a, b, psi_b, psi_a_plus_b;
        Arrays.sort(dk_k, topicsorter);
        E_log_one_minus_pi_tilde = 0;
        for(k = 0;k < K;k++) {
            a = 0; b = 0; dk = dk_k[k];
            for(d = 0;d < D;d++) {
                a += E_s_dk[d][dk];
                for(l = k+1;l < K;l++)
                    b += E_s_dk[d][dk_k[l]];
            }
            E_s_k[dk] = a; E_s_gt_k[dk] = b;
            a += 1.0; b += E_gamma; psi_b = SpecFun.digamma(b); psi_a_plus_b = SpecFun.digamma(a + b);
            G_pi_tilde_k[dk] = Math.exp(SpecFun.digamma(a) - psi_a_plus_b);
            G_one_minus_pi_tilde_k[dk] = Math.exp(psi_b - psi_a_plus_b);
            E_log_one_minus_pi_tilde += psi_b - psi_a_plus_b;
            E_pi_tilde_k[dk] = a / (a + b);
            E_one_minus_pi_tilde_k[dk] = b / (a + b);
        }
        for(k = 0;k < K;k++) {
            dk = dk_k[k];
            G_pi_k[dk] = G_pi_tilde_k[dk];
            E_pi_k[dk] = E_pi_tilde_k[dk];
            for(l = 0;l < k;l++) {
                G_pi_k[dk] *= G_one_minus_pi_tilde_k[dk_k[l]];
                E_pi_k[dk] *= E_one_minus_pi_tilde_k[dk_k[l]];
            }
        }
    }

    // update order 4
    private void updateGamma() {
        double a, b;
        a = a_gamma + (double)K; b = b_gamma - E_log_one_minus_pi_tilde;
        E_gamma = a / b;
    }

    // update order 5
    private void updateTau() {
        int w;
        double exp_psi = Math.exp(SpecFun.digamma(E_t + a_tau * (double)W));
        for(w = 0;w < W;w++) {
            G_tau_w[w] = Math.exp(SpecFun.digamma(a_tau + E_t_w[w])) / exp_psi;
            E_tau_w[w] = (a_tau + E_t_w[w]) / (E_t + a_tau * (double)W);
        }
    }

    private void estimateMixtures() {
        int d, k, w;
        double sum;
        theta_dk = new double[D][K];
        phi_kw = new double[K][W];
        for(d = 0;d < D;d++) {
            sum = 0;
            for(k = 0;k < K;k++)
                sum += theta_dk[d][k] = (E_alpha * G_pi_k[k] + E_n_dk[d][k]);
            for(k = 0;k < K;k++)
                theta_dk[d][k] /= sum;
        }
        for(k = 0;k < K;k++) {
            sum = 0;
            for(w = 0;w < W;w++)
                sum += phi_kw[k][w] = (E_beta * G_tau_w[w] + E_n_kw[k][w]);
            for(w = 0;w < W;w++)
                phi_kw[k][w] /= sum;
        }
    }

    private double avgLogLikelihood() {
        int d, k, p, pe, N = 0;
        double sum = 0, _sum, n;
        for(d = 0;d < D;d++) {
            pe = ho_wd.pointr[d+1];
            for(p = ho_wd.pointr[d];p < pe;p++) {
                _sum = 0;
                for(k = 0;k < K;k++)
                    _sum += theta_dk[d][k] * phi_kw[k][ho_wd.rowind[p]];
                N += (int)(n = Math.ceil(ho_wd.value[p]));
                sum += Math.log(_sum) * n;
            }
        }
        return sum / (double)N;
    }

    private void describeTiming() {
        StringBuilder sb = new StringBuilder();
        sb.append("real variational posterior update milliseconds: ");
        sb.append(real_variational_posterior_update_millis);
        sb.append("\n");
        sb.append("real hyperparamter update milliseconds: ");
        sb.append(real_hyperparameter_update_millis);
        sb.append("\n");
        sb.append("real bookkeeping update milliseconds: ");
        sb.append(real_bookkeeping_millis);
        sb.append("\n");
        sb.append("real total time: ");
        sb.append(real_variational_posterior_update_millis + real_hyperparameter_update_millis + real_bookkeeping_millis);
        sb.append("\n");
        timing = sb.toString();
    }
}
