package edu.stanford.nlp.sentidev;

import java.util.List;

import edu.stanford.nlp.neural.rnn.RNNCoreAnnotations;
import edu.stanford.nlp.trees.MemoryTreebank;
import edu.stanford.nlp.trees.Tree;
import edu.stanford.nlp.util.CollectionUtils;
import edu.stanford.nlp.util.Filter;
import edu.stanford.nlp.util.Generics;

/**
 * In the Sentiment dataset converted to tree form, the labels on the
 * intermediate nodes are the sentiment scores and the leaves are the
 * text of the sentence.  This class provides routines to read a file
 * of those trees and attach the sentiment score as the GoldLabel
 * annotation.
 *
 * @author John Bauer
 */
public class SentimentUtils {
	private SentimentUtils() {} // static methods only

	public static void attachGoldLabels(Tree tree) {
		if (tree.isLeaf()) {
			return;
		}
		for (Tree child : tree.children()) {
			attachGoldLabels(child);
		}

		// In the sentiment data set, the node labels are simply the gold
		// class labels.  There are no categories encoded.
		RNNCoreAnnotations.setGoldClass(tree, Integer.valueOf(tree.label().value()));
	}

	/**
	 * Given a file name, reads in those trees and returns them as a List
	 */
	public static List<Tree> readTreesWithGoldLabels(String path) {
		List<Tree> trees = Generics.newArrayList();
		MemoryTreebank treebank = new MemoryTreebank();
		treebank.loadPath(path, null);
		for (Tree tree : treebank) {
			attachGoldLabels(tree);
			trees.add(tree);
		}
		return trees;
	}

	static final Filter<Tree> UNKNOWN_ROOT_FILTER = new Filter<Tree>() {
		public boolean accept(Tree tree) {
			int gold = RNNCoreAnnotations.getGoldClass(tree);
			return gold != -1;
		}
	};

	public static List<Tree> filterUnknownRoots(List<Tree> trees) {
		return CollectionUtils.filterAsList(trees, UNKNOWN_ROOT_FILTER);
	}

	public static String sentimentString(SentimentModel model, int sentiment) {
		String[] classNames = model.op.classNames;
		if (sentiment < 0 || sentiment > classNames.length) {
			return "Unknown sentiment label " + sentiment;
		}
		return classNames[sentiment];
	}
}
