/**
 * This is a simple class for formatting a data set for internal use of LDA
 * Expects a raw text corpus with one document per line
 * @author arjun
 *
 */
import java.io.*;
import java.util.*;

public class LDADatasetFormatter {
	/**
	 * @param path path of file which is a collection of documents, 1 per line
	 * @param  M number of documents in the file referred to by path
	 * @return int [][] the matrix of documents with vocab serialized as the vocabulary
	 */
	HashSet<String> stopwords;
	HashSet<String> vocab;
	HashMap<Integer, String> map;
	
	public LDADatasetFormatter() throws Exception {
		stopwords = new HashSet<String>();
		// Set up stop words
		BufferedReader br = new BufferedReader(new FileReader("stopwords"));
		String line;
		while((line = br.readLine())!=null){
			stopwords.add(line);
		}
		br.close();
	}
	int [][] textCorpusFormatter(String path) throws Exception {
		vocab = new HashSet<String>();
		HashMap<String, Integer> counts = new HashMap<String, Integer>();
		map = new HashMap<Integer, String>();
		HashMap<String, Integer> inverse_map = new HashMap<String, Integer>();

		// Do a raw count of the number of documents in this corpus
		BufferedReader br = new BufferedReader(new FileReader(path));
		String doc;
		int m = 0; // doc count
		while((doc = br.readLine())!=null){
			if(doc.length()> 1)
				m++;
		}
		br.close();
		// Now parse documents
		int documents [][] = new int [m][]; // dummy initalization.
		ArrayList<String []> documents_list = new ArrayList<String []>();
		br = new BufferedReader(new FileReader(path));

		while((doc = br.readLine())!=null){
			if(doc.length()>1){
				doc = doc.toLowerCase();
				String tokens [] = doc.split("[[:][ ][\"][,][;][.][?][!]]+");
				documents_list.add(tokens);
				for(String s : tokens){
					if(!vocab.contains(s)){
						counts.put(s, 1);
						vocab.add(s);
					} else{
						int count = counts.get(s);
						counts.put(s, count+1);
					}
				}
			}

		}
		System.out.println("Total number of documents in the corpus: " + documents.length);

		//Remove tokens appearing less than 5 times in the corpus and stopwords
		HashSet<String> toRemove = new HashSet<String>();
		for(String s : vocab){
			if(stopwords.contains(s) || counts.get(s)<5)
				toRemove.add(s);
		}
		for(String s : toRemove)
			vocab.remove(s);

		// Generate doc in LDAFormat
		// Build vocab map
		int i = 0;
		for(String s : vocab){
			map.put(i, s);
			inverse_map.put(s, i);
			i++;
		}

		// Now generate documents in LDA Format
		for(int doc_index = 0; doc_index < documents_list.size(); doc_index++){
			String tokens [] = documents_list.get(doc_index);
			int word_count = 0;
			for(String s : tokens){
				if(vocab.contains(s))
					word_count++;
			}
			documents[doc_index] = new int[word_count];
			int j = 0;
			for(String s : tokens){
				if(vocab.contains(s))
					documents[doc_index][j++] = inverse_map.get(s);
			}
		}
		return documents;
	}
	
	public static void main(String a []) throws Exception {
		LDADatasetFormatter formater = new LDADatasetFormatter();
		System.out.println("Formating this corpus...");
		int [][] docs = formater.textCorpusFormatter("../Temp/all.txt");
		UseSerialization.doSave(formater.map, "vocab_map");
		UseSerialization.doSave(docs, "docs");
		System.out.println("Done...objects serialized to file");
	}
}

