//    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
//    "Fast Collapsed Gibbs Sampling For Latent Dirichlet Allocation"
//    written by Ian Porteous, David Newman, Alexander Ihler, Arthur Asuncion, Padhraic Smyth, and Max Welling.

package com.pawjaw.topics;

import com.pawjaw.linalg.Matrix;
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;

public class CGLDA {
    // output
    public long real_burnin_millis, real_estimates_millis;
    public String timing;
    public List<Double> avglls = new ArrayList();  // record of average log likelihood of held out training data during inference
    public List<int[]> n_ks = new ArrayList();  // record of count of topic [k] assignments over all tokens during inference
    public int[] n_k; // count of topic [k] assignments over all tokens
    public Integer[] dk_k; // map topic [k] in descending order of n_k
    public double[][] phi_kw; // estimate of latent variable phi (multinomial params) for topic [k] word [w]
    public double[][] theta_dk; // estimate latent variable theta (multinomial params) for document [d] topic [k]
    
    // defaults
    public static final double DEFAULT_ALPHASUM = 1.0; // alphasum = alpha x # topics
    public static final double DEFAULT_BETA = 0.01;
    public static final int DEFAULT_BURNIN = 1;
    public static final int DEFAULT_SAMPLEFREQ = 1;
    public static final int DEFAULT_SAMPLES = 200;
    
    // constructors, all work is done during construction
    public CGLDA(Matrix word_count_documents_in_columns, int topics, ProgressCallback notifier) {
        K = topics;
        W = word_count_documents_in_columns.rows;
        alpha = DEFAULT_ALPHASUM / (double)topics;
        beta = DEFAULT_BETA;
        beta_wk = null;
        beta_k = null;
        ho_wd = null;
        setup(word_count_documents_in_columns, topics);
        run(DEFAULT_BURNIN, DEFAULT_SAMPLEFREQ, DEFAULT_SAMPLES, 0, notifier);
    }
    
    public CGLDA(Matrix word_count_documents_in_columns_training, Matrix word_count_documents_in_columns_heldout,
            int topics, double alpha, double beta, int burnin, int samplefreq, int samples, int heldout_log_likelihood_estimates,
            ProgressCallback notifier) {
        K = topics;
        W = word_count_documents_in_columns_training.rows;
        this.alpha = alpha;
        this.beta = beta;
        beta_wk = null;
        beta_k = null;
        ho_wd = word_count_documents_in_columns_heldout;
        setup(word_count_documents_in_columns_training, topics);
        run(burnin, samplefreq, samples, heldout_log_likelihood_estimates, notifier);
    }

    public CGLDA(Matrix word_count_documents_in_columns, int topics, double alpha, double beta,
            int burnin, int samplefreq, int samples, ProgressCallback notifier) {
        K = topics;
        W = word_count_documents_in_columns.rows;
        this.alpha = alpha;
        this.beta = beta;
        beta_wk = null;
        beta_k = null;
        ho_wd = null;
        setup(word_count_documents_in_columns, topics);
        run(burnin, samplefreq, samples, 0, notifier);
    }

    public CGLDA(Matrix word_count_documents_in_columns, Matrix word_count_documents_in_columns_heldout,
            int topics, double alpha, double[][] beta_wk, int burnin, int samplefreq, int samples, int heldout_log_likelihood_estimates,
            ProgressCallback notifier) {
        this.alpha = alpha;
        beta = DEFAULT_BETA;
        this.beta_wk = beta_wk;
        beta_k = new double[topics];
        ho_wd = word_count_documents_in_columns_heldout;
        setup(word_count_documents_in_columns, topics);
        run(burnin, samplefreq, samples, heldout_log_likelihood_estimates, notifier);
    }

    public CGLDA(Matrix word_count_documents_in_columns, int topics, double alpha, double[][] beta_wk,
            int burnin, int samplefreq, int samples, ProgressCallback notifier) {
        this.alpha = alpha;
        beta = DEFAULT_BETA;
        this.beta_wk = beta_wk;
        beta_k = new double[topics];
        ho_wd = null;
        setup(word_count_documents_in_columns, topics);
        run(burnin, samplefreq, samples, 0, notifier);
    }

    /************************ private section ************************/
    private int K; // number of topics
    private int D; // number of documents
    private int W; // number of words
    private int E = 0; // estimates taken
    private double alpha; // uniform parameters for Dirichlet prior for document-topic assignment
    private double beta; // default dirichlet parameter for topic-word assignment when beta_wk is not specified
    private Random r; // random number generator for gibbs sampling
    private Matrix n_wd; // training document word counts word [w] document [d]
    private Matrix ho_wd = null; // held-out documents
    private int[] n_d; // length of document [d]
    private int[][] n_kw; // count of topic [k] word [w] assignments
    private int[][] n_dk; // count of document [d] topic [k] assignments
    private int[][] k_di; // current assignment of document [d] token [i] to topic k
    private double[] prs; // temporary storage for calculating token topic assignments
    private double[] beta_k; // sum of beta_wk[w]
    private double[][] beta_wk; // dense dirichlet alpha parameter matrix serving as prior for topic-word assignment word [w] topic [k]

    private void setup(Matrix word_count_documents_in_columns, int topics) {
        int d, p, pe, k, w;
        n_wd = word_count_documents_in_columns;
        K = topics;
        D = n_wd.cols;
        W = n_wd.rows;
        n_d = new int[D];
        n_k = new int[K];
        n_kw = new int[K][W];
        n_dk = new int[D][K];
        prs = new double[K];
        k_di = new int[D][];
        phi_kw = new double[K][W];
        theta_dk = new double[D][K];
        r = new Random();
        for(d = 0;d < D;d++) {
            pe = n_wd.pointr[d+1];
            for(p = n_wd.pointr[d];p < pe;p++)
                n_d[d] += (int)Math.ceil(n_wd.value[p]);
            k_di[d] = new int[n_d[d]];
        }
        dk_k = new Integer[K];
        for(k = 0;k < K;k++)
            dk_k[k] = k;
        if(beta_wk != null)
            for(k = 0;k < K;k++)
                for(w = 0;w < W;w++)
                    beta_k[k] += beta_wk[w][k];
    }
    
    private void run(int burnin, int samplefreq, int samples, int heldout_log_likelihood_estimates, ProgressCallback notifier) {
        boolean use_uniform_beta = beta_wk == null;
        int iteration, sample = 0, count = 1, total = burnin + samples * samplefreq;
        int avgllsample = heldout_log_likelihood_estimates <= 1 ? 0 : (int)Math.max(samples / heldout_log_likelihood_estimates, 1);
        long start;
        initializeTopicAssignments();
        start = System.currentTimeMillis();
        for(iteration = 0;iteration < burnin;iteration++) {
            if(use_uniform_beta)
                sampleUniformBeta();
            else
                sampleSpecifiedBeta();
            if(notifier != null)
                notifier.notify(count++, total);
        }
        real_burnin_millis = System.currentTimeMillis() - start;
        start = System.currentTimeMillis();
        while(sample++ < samples) {
            for(iteration = 0;iteration < samplefreq;iteration++) {
                if(use_uniform_beta)
                    sampleUniformBeta();
                else
                    sampleSpecifiedBeta();
                if(notifier != null)
                    notifier.notify(count++, total);
            }
            if(use_uniform_beta)
                estimateMixturesUniformBeta();
            else
                estimateMixturesSpecifiedBeta();
            if(avgllsample > 0 && sample % avgllsample == 0) {
                avglls.add(avgLogLikelihood());
                n_ks.add(Arrays.copyOf(n_k, K));
            }
        }
        real_estimates_millis = System.currentTimeMillis() - start;
        if(heldout_log_likelihood_estimates == 1) {
            avglls.add(avgLogLikelihood());
            n_ks.add(Arrays.copyOf(n_k, K));
        }
        averageEstimates();
        describeTiming();
    }

    private void initializeTopicAssignments() {
        int d, w, p, pe, i, I, k;
        for(d = 0;d < D;d++) {
            pe = n_wd.pointr[d+1];
            i = I = 0;
            for(p = n_wd.pointr[d];p < pe;p++) {
                w = n_wd.rowind[p];
                for(I += (int)n_wd.value[p];i < I;i++) {
                    k = k_di[d][i] = r.nextInt(K);
                    n_k[k]++;
                    n_kw[k][w]++;
                    n_dk[d][k]++;
                }
            }
        }
    }

    private void sampleUniformBeta() {
        int d, w, p, pe, i, I, k;
        double Z, betaW = beta * (double)W;
        for(d = 0;d < D;d++) {
            i = I = 0;
            pe = n_wd.pointr[d+1];
            for(p = n_wd.pointr[d];p < pe;p++) {
                w = n_wd.rowind[p];
                for(I += (int)n_wd.value[p];i < I;i++) {
                    k = k_di[d][i];
                    n_k[k]--;
                    n_kw[k][w]--;
                    n_dk[d][k]--;
                    Z = 0;
                    for(k = 0;k < K;k++)
                        Z += prs[k] = ((double)n_kw[k][w] + beta) / ((double)n_k[k] + betaW) * ((double)n_dk[d][k] + alpha);
                    Z *= r.nextDouble();
                    for(k = 0;k < K;k++)
                        if((Z -= prs[k]) <= 0)
                            break;
                    if(k == K)
                        k = K - 1;
                    n_k[k]++;
                    n_kw[k][w]++;
                    n_dk[d][k]++;
                    k_di[d][i] = k;
                }
            }
        }
    }

    private void sampleSpecifiedBeta() {
        int d, w, p, pe, i, I, k;
        double Z;
        for(d = 0;d < D;d++) {
            i = I = 0;
            pe = n_wd.pointr[d+1];
            for(p = n_wd.pointr[d];p < pe;p++) {
                w = n_wd.rowind[p];
                for(I += (int)n_wd.value[p];i < I;i++) {
                    k = k_di[d][i];
                    n_k[k]--;
                    n_kw[k][w]--;
                    n_dk[d][k]--;
                    Z = 0;
                    for(k = 0;k < K;k++)
                        Z += prs[k] =  ((double)n_kw[k][w] + beta_wk[w][k]) / ((double)n_k[k] + beta_k[k]) * ((double)n_dk[d][k] + alpha);
                    Z *= r.nextDouble();
                    for(k = 0;k < K;k++)
                        if((Z -= prs[k]) <= 0)
                            break;
                    if(k == K)
                        k = K - 1;
                    n_k[k]++;
                    n_kw[k][w]++;
                    n_dk[d][k]++;
                    k_di[d][i] = k;
                }
            }
        }
    }
    
    private void estimateMixturesUniformBeta() {
        int d, k, w;
        double Z;
        E++;
        for(d = 0;d < D;d++) {
            Z = (double)n_d[d] + (double)K * alpha;
            for(k = 0;k < K;k++)
                theta_dk[d][k] += ((double)n_dk[d][k] + alpha) / Z;
        }
        for(k = 0;k < K;k++) {
            Z = (double)n_k[k] + (double)W * beta;
            for(w = 0;w < W;w++)
                phi_kw[k][w] += ((double)n_kw[k][w] + beta) / Z;
        }
    }

    private void estimateMixturesSpecifiedBeta() {
        int d, k, w;
        double Z;
        E++;
        for(d = 0;d < D;d++) {
            Z = (double)n_d[d] + (double)K * alpha;
            for(k = 0;k < K;k++)
                theta_dk[d][k] += ((double)n_dk[d][k] + alpha) / Z;
        }
        for(k = 0;k < K;k++) {
            Z = (double)n_k[k] + beta_k[k];
            for(w = 0;w < W;w++)
                phi_kw[k][w] += ((double)n_kw[k][w] + beta_wk[w][k]) / Z;
        }
    }

    private double avgLogLikelihood() {
        int d, k, p, pe, N = 0;
        double sum = 0, _sum, n, den = E * E;
        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 / den) * n;
            }
        }
        return sum / (double)N;
    }

    private void averageEstimates() {
        int d, k, w;
        for(k = 0;k < K;k++)
            for(w = 0;w < W;w++)
                phi_kw[k][w] /= (double)E;
        for(d = 0;d < D;d++)
            for(k = 0;k < K;k++)
                theta_dk[d][k] /= (double)E;
        Arrays.sort(dk_k, new Comparator() {
            public int compare(Object o1, Object o2) {
                double n_j1 = n_k[(Integer)o1], n_j2 = n_k[(Integer)o2];
                return n_j1 < n_j2 ? 1 : n_j1 > n_j2 ? -1 : 0;
            }});
    }

    private void describeTiming() {
        StringBuilder sb = new StringBuilder();
        sb.append("real burnin sampling milliseconds: ");
        sb.append(real_burnin_millis);
        sb.append("\n");
        sb.append("real estimate sampling milliseconds: ");
        sb.append(real_estimates_millis);
        sb.append("\n");
        sb.append("real total time: ");
        sb.append(real_burnin_millis + real_estimates_millis);
        sb.append("\n");
        timing = sb.toString();
    }
}
