/**
 * 
 */
package edu.berkeley.nlp.PCFGLA;

import java.util.Collection;

import edu.berkeley.nlp.syntax.StateSet;
import edu.berkeley.nlp.syntax.Tree;
import edu.berkeley.nlp.util.Counter;
import edu.berkeley.nlp.util.CounterMap;
import edu.berkeley.nlp.util.Numberer;
import edu.berkeley.nlp.util.PriorityQueue;

/**
 * CorpusStatistics calculates symbol counts for a corpus.
 * 
 * @author leon
 * 
 */
public class CorpusStatistics {
	// i dont know how to initialize shorts...
	short zero = 0, one = 1;

	int[] counts;
	Collection<Tree<StateSet>> trees;
	Counter<UnaryRule> unaryRuleCounter;
	Counter<BinaryRule> binaryRuleCounter;

	int[] contexts;
	CounterMap<Integer, String> posCounter;

	/**
	 * Count statistics for a collection of StateSet trees.
	 */
	public CorpusStatistics(Numberer tagNumberer,
			Collection<Tree<StateSet>> trees) {
		counts = new int[tagNumberer.objects().size()];
		this.trees = trees;
		unaryRuleCounter = new Counter<UnaryRule>();
		binaryRuleCounter = new Counter<BinaryRule>();
		contexts = new int[tagNumberer.objects().size()];
		posCounter = new CounterMap<Integer, String>();
	}

	public void countSymbols() {
		for (Tree<StateSet> tree : trees) {
			addCount(tree);
		}
	}

	private void addCount(Tree<StateSet> tree) {
		counts[tree.getLabel().getState()] += 1.0;
		if (!tree.isPreTerminal()) {
			for (Tree<StateSet> child : tree.getChildren()) {
				addCount(child);
			}
		}
	}

	/*
	 * Counts how many different 'things' (non-terminals or terminals for the
	 * POS) appear under a given nonterminal symbol. Currently POS and other
	 * nonterminals are handled the same way. We might to change that.
	 */

	public void countRuleParents() {
		for (Tree<StateSet> tree : trees) {
			addParent(tree);
		}
		for (BinaryRule br : binaryRuleCounter.keySet()) {
			contexts[br.parentState]++;
			contexts[br.leftChildState]++;
			contexts[br.rightChildState]++;
		}
		for (UnaryRule ur : unaryRuleCounter.keySet()) {
			contexts[ur.parentState]++;
			contexts[ur.childState]++;
		}
		for (int i = 0; i < contexts.length; i++) {
			Counter<String> tempC = posCounter.getCounter(i);
			contexts[i] += tempC.size();

		}
	}

	public int[] getContextCounts() {
		return contexts;
	}

	private void addParent(Tree<StateSet> tree) {
		short parentState = tree.getLabel().getState();
		counts[parentState] += 1.0;
		if (!tree.isPreTerminal()) {
			if (tree.getChildren().size() == 1) {
				UnaryRule r = new UnaryRule(parentState, tree.getChildren()
						.get(0).getLabel().getState(), new double[1][1]);
				unaryRuleCounter.incrementCount(r, 1.0);
			} else {
				BinaryRule r = new BinaryRule(parentState, tree.getChildren()
						.get(0).getLabel().getState(), tree.getChildren()
						.get(1).getLabel().getState(), new double[1][1][1]);
				binaryRuleCounter.incrementCount(r, 1.0);
			}
			for (Tree<StateSet> child : tree.getChildren()) {
				addParent(child);
			}
		} else {
			posCounter.incrementCount((int) parentState, tree.getChildren()
					.get(0).getLabel().getWord(), 1.0);
		}
	}

	/**
	 * Get the number of times each state appeared.
	 * 
	 * @return
	 */
	public int[] getSymbolCounts() {
		countSymbols();
		return counts;
	}

	public void printStateCountArray(Numberer tagNumberer, int[] array) {
		PriorityQueue<String> pq = new PriorityQueue<String>(array.length);
		for (int i = 0; i < array.length; i++) {
			pq.add((String) tagNumberer.object(i), array[i]);
			// System.out.println(i+". "+(String)tagNumberer.object(i)+"\t "+symbolCounter.getCount(i,0));
		}
		int i = 0;
		while (pq.hasNext()) {
			i++;
			int p = (int) pq.getPriority();
			System.out.println(i + ". " + pq.next() + "\t " + p);
		}
	}

}
