/*
 * (C) Copyright 2005-2011, Gregor Heinrich (gregor :: arbylon : net) \
 * (This file is part of the knowceans-ilda experimental software package
 */
/*
 * knowceans-ilda 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.
 */
/*
 * knowceans-ilda 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, write to the Free Software Foundation, Inc., 59 Temple
 * Place, Suite 330, Boston, MA 02111-1307 USA
 */
package org.knowceans.topics.simple;

import static java.lang.Math.log;

import java.io.FileNotFoundException;
import java.io.PrintStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.SortedSet;
import java.util.TreeSet;

import org.knowceans.corpus.NumCorpus;
import org.knowceans.corpus.VisCorpus;
import org.knowceans.util.ArrayIo;
import org.knowceans.util.ArrayUtils;
import org.knowceans.util.CokusRandom;
import org.knowceans.util.DirichletEstimation;
import org.knowceans.util.IndexQuickSort;
import org.knowceans.util.RandomSamplers;
import org.knowceans.util.StopWatch;
import org.knowceans.util.Vectors;

/**
 * LDA Gibbs sampler with nonparametric prior (HDP):
 * <p>
 * (m,k | alpha * tau | gamma), k->inf, (k,t | beta)
 * <p>
 * using Teh et al. (2006) approach for the direct assignment sampler, with
 * modular LDA parametric sampler first published by Griffiths (2002) and
 * explained in Heinrich (2005). For the original LDA paper, see Blei et al.
 * (2002).
 * <p>
 * The general idea is to retain as much as possible of the standard LDA Gibbs
 * sampler, which is possible by alternatingly sampling the finite case with K +
 * 1 topics and resampling the topic weights taking into account the current
 * assignments of data items to topics and pruning or expanding the topic set
 * accordingly.
 * <p>
 * I tried to find the (subjectively) best tradeoff between simplicity and the
 * JASA paper (Teh et al. 2006). Therefore I have only used the direct
 * assignment method.
 * <p>
 * The implementation uses lists instead of primitive arrays, but for
 * performance reasons, this may be changed to have a bound Kmax to allocate
 * fixed-size arrays, similar to a truncated DP.
 * <p>
 * Caveats: (1) Performance is not a core criterion, and OOP encapsulation is
 * ignored for compactness' sake. (2) Code still uses the likelihood function of
 * LDA, and without the hyperparameter terms.
 * <p>
 * LICENSE: GPL3, see: http://www.gnu.org/licenses/gpl-3.0.html
 * <p>
 * References:
 * <p>
 * D.M. Blei, A. Ng, M.I. Jordan. Latent Dirichlet Allocation. NIPS, 2002
 * <p>
 * T. Griffiths. Gibbs sampling in the generative model of Latent Dirichlet
 * Allocation. TR, 2002, www-psych.stanford.edu/~gruffydd/cogsci02/lda.ps
 * <p>
 * G. Heinrich. Parameter estimation for text analysis. TR, 2009,
 * www.arbylon.net/publications/textest2.pdf
 * <p>
 * G. Heinrich. "Infinite LDA" -- implementing the HDP with minimum code
 * complexity. TN2011/1, www.arbylon.net/publications/ilda.pdf
 * <p>
 * Y.W. Teh, M.I. Jordan, M.J. Beal, D.M. Blei. Hierarchical Dirichlet
 * Processes. JASA, 101:1566-1581, 2006
 * 
 * @author (c) 2008-2011 Gregor Heinrich, gregor :: arbylon : net
 * @version 0.95
 */
public class IldaGibbs implements ISimpleGibbs, ISimpleQueryGibbs, ISimplePpx {

	private static TopicMatrixPanel vis;

	/**
	 * test driver for mixture network Gibbs sampler
	 * 
	 * @param args
	 */
	public static void main(String[] args) {

		int niter = 500, niterq = 10;
		String filebase = "nips/nips";
		// file or synthetic
		boolean usefile = true;
		// topic display panel
		boolean display = true;

		Random rand = new CokusRandom(56567651);
		NumCorpus corpus;
		if (usefile) {
			corpus = new NumCorpus(filebase + ".corpus");
			if (display) {
				// panel is K x K terms large, so let's use the sqrt.
				vis = new TopicMatrixPanel(900, 400, (int) Math.sqrt(corpus
						.getNumTerms()), 1);
			}

		} else {
			// test with generated corpus
			int K = 10;
			corpus = VisCorpus.generateLdaCorpus(K, 1000, 200);
			if (display) {
				vis = new TopicMatrixPanel(900, 400, K, 300 / K);
			}
		}

		// corpus.reduce(100, rand);
		corpus.split(10, 2, rand);
		NumCorpus train = (NumCorpus) corpus.getTrainCorpus();
		NumCorpus test = (NumCorpus) corpus.getTestCorpus();
		int[][] w = train.getDocWords(rand);
		int[][] wq = test.getDocWords(rand);
		int K0 = 0;
		int V = corpus.getNumTerms();
		double alpha = 1.;
		// beta = 1 --> K = 12,.5-->16, .1-->26@200, 75@500, 115@645 (beta
		// should be larger),
		//
		double beta = .1;
		double gamma = 1.5;

		// run sampler
		IldaGibbs gs = new IldaGibbs(w, wq, K0, V, alpha, beta, gamma, rand);
		gs.init();
		System.out.println("initialised");
		System.out.println(gs);

		// initial test
		gs.initq();
		gs.runq(niterq);
		System.out.println("perplexity = " + gs.ppx());

		StopWatch.start();
		System.out.println("starting Gibbs sampler with " + niter
				+ " iterations");

		gs.run(niter);
		System.out.println(StopWatch.format(StopWatch.stop()));

		// test
		gs.initq();
		gs.runq(niterq);
		System.out.println("perplexity = " + gs.ppx());
		System.out.println(gs);

		gs.packTopics();

		System.out.println("finished");
		System.out.println(gs);
		if (!usefile)
			try {
				PrintStream bw = new PrintStream(filebase + ".ilda.result");
				gs.print(bw, filebase, corpus.getOrigDocIds()[0],
						train.getNumWords());
				bw.close();
				System.out.println("done");
			} catch (FileNotFoundException e) {
				e.printStackTrace();
			}

	} // main

	/**
	 * print the result of the model, save topics to files
	 * 
	 * @param filebase
	 * @param docids
	 * @param W
	 */
	private void print(PrintStream out, String filebase, int[] docids, int W) {
		phi = new double[K][V];
		double[][] theta = new double[M][K];
		for (int k = 0; k < K; k++) {
			for (int t = 0; t < V; t++) {
				phi[k][t] = (nkt.get(k)[t] + beta) / (nk.get(k) + beta * V);
			}
		}
		for (int m = 0; m < M; m++) {
			for (int k = 0; k < K; k++) {
				theta[m][k] = (nmk[m].get(k) + alpha)
						/ (w[m].length + alpha * K);
			}
		}
		LdaTopics lt = new LdaTopics(filebase, theta, phi);
		for (int m = 0; m < M; m++) {
			// no labels
			out.println(lt.printDocument(m, docids[m], 10, false, true));
			// out.println(lt.printDocument(m, docids[m], 10, true, true));
		}
		int[] kk = IndexQuickSort.sort(nk);
		IndexQuickSort.reverse(kk);
		for (int k = 0; k < K; k++) {
			out.println(String.format("%d (%2.5f / %d): %s", k, nk.get(kk[k])
					/ (double) W * K, K, lt.printTopic(kk[k], 20)));
		}
		ArrayIo.saveBinaryMatrix(filebase + ".ilda.theta.zip", theta);
		ArrayIo.saveBinaryMatrix(filebase + ".ilda.phi.zip", phi);
	}

	private int[][] w;
	private int[][] wq;
	/**
	 * inactive components: index = index in count arrays, element = value in z.
	 * Filled on component removal.
	 */
	// reuse the component emptied last
	// private List<Integer> kgaps;
	// reuse component with the lowest index
	private SortedSet<Integer> kgaps;
	/**
	 * active components: value = value in z and count arrays, which should
	 * always match. This array is never removed elements from but inactive
	 * elements are set to -1. This allows to reuse components (gaps) and to
	 * keep the indices of z and counts identical.
	 */
	private List<Integer> kactive;
	private List<Integer>[] nmk;
	private int[][] nmkq;
	private List<int[]> nkt;
	private List<Integer> nk;
	private double[][] phi;
	private int[][] z;
	private int[][] zq;

	private double[] pp;
	/**
	 * step to increase the sampling array
	 */
	public final int ppstep = 10;
	/**
	 * precision of the 2nd-level DP
	 */
	private double alpha;
	/**
	 * mean of the 2nd-level DP = sample from 1st-level DP
	 */
	private ArrayList<Double> tau;
	/**
	 * parameter of root base measure (= component Dirichlet)
	 */
	private double beta;
	/**
	 * precision of root DP
	 */
	private double gamma;

	// hyperparameters for DP and Dirichlet samplers
	// Teh+06: Docs: (1, 1), M1-3: (0.1, 0.1); HMM: (1, 1)
	double aalpha = 5;
	double balpha = 0.1;
	double abeta = 0.1;
	double bbeta = 0.1;
	// Teh+06: Docs: (1, 0.1), M1-3: (5, 0.1), HMM: (1, 1)
	double agamma = 5;
	double bgamma = 0.1;
	// number of samples for parameter samplers
	int R = 10;

	/**
	 * total number of tables
	 */
	private double T;

	private Random rand;
	RandomSamplers samp;
	private int iter;
	/**
	 * current number of non-empty components
	 */
	private int K;
	private int M;
	private int Mq;
	private int Wq;
	private int V;
	private boolean inited = false;
	private boolean fixedK = false;
	private boolean fixedHyper = false;

	/**
	 * parametrise gibbs sampler
	 * 
	 * @param w
	 *            word tokens
	 * @param wq
	 *            word tokens (testing)
	 * @param K
	 *            initial number of topics: may be 0 if gamma > 0.
	 * @param V
	 *            number of terms
	 * @param alpha
	 *            node A precision (document DP)
	 * @param gamma
	 *            node A precision (root DP), 0 for fixed K: plain LDA.
	 * @param beta
	 *            node B hyperparam
	 * @param rand
	 *            random number generator
	 */
	public IldaGibbs(int[][] w, int[][] wq, int K, int V, double alpha,
			double beta, double gamma, Random rand) {
		// assign
		this.w = w;
		this.wq = wq;
		// start with 0 or more topics
		this.K = K;
		this.alpha = alpha;
		this.beta = beta;
		this.gamma = gamma;
		if (gamma == 0) {
			this.fixedK = true;
		}
		this.M = w.length;
		this.Mq = wq.length;
		this.V = V;
		this.rand = rand;
		this.samp = new RandomSamplers(rand);
	}

	/**
	 * initialise Markov chain
	 */
	@SuppressWarnings("unchecked")
	public void init() {
		// allocate
		nmk = new ArrayList[M];
		nkt = new ArrayList<int[]>();
		nk = new ArrayList<Integer>();
		z = new int[M][];
		for (int m = 0; m < M; m++) {
			nmk[m] = new ArrayList<Integer>();
			for (int k = 0; k < K; k++) {
				nmk[m].add(0);
			}
			z[m] = new int[w[m].length];
		}
		// indexing lists
		kactive = new ArrayList<Integer>();
		// kgaps = new ArrayList<Integer>();
		kgaps = new TreeSet<Integer>();
		// create mean weights
		tau = new ArrayList<Double>();
		for (int k = 0; k < K; k++) {
			kactive.add(k);
			nkt.add(new int[V]);
			nk.add(0);
			// set to value for fixed K
			tau.add(1. / K);
		}
		// tau has one dimension more
		tau.add(1. / K);
		pp = new double[K + ppstep];
		// initialise (run without decrements because z[*][*] = -1)
		run(1);
		if (!fixedK) {
			updateTau();
		}
		inited = true;
	}

	/**
	 * initialise Markov chain for querying
	 */
	public void initq() {
		// compute parameters
		int Kg = K + kgaps.size();
		phi = new double[Kg][V];
		for (int kk = 0; kk < K; kk++) {
			int k = kactive.get(kk);
			for (int t = 0; t < V; t++) {
				phi[k][t] = (nkt.get(k)[t] + beta) / (nk.get(k) + V * beta);
			}
		}
		// allocate
		nmkq = new int[Mq][Kg];
		zq = new int[Mq][];
		Wq = 0;
		// initialise
		for (int m = 0; m < Mq; m++) {
			zq[m] = new int[wq[m].length];
			for (int n = 0; n < wq[m].length; n++) {
				int k = rand.nextInt(K);
				zq[m][n] = k;
				nmkq[m][k]++;
				Wq++;
			}
		}
	}

	/**
	 * run Gibbs sampler
	 * 
	 * @param niter
	 *            number of Gibbs iterations
	 */
	public void run(int niter) {

		for (iter = 0; iter < niter; iter++) {
			System.out.println(iter);
			System.out.println(this);
			for (int m = 0; m < M; m++) {
				for (int n = 0; n < w[m].length; n++) {
					// sampling z
					int k, kold = -1;
					int t = w[m][n];
					if (inited) {
						k = z[m][n];
						// decrement
						nmk[m].set(k, nmk[m].get(k) - 1);
						nkt.get(k)[t]--;
						nk.set(k, nk.get(k) - 1);
						kold = k;
					}
					// compute weights
					double psum = 0;
					// (37)
					for (int kk = 0; kk < K; kk++) {
						k = kactive.get(kk);
						pp[kk] = (nmk[m].get(k) + alpha * tau.get(k)) * //
								(nkt.get(k)[t] + beta) / (nk.get(k) + V * beta);
						psum += pp[kk];
					}
					// likelihood of new component
					if (!fixedK) {
						pp[K] = alpha * tau.get(K) / V;
						psum += pp[K];
					}
					double u = rand.nextDouble();
					u *= psum;
					psum = 0;
					int kk = 0;
					for (; kk < K + 1; kk++) {
						psum += pp[kk];
						if (u <= psum) {
							break;
						}
					}
					// reassign and increment
					if (kk < K) {
						k = kactive.get(kk);
						z[m][n] = k;
						nmk[m].set(k, nmk[m].get(k) + 1);
						nkt.get(k)[t]++;
						nk.set(k, nk.get(k) + 1);
					} else {
						assert (!fixedK);
						z[m][n] = spawnTopic(m, t);
						updateTau();
						System.out.println("K = " + K);
					}
					// empty topic?
					if (inited && nk.get(kold) == 0) {
						// remove the object not the index
						kactive.remove((Integer) kold);
						kgaps.add(kold);
						assert (Vectors.sum(nkt.get(kold)) == 0
								&& nk.get(kold) == 0 && nmk[m].get(kold) == 0);
						K--;
						System.out.println("K = " + K);
						updateTau();
					}
				} // n
			} // m
			if (vis != null) {
				vis.setTopics(nkt);
			}
			if (!fixedK) {
				updateTau();
			}
			if (iter > 10 && !fixedHyper) {
				updateHyper();
			}
		} // i
	}

	/**
	 * query Gibbs sampler. This assumes the standard LDA model as we know the
	 * dimensionality from the training set, therefore topics need to be pruned.
	 * 
	 * @param niter
	 *            number of Gibbs iterations
	 */
	public void runq(int niter) {
		for (int qiter = 0; qiter < niter; qiter++) {
			for (int m = 0; m < nmkq.length; m++) {
				for (int n = 0; n < wq[m].length; n++) {
					// decrement
					int k = zq[m][n];
					int t = wq[m][n];
					nmkq[m][k]--;
					// compute weights
					double psum = 0;
					for (int kk = 0; kk < K; kk++) {
						pp[kk] = (nmkq[m][kk] + alpha) * phi[kk][t];
						psum += pp[kk];
					}
					// sample
					double u = rand.nextDouble() * psum;
					psum = 0;
					int kk = 0;
					for (; kk < K; kk++) {
						psum += pp[kk];
						if (u <= psum) {
							break;
						}
					}
					// reassign and increment
					zq[m][n] = kk;
					nmkq[m][kk]++;
				} // n
			} // m
		} // i
	}

	/**
	 * adds a topic to the list of active topics, either by reusing an existing
	 * inactive index (gap) or increasing the count arrays. NB: Within this
	 * method, the state is inconsistent.
	 * 
	 * @param m
	 *            current document
	 * @param t
	 *            current term
	 * @return index of topic spawned
	 */
	private int spawnTopic(int m, int t) {
		int k;
		if (kgaps.size() > 0) {
			// reuse gap
			// k = kgaps.remove(kgaps.size() - 1);
			k = kgaps.first();
			kgaps.remove(k);
			kactive.add(k);
			nmk[m].set(k, 1);
			nkt.get(k)[t] = 1;
			nk.set(k, 1);
		} else {
			// add element to count arrays
			k = K;
			for (int i = 0; i < M; i++) {
				nmk[i].add(0);
			}
			kactive.add(K);
			nmk[m].set(K, 1);
			nkt.add(new int[V]);
			nkt.get(K)[t] = 1;
			nk.add(1);
			tau.add(0.);
		}
		K++;
		if (pp.length <= K) {
			pp = new double[K + ppstep];
		}
		return k;
	}

	/**
	 * reorders topics such that no gaps exist in the count arrays and topics
	 * are ordered with their counts descending. Removes any gap dimensions.
	 */
	public void packTopics() {
		// sort topics by size
		int[] knew2k = IndexQuickSort.sort(nk);
		IndexQuickSort.reverse(knew2k);
		// reorder and weed out empty count arrays
		IndexQuickSort.reorder(nk, knew2k);
		IndexQuickSort.reorder(nkt, knew2k);
		for (int i = 0; i < kgaps.size(); i++) {
			nk.remove(nk.size() - 1);
			nkt.remove(nkt.size() - 1);
		}
		for (int m = 0; m < M; m++) {
			IndexQuickSort.reorder(nmk[m], knew2k);
			for (int i = 0; i < kgaps.size(); i++) {
				nmk[m].remove(nmk[m].size() - 1);
			}
		}
		// any new topics will be appended
		kgaps.clear();
		int[] k2knew = IndexQuickSort.inverse(knew2k);
		// rewrite topic labels
		for (int i = 0; i < K; i++) {
			kactive.set(i, k2knew[kactive.get(i)]);
		}
		for (int m = 0; m < M; m++) {
			for (int n = 0; n < w[m].length; n++) {
				z[m][n] = k2knew[z[m][n]];
			}
		}
	}

	/**
	 * prune topics and update tau, the root DP mixture weights.
	 */
	private void updateTau() {
		// (40) sample mk
		double[] mk = new double[K + 1];
		// TODO: average multi-sample?
		for (int kk = 0; kk < K; kk++) {
			int k = kactive.get(kk);
			for (int m = 0; m < M; m++) {
				if (nmk[m].get(k) > 1) {
					// number of tables a CRP(alpha tau) produces for nmk items
					mk[kk] += samp.randAntoniak(alpha * tau.get(k), //
							nmk[m].get(k));
				} else {
					mk[kk] += nmk[m].get(k);
				}
			}
		}
		// number of tables
		T = Vectors.sum(mk);
		mk[K] = gamma;
		// (36) sample tau
		double[] tt = samp.randDir(mk);
		for (int kk = 0; kk < K; kk++) {
			int k = kactive.get(kk);
			tau.set(k, tt[kk]);
		}
		tau.set(K, tt[K]);
	}

	/**
	 * update scalar DP hyperparameters alpha, gamma and Dirichlet
	 * hyperparameter beta. Assumes that T is updated (by updateTau).
	 */
	private void updateHyper() {
		for (int r = 0; r < R; r++) {
			// gamma: root level (Escobar+West95) with n = T
			// (14)
			double eta = samp.randBeta(gamma + 1, T);
			double bloge = bgamma - log(eta);
			// (13')
			// double pie = 1. / (1. - (T * bloge / (gamma + K - 1)));
			double pie = 1. / (1. + (T * bloge / (agamma + K - 1)));
			// (13)
			int u = samp.randBernoulli(pie);
			gamma = samp.randGamma(agamma + K - 1 + u, 1. / bloge);

			// alpha: document level (Teh+06)
			double qs = 0;
			double qw = 0;
			for (int m = 0; m < M; m++) {
				// (49) (corrected)
				qs += samp.randBernoulli(w[m].length / (w[m].length + alpha));
				// (48)
				qw += log(samp.randBeta(alpha + 1, w[m].length));
			}
			// (47)
			alpha = samp.randGamma(aalpha + T - qs, 1. / (balpha - qw));
		}
		int[] ak = (int[]) ArrayUtils.asPrimitiveArray(nk);
		int[][] akt = new int[K][V];
		for (int k = 0; k < K; k++) {
			akt[k] = nkt.get(k);
		}
		beta = DirichletEstimation
				.estimateAlphaMap(akt, ak, beta, abeta, bbeta);
	}

	/**
	 * @return the perplexity of the last query sample.
	 */
	public double ppx() {
		// TODO: this uses LDA's perplexity --> add hyperparameters and DP stuff
		double loglik = 0;
		// compute thetaq
		double[][] thetaq = new double[Mq][K];
		for (int m = 0; m < Mq; m++) {
			for (int k = 0; k < K; k++) {
				thetaq[m][k] = (nmkq[m][k] + alpha)
						/ (wq[m].length + K * alpha);
			}
		}
		// compute ppx
		for (int m = 0; m < Mq; m++) {
			for (int n = 0; n < wq[m].length; n++) {
				double sum = 0;
				for (int k = 0; k < K; k++) {
					sum += thetaq[m][k] * phi[k][wq[m][n]];
				}
				loglik += Math.log(sum);
			}
		}
		return Math.exp(-loglik / Wq);
	}

	// ////////////// output routines ////////////////

	/**
	 * assemble a string of overview information.
	 */
	@Override
	public String toString() {
		return String.format("ILDA: M = %d, K = %d, V = %d, "
				+ "alpha = %2.5f, beta = %2.5f, gamma = %2.5f", //
				M, K, V, alpha, beta, gamma);
	}
}