/*
 * (C) Copyright 2005, Gregor Heinrich (gregor :: arbylon : net) (This file is
 * part of the org.knowceans experimental software packages.)
 */
/*
 * LdaGibbsSampler 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 2 of the License, or (at your option) any
 * later version.
 */
/*
 * LdaGibbsSampler 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
 */

/*
 * Created on Mar 6, 2005
 */


import java.text.DecimalFormat;
import java.text.NumberFormat;
import java.io.*;
import java.util.*;


/**
 * Gibbs sampler for estimating the best assignments of topics for words and
 * documents in a corpus. The algorithm is introduced in Tom Griffiths' paper
 * "Gibbs sampling in the generative model of Latent Dirichlet Allocation"
 * (2002).
 * 
 * @author heinrich
 */
public class LdaGibbsSampler {

	/**
	 * document data (term lists)
	 */
	int[][] documents;

	/**
	 * vocabulary size
	 */
	int V;
	
	/**
	 * vocabulary map <int, string>
	 */
	HashMap<Integer, String> V_Map;
	/**
	 * number of topics
	 */
	int K;

	/**
	 * Dirichlet parameter (document--topic associations)
	 */
	double alpha;

	/**
	 * Dirichlet parameter (topic--term associations)
	 */
	double beta;

	/**
	 * topic assignments for each word.
	 */
	int z[][];

	/**
	 * cwt[i][j] number of instances of word i (term?) assigned to topic j.
	 */
	int[][] nw;

	/**
	 * na[i][j] number of words in document i assigned to topic j.
	 */
	int[][] nd;

	/**
	 * nwsum[j] total number of words assigned to topic j.
	 */
	int[] nwsum;

	/**
	 * nasum[i] total number of words in document i.
	 */
	int[] ndsum;

	/**
	 * cumulative statistics of theta
	 */
	double[][] thetasum;

	/**
	 * cumulative statistics of phi
	 */
	double[][] phisum;

	/**
	 * size of statistics
	 */
	int numstats;

	/**
	 * sampling lag (?)
	 */
	private static int THIN_INTERVAL = 20;

	/**
	 * burn-in period
	 */
	private static int BURN_IN = 100;

	/**
	 * max iterations
	 */
	private static int ITERATIONS = 1000;

	/**
	 * sample lag (if -1 only one sample taken)
	 */
	private static int SAMPLE_LAG;

	private static int dispcol = 0;

	/**
	 * Initialise the Gibbs sampler with data.
	 * 
	 * @param V
	 *            vocabulary size
	 * @param data
	 */
	public LdaGibbsSampler(int[][] documents, int V) {

		this.documents = documents;
		this.V = V;
	}

	/**
	 * Initialisation: Must start with an assignment of observations to topics ?
	 * Many alternatives are possible, I chose to perform random assignments
	 * with equal probabilities
	 * 
	 * @param K
	 *            number of topics
	 * @return z assignment of topics to words
	 */
	public void initialState(int K) {
		int i;

		int M = documents.length;

		// initialise count variables.
		nw = new int[V][K];
		nd = new int[M][K];
		nwsum = new int[K];
		ndsum = new int[M];

		// The z_i are are initialised to values in [1,K] to determine the
		// initial state of the Markov chain.

		z = new int[M][];
		for (int m = 0; m < M; m++) {
			int N=5;
			try {
				N = documents[m].length;
			} catch (Exception e) {
				// TODO Auto-generated catch block
				System.out.println("\n m: " + m + "\n");
				e.printStackTrace();
				
			}
			z[m] = new int[N];
			for (int n = 0; n < N; n++) {
				int topic = (int) (Math.random() * K);
				z[m][n] = topic;
				// number of instances of word i assigned to topic j
				nw[documents[m][n]][topic]++;
				// number of words in document i assigned to topic j.
				nd[m][topic]++;
				// total number of words assigned to topic j.
				nwsum[topic]++;
			}
			// total number of words in document i
			ndsum[m] = N;
		}


	}

	/**
	 * Main method: Select initial state ? Repeat a large number of times: 1.
	 * Select an element 2. Update conditional on other elements. If
	 * appropriate, output summary for each run.
	 * 
	 * @param K
	 *            number of topics
	 * @param alpha
	 *            symmetric prior parameter on document--topic associations
	 * @param beta
	 *            symmetric prior parameter on topic--term associations
	 */
	private void gibbs(int K, double alpha, double beta) {
		this.K = K;
		this.alpha = alpha;
		this.beta = beta;

		// init sampler statistics
		if (SAMPLE_LAG > 0) {
			thetasum = new double[documents.length][K];
			phisum = new double[K][V];
			numstats = 0;
		}

		// initial state of the Markov chain:
		initialState(K);

		System.out.println("Sampling " + ITERATIONS
				+ " iterations with burn-in of " + BURN_IN + " (B/S="
				+ THIN_INTERVAL + ").");

		for (int i = 0; i < ITERATIONS; i++) {

			// for all z_i
			for (int m = 0; m < z.length; m++) {
				for (int n = 0; n < z[m].length; n++) {

					// (z_i = z[m][n])
					// sample from p(z_i|z_-i, w)
					int topic;
					topic = sampleFullConditional(m, n);
					z[m][n] = topic;
				}
			}

			
			if(i==BURN_IN)
				System.out.println("\nBURN_IN Time: " + new Date().toString());

			if ((i < BURN_IN) && (i % THIN_INTERVAL == 0)) {
				System.out.print("B");
				//System.out.println(getLogLiklihood());
				printCurrentPhi(i);
				dispcol++;
			}
			// display progress
			if ((i > BURN_IN) && (i % THIN_INTERVAL == 0)) {
				System.out.print("S");
				//System.out.println(getLogLiklihood());
				dispcol++;
			}
			// get statistics after burn-in
			if ((i > BURN_IN) && (SAMPLE_LAG > 0) && (i % SAMPLE_LAG == 0)) {
				updateParams();
				System.out.print("|");
				printCurrentPhi(i);
				//System.out.println(getLogLiklihood());
				if (i % THIN_INTERVAL != 0)
					dispcol++;
			}
			if (dispcol >= 100) {
				System.out.println();
				dispcol = 0;
			}
		}
	}
	/**
	 * Compute the current log-liklihood
	 * P(W|Z) = Sum_k=1_K( log( B(beta + n_k)/ B(beta) ) )
	 * @return
	 * log-liklihood
	 */
	/*
	private double getLogLiklihood(){
		//System.out.println("Computing liklihood");
		double log_liklihood = 0;
		double [] beta_v = new double[V]; // beta_1 to V which is symmetric
		double [][] beta_n_k = new double [K][V]; // beta + n_k
		for(int k=0; k<K; k++)
			for(int v=0; v<V; v++)
				beta_n_k[k][v] = beta + nw[v][k];
		for(int v=0; v<V; v++)
			beta_v[v] = beta;
		for(int k=0; k<K; k++){
			//System.out.println("beta_n_k : " + ArrayToString(beta_n_k[k]));
			double B_beta_n_k = Gamma.B(beta_n_k[k]);
			//System.out.println("B(beta_n_k): " + B_beta_n_k);
			double B_beta = Gamma.B(beta_v);
			//System.out.println("B(beta): " + B_beta);
			double val = Math.log(B_beta_n_k/B_beta);
			//System.out.println("liklihood comp sum: " + val);
			log_liklihood += val;
		}
		return log_liklihood;
	}
	*/
	
	/**
	 * Sample a topic z_i from the full conditional distribution: p(z_i = j |
	 * z_-i, w) = (n_-i,j(w_i) + beta)/(n_-i,j(.) + W * beta) * (n_-i,j(d_i) +
	 * alpha)/(n_-i,.(d_i) + K * alpha)
	 * 
	 * @param m
	 *            document
	 * @param n
	 *            word
	 */
	private int sampleFullConditional(int m, int n) {

		// remove z_i from the count variables
		int topic = z[m][n];
		nw[documents[m][n]][topic]--;
		nd[m][topic]--;
		nwsum[topic]--;
		ndsum[m]--;

		// do multinomial sampling via cumulative method:
		double[] p = new double[K];
		int [] labels = new int[K];
		for (int k = 0; k < K; k++) {
			p[k] = (nw[documents[m][n]][k] + beta) / (nwsum[k] + V * beta)
			* (nd[m][k] + alpha) / (ndsum[m] + K * alpha);
			labels[k] = k;

		}

		//topic = new SampleMultinomial(labels, p).sample();
		// cumulate multinomial parameters
		for (int k = 1; k < p.length; k++) {
			p[k] += p[k - 1];
		}

		// scaled sample because of unnormalised p[]
		double u = Math.random() * p[K - 1];
		for (topic = 0; topic < p.length; topic++) {
			if (u < p[topic])
				break;
		}

		// add newly estimated z_i to count variables
		nw[documents[m][n]][topic]++;
		nd[m][topic]++;
		nwsum[topic]++;
		ndsum[m]++;

		return topic;
	}

	/**
	 * Add to the statistics the values of theta and phi for the current state.
	 */
	private void updateParams() {
		for (int m = 0; m < documents.length; m++) {
			for (int k = 0; k < K; k++) {
				thetasum[m][k] += (nd[m][k] + alpha) / (ndsum[m] + K * alpha);
			}
		}
		for (int k = 0; k < K; k++) {
			for (int w = 0; w < V; w++) {
				phisum[k][w] += (nw[w][k] + beta) / (nwsum[k] + V * beta);
			}
		}
		numstats++;
	}

	/**
	 * Retrieve estimated document--topic associations. If sample lag > 0 then
	 * the mean value of all sampled statistics for theta[][] is taken.
	 * 
	 * @return theta multinomial mixture of document topics (M x K)
	 */
	public double[][] getTheta() {
		double[][] theta = new double[documents.length][K];

		if (SAMPLE_LAG > 0) {
			for (int m = 0; m < documents.length; m++) {
				for (int k = 0; k < K; k++) {
					theta[m][k] = thetasum[m][k] / numstats;
				}
			}

		} else {
			for (int m = 0; m < documents.length; m++) {
				for (int k = 0; k < K; k++) {
					theta[m][k] = (nd[m][k] + alpha) / (ndsum[m] + K * alpha);
				}
			}
		}

		return theta;
	}

	/**
	 * Retrieve estimated topic--word associations. If sample lag > 0 then the
	 * mean value of all sampled statistics for phi[][] is taken.
	 * 
	 * @return phi multinomial mixture of topic words (K x V)
	 */
	public double[][] getPhi() {
		double[][] phi = new double[K][V];
		if (SAMPLE_LAG > 0) {
			for (int k = 0; k < K; k++) {
				for (int w = 0; w < V; w++) {
					phi[k][w] = phisum[k][w] / numstats;
				}
			}
		} else {
			for (int k = 0; k < K; k++) {
				for (int w = 0; w < V; w++) {
					phi[k][w] = (nw[w][k] + beta) / (nwsum[k] + V * beta);
				}
			}
		}
		return phi;
	}
	
	public void printCurrentPhi(int i){
		double [][] phi = new double [K][V];
		if(i<BURN_IN)
		for (int k = 0; k < K; k++) {
			for (int w = 0; w < V; w++) {
				phi[k][w] = (nw[w][k] + beta) / (nwsum[k] + V * beta);
			}
		}
		else
			for (int k = 0; k < K; k++) {
				for (int w = 0; w < V; w++) {
					phi[k][w] = phisum[k][w] / numstats;
				}
			}
		
		System.out.println("\nPrinting Topic Distribution for iteration : " + i + "\n");
		for(int j=0; j<K; j++){
			Phi_k_w [] phi_j = new Phi_k_w [V];
			for(int l=0; l<V; l++)
				phi_j[l] = new Phi_k_w(l, phi[j][l]);
			Arrays.sort(phi_j);

			System.out.println();
			System.out.print("Topic " + (j+1) + ". ");
			for(int m=0; m<125; m++)
				System.out.print(V_Map.get(phi_j[m].w) + ", " );
		}
		
		//new GibbsVisualizer().saveDist(Integer.toString(i), phi);
	}
	/**
	 * Print table of multinomial data
	 * 
	 * @param data
	 *            vector of evidence
	 * @param fmax
	 *            max frequency in display
	 * @return the scaled histogram bin values
	 */
	public static void hist(double[] data, int fmax) {

		double[] hist = new double[data.length];
		// scale maximum
		double hmax = 0;
		for (int i = 0; i < data.length; i++) {
			hmax = Math.max(data[i], hmax);
		}
		double shrink = fmax / hmax;
		for (int i = 0; i < data.length; i++) {
			hist[i] = shrink * data[i];
		}

		NumberFormat nf = new DecimalFormat("00");
		String scale = "";
		for (int i = 1; i < fmax / 10 + 1; i++) {
			scale += "    .    " + i % 10;
		}

		System.out.println("x" + nf.format(hmax / fmax) + "\t0" + scale);
		for (int i = 0; i < hist.length; i++) {
			System.out.print(i + "\t|");
			for (int j = 0; j < Math.round(hist[i]); j++) {
				if ((j + 1) % 10 == 0)
					System.out.print("]");
				else
					System.out.print("|");
			}
			System.out.println();
		}
	}

	/**
	 * Configure the gibbs sampler
	 * 
	 * @param iterations
	 *            number of total iterations
	 * @param burnIn
	 *            number of burn-in iterations
	 * @param thinInterval
	 *            update statistics interval
	 * @param sampleLag
	 *            sample interval (-1 for just one sample at the end)
	 */
	public void configure(int iterations, int burnIn, int thinInterval,
			int sampleLag) {
		ITERATIONS = iterations;
		BURN_IN = burnIn;
		THIN_INTERVAL = thinInterval;
		SAMPLE_LAG = sampleLag;
	}

	/**
	 * Driver with example data.
	 * 
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
		// Format the data in the LDAFormat
		LDADatasetFormatter formater = new LDADatasetFormatter();
		System.out.println("Formating this corpus...");
		int [][] docs = formater.textCorpusFormatter("./all.txt");
		UseSerialization.doSave(formater.map, "vocab_map");
		UseSerialization.doSave(docs, "docs");
		System.out.println("Done...objects serialized to file");
		
		boolean flag = true;
		//Default eclipse path ./
		//String path_prefix = "D:\\eclipse\\workspace\\LatentDirichletAllocation\\";
		String path_prefix = "";
		System.out.println("\nStart Time: " + new Date().toString() + "\n ");
		// vocabulary + documents
		HashMap<Integer, String> map = (HashMap<Integer, String>) UseSerialization.doLoad(path_prefix + "vocab_map");
		int[][] documents = (int [][] )UseSerialization.doLoad(path_prefix + "docs");
		int V = map.size();
		int M = documents.length;
		System.out.println("Data Loaded. # of docs: " + M + ", # of words: " + V);
			
		// # topics
		int K = 20;
		// good values alpha = 2, beta = .5
		double alpha = 50.0/K;
		double beta = .1;

		System.out.println("Latent Dirichlet Allocation using Gibbs Sampling.");
		double[][] theta;
		double[][] phi = new double[K][V];
		if(flag){
			LdaGibbsSampler lda = new LdaGibbsSampler(documents, V);
			lda.V_Map = map;
			//lda.configure(25000, 10000, 1000, 500);
			lda.configure(1200, 300, 400, 75);
			lda.gibbs(K, alpha, beta);
			System.out.println("\nFinish Time: " + new Date().toString());
			theta = lda.getTheta();
			phi = lda.getPhi();
			// Save the posterior distributions in files
			UseSerialization.doSave(phi, path_prefix + "phi.ser");
			UseSerialization.doSave(theta, path_prefix + "theta.ser");
			// Output the topics in the topics.txt
			try{
				BufferedWriter bw = new BufferedWriter(new FileWriter(path_prefix + "topics.txt"));
				bw.write("\nPrinting Topic Distribution for iteration : " + "\n");
				for(int j=0; j<K; j++){
					Phi_k_w [] phi_j = new Phi_k_w [V];
					for(int l=0; l<V; l++)
						phi_j[l] = new Phi_k_w(l, phi[j][l]);
					Arrays.sort(phi_j);

					bw.newLine();
					bw.write("Topic " + (j+1) + ". ");
					for(int m=0; m<125; m++) // top 125 terms for each topic
						bw.write(lda.V_Map.get(phi_j[m].w) + ", " );
				}
				bw.flush();
				bw.close();
			}catch(Exception e){
				e.printStackTrace();
			}

			System.out.println();
			System.out.println();

			//System.out.println("Printing per topic word distribution");
		}

		//System.out.println("Printing Topic Distribution\n");
				
		/*
		phi = (double [][])UseSerialization.doLoad("phi.ser");
		for(int i=0; i<K; i++){
			Phi_k_w [] phi_i = new Phi_k_w [V];
			for(int l=0; l<V; l++)
				phi_i[l] = new Phi_k_w(l, phi[i][l]);
			Arrays.sort(phi_i);

			System.out.println();
			System.out.print("Topic " + i+1 + ". ");
			for(int j=0; j<25; j++)
				System.out.print(V_Map.get(phi_i[j].w) + "  " );
		}

		*/
	}

}