package net.woodyfolsom.cs6601.p3.ngram;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.Collections;
import java.util.Comparator;
import java.util.Date;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import net.woodyfolsom.cs6601.p3.StockUtil;
import net.woodyfolsom.cs6601.p3.domain.Headline;
import net.woodyfolsom.cs6601.p3.domain.StockPrice;

public class NGramModel {
	static final int MAX_N_GRAM_LENGTH = 3;
	static final int RANDOM_LENGTH = 100;

	private static final String START = "<start>";
	private static final String END = "<end>";
	private static final String UNK = "<unk>";

	private Map<Integer, NGramDistribution> nGrams;
	private Map<Integer, Map<NGram,Double>> nGramPriceAvg;
	
	private int[] totalNGramCounts = new int[MAX_N_GRAM_LENGTH + 1];

	private Pattern wordPattern = Pattern.compile("\\w+");
	private Set<String> nonWords = new HashSet<String>();
	private Set<String> words = new HashSet<String>();
	private Set<String> wordsSeenOnce = new HashSet<String>();

	double getProb(NGram nGram, NGram nMinusOneGram) {
		NGramDistribution nGramDist = nGrams.get(nGram.size());
		double prob;
		if (nGramDist.containsKey(nGram)) {
			// impossible for the bigram not to exist if the trigram does
			NGramDistribution nMinusOneGramDist = nGrams.get(nGram.size() - 1);
			int nMinusOneGramCount = nMinusOneGramDist.get(nMinusOneGram);
			int nGramCount = nGramDist.get(nGram);
			prob = (double) (nGramCount) / nMinusOneGramCount;
		} else {
			// Laplace smoothing
			// prob = 1.0 / ((totalNGramCounts[nGram.size()] + 2.0));
			NGram backoff = nGram.subNGram(1);
			NGram backoffMinusOne = backoff.subNGram(0, backoff.size() - 1);
			return getProb(backoff, backoffMinusOne);
		}

		return prob;
	}

	static String getRandomToken(Map<Integer, NGramDistribution> nGrams, int n,
			NGram nMinusOneGram) {
		List<Entry<NGram, Integer>> matchingNgrams = new ArrayList<Entry<NGram, Integer>>();
		NGramDistribution ngDist = nGrams.get(n);

		for (Entry<NGram, Integer> entry : ngDist.entrySet()) {
			if (entry.getKey().startsWith(nMinusOneGram)) {
				matchingNgrams.add(entry);
			}
		}

		int totalCount = 0;
		for (int i = 0; i < matchingNgrams.size(); i++) {
			totalCount += matchingNgrams.get(i).getValue();
		}
		double random = Math.random();
		int randomNGramIndex = (int) (random * totalCount);

		totalCount = 0;
		for (int i = 0; i < matchingNgrams.size(); i++) {
			int currentCount = matchingNgrams.get(i).getValue();
			if (randomNGramIndex < totalCount + currentCount) {
				NGram ngram = matchingNgrams.get(i).getKey();
				return ngram.get(ngram.size() - 1);
			}
			totalCount += currentCount;
		}
		return getRandomToken(nGrams, n - 1, nMinusOneGram.subNGram(1));
	}

	public NGramModel() {
		// initialize Maps of 0-grams, unigrams, bigrams, trigrams...
		nGrams = new HashMap<Integer, NGramDistribution>();
		for (int i = 0; i <= MAX_N_GRAM_LENGTH; i++) {
			nGrams.put(i, new NGramDistribution());
		}
		nGramPriceAvg = new HashMap<Integer, Map<NGram,Double>>();
		for (int i = 0; i <= MAX_N_GRAM_LENGTH; i++) {
			nGramPriceAvg.put(i, new HashMap<NGram,Double>());
		}
	}
	
	private void addNGram(int nGramLength, NGram nGram, String stockName, Date date, Map<String, Map<Date,StockPrice>> stockTrends) {
		if (nGram.size() < nGramLength) {
			System.out.println("Cannot create " + nGramLength + "-gram from: "
					+ nGram);
		}

		Map<NGram, Integer> nGramCounts = nGrams.get(nGramLength);
		NGram nGramCopy = nGram.copy(nGramLength);

		if (nGramCounts.containsKey(nGramCopy)) {
			int nGramCount = nGramCounts.get(nGramCopy);
			nGramCounts.put(nGramCopy, nGramCount + 1);
		} else {
			nGramCounts.put(nGramCopy, 1);
		}
		
		Map<NGram, Double> nGramPriceAvgs = nGramPriceAvg.get(nGramLength);
		
		NGram nGramCopy2 = nGram.copy(nGramLength);
		
		//TODO GET NEXT TRADING DAY'S DATE
		Date nextDay = StockUtil.getNextTradingDay(date);
		StockPrice stockPrice = stockTrends.get(stockName).get(nextDay);
		double percentChange;
		if (stockPrice == null) {
			percentChange = 0.0;
		} else {
			percentChange = StockUtil.getPercentChange(stockPrice);
		}
				
		if (nGramPriceAvgs.containsKey(nGramCopy2)) {
			double totalPercentChange = nGramPriceAvgs.get(nGramCopy);
			nGramPriceAvgs.put(nGramCopy2, totalPercentChange + percentChange);
		} else {
			nGramPriceAvgs.put(nGramCopy2, percentChange);
		}
	}

	/**
	 * Given an arbitrary String, replace punctuation with spaces, remove
	 * non-alphanumeric characters, prepend with <START> token, append <END>
	 * token.
	 * 
	 * @param rawHeadline
	 * @return
	 */
	private String sanitize(String rawHeadline) {
		String nonpunctuatedHeadline = rawHeadline.replaceAll(
				"[\'\";:,\\]\\[]", " ");
		String alphaNumericHeadline = nonpunctuatedHeadline.replaceAll(
				"[^A-Za-z0-9 ]", "");
		return START + " " + alphaNumericHeadline + " " + END;
	}

	private void calcPerplexity(List<Headline> validationSet, int nGramLimit,
			boolean useUnk) throws FileNotFoundException, IOException {
		List<String> fileByLines = new ArrayList<String>();
		StringBuilder currentLine = new StringBuilder();

		for (Headline headline : validationSet) {
			String sanitizedLine = sanitize(headline.getText());
			// split on whitespace
			String[] tokens = sanitizedLine.toLowerCase().split("\\s+");
			for (String token : tokens) {
				if (!isWord(token)) {
					continue;
				}

				String word;
				if (!words.contains(token) && useUnk) {
					word = UNK;
					// words.add(token);
				} else {
					word = token;
				}

				if (END.equals(word)) {
					currentLine.append(word);
					fileByLines.add(currentLine.toString());
					currentLine = new StringBuilder();
				} else {
					currentLine.append(word);
					currentLine.append(" ");
				}
			}
		}

		int wordNum = 0;
		int sentNum = 0;

		double logProbT = 0.0;
		for (String str : fileByLines) {
			double logProbS = 0.0;
			NGram currentNgram = new NGram(nGramLimit);
			String[] tokens = str.split("\\s+");
			for (String token : tokens) {
				String word;
				if (!words.contains(token) && useUnk) {
					word = UNK;
				} else {
					word = token;
				}
				currentNgram.add(word);
				NGram nMinusOneGram = currentNgram.subNGram(0,
						currentNgram.size() - 1);

				double prob = getProb(currentNgram, nMinusOneGram);

				logProbS += Math.log(prob);
			}
			logProbT += logProbS;
			wordNum += tokens.length;
			sentNum++;
		}

		System.out.println("Evaluated " + sentNum + " sentences and " + wordNum
				+ " words using " + nGramLimit + "-gram model.");
		int N = wordNum + sentNum;
		System.out.println("Total n-grams: " + N);
		double perplexity = Math.pow(Math.E, (-1.0 / N) * logProbT);
		System.out.println("Perplexity: over " + N
				+ " recognized n-grams in verification corpus: " + perplexity);
	}

	private void generateModel(List<Headline> trainingSet, boolean genRandom,
			boolean useUnk, Map<String,Map<Date,StockPrice>> stockTrends) throws FileNotFoundException, IOException {
		//List<String> fileByLines = new ArrayList<String>();

		for (Headline headline : trainingSet) {
			String headlineText = headline.getText();
			if (headlineText.length() == 0) {
				continue;
			}
			String sanitizedLine = sanitize(headline.getText());
			// split on whitespace
			String[] tokens = sanitizedLine.toLowerCase().split("\\s+");
			
			StringBuilder currentLine = new StringBuilder();
			
			for (String token : tokens) {
				if (!isWord(token)) {
					continue;
				}

				String word;
				if (!wordsSeenOnce.contains(token) && useUnk) {
					word = UNK;
					wordsSeenOnce.add(token);
				} else {
					words.add(token);
					word = token;
				}

				if (END.equals(word)) {
					currentLine.append(word);
					//fileByLines.add(currentLine.toString());
					//currentLine = new StringBuilder();
				} else {
					currentLine.append(word);
					currentLine.append(" ");
				}
			}

			String str = currentLine.toString();
			System.out.println(str);
			NGram currentNgram = new NGram(MAX_N_GRAM_LENGTH);
			for (String token : str.split("\\s+")) {
				currentNgram.add(token);
				for (int i = 0; i <= currentNgram.size(); i++) {
					addNGram(currentNgram.size() - i,
							currentNgram.subNGram(i, currentNgram.size()), headline.getStock(), headline.getDate(), stockTrends);
					totalNGramCounts[currentNgram.size() - i]++;
				}
			}
		}

		DecimalFormat decFmt = new DecimalFormat("###0.0000");
		for (int modelIndex = 1; modelIndex <= 3; modelIndex++) {
			File file = new File(modelIndex + "grams.txt");
			BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(file)));
			List<Entry<NGram, Integer>> ngrams = new ArrayList<Entry<NGram, Integer>>(
					nGrams.get(modelIndex).entrySet());
			Collections.sort(ngrams, new NGramComparator());
			System.out.println("Highest frequency " + modelIndex + "-grams:");
			for (int i = 1; i <= 10; i++) {
				System.out
						.println(i
								+ ". "
								+ ngrams.get(i - 1).getKey()
								+ " : "
								+ (((double) (ngrams.get(i - 1).getValue()) / totalNGramCounts[1])));
			}
			Map<NGram,Double> pricesForModel = nGramPriceAvg.get(modelIndex);
			for (int nGramIndex = 1; nGramIndex <= ngrams.size(); nGramIndex++) {
				NGram key = ngrams.get(nGramIndex - 1).getKey();
				writer.write(key.toString());
				writer.write(",");
				int count = ngrams.get(nGramIndex - 1).getValue();
				writer.write(Integer.toString(count));
				writer.write(",");
				double avgPrice;
				try {
					avgPrice = pricesForModel.get(key);
					System.out.println("Avg price for " + modelIndex + "-gram " + key +": " + avgPrice);
				} catch (NullPointerException npe) {
					System.out.println("null avgPrice for " + modelIndex + "-gram " + key);
					avgPrice = 0.0;
				}
				writer.write(decFmt.format(avgPrice/(double)count));
				writer.write("\n");
			}
			try {
				writer.close();
			} catch (IOException ioe) {
				System.out.println(ioe.getMessage());
			}
		}
		if (genRandom) {
			for (int nGramLength = 1; nGramLength <= MAX_N_GRAM_LENGTH; nGramLength++) {
				System.out.println("Random sentence of length " + RANDOM_LENGTH
						+ " using " + nGramLength + "-gram language model:");
				StringBuilder randomText = new StringBuilder();

				NGram randomNminusOneGram = new NGram(nGramLength);
				for (int i = 0; i < RANDOM_LENGTH; i++) {
					String randomToken = getRandomToken(nGrams,
							randomNminusOneGram.size() + 1, randomNminusOneGram);
					NGram randomNGram = randomNminusOneGram.copy();
					randomNGram.add(randomToken);
					randomText.append(randomNGram.get(randomNGram.size() - 1));
					randomText.append(" ");

					if (randomNGram.size() < nGramLength) {
						randomNminusOneGram = randomNGram;
					} else {
						randomNminusOneGram = randomNGram.subNGram(1);
					}
				}
				System.out.println(randomText);
				System.out.println();
			}
		}
	}

	private boolean isWord(String aString) {
		if (aString == null || aString.length() == 0) {
			return false;
		}

		if (nonWords.contains(aString)) {
			return false;
		}

		if (START.equals(aString) || END.equals(aString) || UNK.equals(aString)) {
			return true;
		}

		Matcher wordMatcher = wordPattern.matcher(aString);
		if (wordMatcher.find()) {
			return true;
		}
		nonWords.add(aString);
		return false;
	}

	public void reportModel(List<Headline> trainingSet,
			List<Headline> validationSet, Map<String,Map<Date,StockPrice>> stockTrends) {
		try {
			NGramModel ngm = new NGramModel();
			boolean doCalcPerplexity = true;
			ngm.generateModel(trainingSet, !doCalcPerplexity, doCalcPerplexity, stockTrends);
			if (doCalcPerplexity) {
				for (int i = 1; i <= MAX_N_GRAM_LENGTH; i++) {
					ngm.calcPerplexity(validationSet, i, true);
				}
			}
		} catch (FileNotFoundException fnfe) {
			fnfe.printStackTrace(System.err);
		} catch (IOException ioe) {
			ioe.printStackTrace(System.err);
		}
	}

	private class NGramComparator implements Comparator<Entry<NGram, Integer>> {

		@Override
		public int compare(Entry<NGram, Integer> o1, Entry<NGram, Integer> o2) {
			return o2.getValue().compareTo(o1.getValue());
		}

	}
}