package de.ugoe.cs.eventbench.models;

import java.io.Serializable;
import java.util.Collection;
import java.util.HashSet;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import de.ugoe.cs.util.StringTools;

import edu.uci.ics.jung.graph.DelegateTree;
import edu.uci.ics.jung.graph.Graph;
import edu.uci.ics.jung.graph.Tree;

/**
 * <p>
 * This class implements a <it>trie</it>, i.e., a tree of sequences that the
 * occurence of subsequences up to a predefined length. This length is the trie
 * order.
 * </p>
 * 
 * @author Steffen Herbold
 * 
 * @param <T>
 *            Type of the symbols that are stored in the trie.
 * 
 * @see TrieNode
 */
public class Trie<T> implements IDotCompatible, Serializable {

	/**
	 * <p>
	 * Id for object serialization.
	 * </p>
	 */
	private static final long serialVersionUID = 1L;

	/**
	 * <p>
	 * Collection of all symbols occuring in the trie.
	 * </p>
	 */
	private Collection<T> knownSymbols;

	/**
	 * <p>
	 * Reference to the root of the trie.
	 * </p>
	 */
	private final TrieNode<T> rootNode;

	/**
	 * <p>
	 * Contructor. Creates a new Trie.
	 * </p>
	 */
	public Trie() {
		rootNode = new TrieNode<T>();
		knownSymbols = new LinkedHashSet<T>();
	}

	/**
	 * <p>
	 * Returns a collection of all symbols occuring in the trie.
	 * </p>
	 * 
	 * @return symbols occuring in the trie
	 */
	public Collection<T> getKnownSymbols() {
		return new LinkedHashSet<T>(knownSymbols);
	}

	/**
	 * <p>
	 * Trains the current trie using the given sequence and adds all subsequence
	 * of length {@code maxOrder}.
	 * </p>
	 * 
	 * @param sequence
	 *            sequence whose subsequences are added to the trie
	 * @param maxOrder
	 *            maximum length of the subsequences added to the trie
	 */
	public void train(List<T> sequence, int maxOrder) {
		if (maxOrder < 1) {
			return;
		}
		IncompleteMemory<T> latestActions = new IncompleteMemory<T>(maxOrder);
		int i = 0;
		for (T currentEvent : sequence) {
			latestActions.add(currentEvent);
			knownSymbols.add(currentEvent);
			i++;
			if (i >= maxOrder) {
				add(latestActions.getLast(maxOrder));
			}
		}
		int sequenceLength = sequence.size();
		for (int j = maxOrder - 1; j > 0; j--) {
			add(sequence.subList(sequenceLength - j, sequenceLength));
		}
	}

	/**
	 * <p>
	 * Adds a given subsequence to the trie and increases the counters
	 * accordingly.
	 * </p>
	 * 
	 * @param subsequence
	 *            subsequence whose counters are increased
	 * @see TrieNode#add(List)
	 */
	protected void add(List<T> subsequence) {
		if (subsequence != null && !subsequence.isEmpty()) {
			knownSymbols.addAll(subsequence);
			subsequence = new LinkedList<T>(subsequence); // defensive copy!
			T firstSymbol = subsequence.get(0);
			TrieNode<T> node = getChildCreate(firstSymbol);
			node.add(subsequence);
		}
	}

	/**
	 * <p>
	 * Returns the child of the root node associated with the given symbol or
	 * creates it if it does not exist yet.
	 * </p>
	 * 
	 * @param symbol
	 *            symbol whose node is required
	 * @return node associated with the symbol
	 * @see TrieNode#getChildCreate(Object)
	 */
	protected TrieNode<T> getChildCreate(T symbol) {
		return rootNode.getChildCreate(symbol);
	}

	/**
	 * <p>
	 * Returns the child of the root node associated with the given symbol or
	 * null if it does not exist.
	 * </p>
	 * 
	 * @param symbol
	 *            symbol whose node is required
	 * @return node associated with the symbol; null if no such node exists
	 * @see TrieNode#getChild(Object)
	 */
	protected TrieNode<T> getChild(T symbol) {
		return rootNode.getChild(symbol);
	}

	/**
	 * <p>
	 * Returns the number of occurences of the given sequence.
	 * </p>
	 * 
	 * @param sequence
	 *            sequence whose number of occurences is required
	 * @return number of occurences of the sequence
	 */
	public int getCount(List<T> sequence) {
		int count = 0;
		TrieNode<T> node = find(sequence);
		if (node != null) {
			count = node.getCount();
		}
		return count;
	}

	/**
	 * <p>
	 * Returns the number of occurences of the given prefix and a symbol that
	 * follows it.<br>
	 * Convenience function to simplify usage of {@link #getCount(List)}.
	 * </p>
	 * 
	 * @param sequence
	 *            prefix of the sequence
	 * @param follower
	 *            suffix of the sequence
	 * @return number of occurences of the sequence
	 * @see #getCount(List)
	 */
	public int getCount(List<T> sequence, T follower) {
		List<T> tmpSequence = new LinkedList<T>(sequence);
		tmpSequence.add(follower);
		return getCount(tmpSequence);

	}

	/**
	 * <p>
	 * Searches the trie for a given sequence and returns the node associated
	 * with the sequence or null if no such node is found.
	 * </p>
	 * 
	 * @param sequence
	 *            sequence that is searched for
	 * @return node associated with the sequence
	 * @see TrieNode#find(List)
	 */
	public TrieNode<T> find(List<T> sequence) {
		if (sequence == null || sequence.isEmpty()) {
			return rootNode;
		}
		List<T> sequenceCopy = new LinkedList<T>(sequence);
		TrieNode<T> result = null;
		TrieNode<T> node = getChild(sequenceCopy.get(0));
		if (node != null) {
			sequenceCopy.remove(0);
			result = node.find(sequenceCopy);
		}
		return result;
	}

	/**
	 * <p>
	 * Returns a collection of all symbols that follow a given sequence in the
	 * trie. In case the sequence is not found or no symbols follow the sequence
	 * the result will be empty.
	 * </p>
	 * 
	 * @param sequence
	 *            sequence whose followers are returned
	 * @return symbols following the given sequence
	 * @see TrieNode#getFollowingSymbols()
	 */
	public Collection<T> getFollowingSymbols(List<T> sequence) {
		Collection<T> result = new LinkedList<T>();
		TrieNode<T> node = find(sequence);
		if (node != null) {
			result = node.getFollowingSymbols();
		}
		return result;
	}

	/**
	 * <p>
	 * Returns the longest suffix of the given context that is contained in the
	 * tree and whose children are leaves.
	 * </p>
	 * 
	 * @param context
	 *            context whose suffix is searched for
	 * @return longest suffix of the context
	 */
	public List<T> getContextSuffix(List<T> context) {
		List<T> contextSuffix = new LinkedList<T>(context); // defensive copy
		boolean suffixFound = false;

		while (!suffixFound) {
			if (contextSuffix.isEmpty()) {
				suffixFound = true; // suffix is the empty word
			} else {
				TrieNode<T> node = find(contextSuffix);
				if (node != null) {
					if (!node.getFollowingSymbols().isEmpty()) {
						suffixFound = true;
					}
				}
				if (!suffixFound) {
					contextSuffix.remove(0);
				}
			}
		}

		return contextSuffix;
	}

	/**
	 * <p>
	 * Helper class for graph visualization of a trie.
	 * </p>
	 * 
	 * @author Steffen Herbold
	 * @version 1.0
	 */
	static public class Edge {
	}

	/**
	 * <p>
	 * Helper class for graph visualization of a trie.
	 * </p>
	 * 
	 * @author Steffen Herbold
	 * @version 1.0
	 */
	static public class TrieVertex {

		/**
		 * <p>
		 * Id of the vertex.
		 * </p>
		 */
		private String id;

		/**
		 * <p>
		 * Contructor. Creates a new TrieVertex.
		 * </p>
		 * 
		 * @param id
		 *            id of the vertex
		 */
		protected TrieVertex(String id) {
			this.id = id;
		}

		/**
		 * <p>
		 * Returns the id of the vertex.
		 * </p>
		 * 
		 * @see java.lang.Object#toString()
		 */
		@Override
		public String toString() {
			return id;
		}
	}

	/**
	 * <p>
	 * Returns a {@link Graph} representation of the trie.
	 * </p>
	 * 
	 * @return {@link Graph} representation of the trie
	 */
	protected Tree<TrieVertex, Edge> getGraph() {
		DelegateTree<TrieVertex, Edge> graph = new DelegateTree<TrieVertex, Edge>();
		rootNode.getGraph(null, graph);
		return graph;
	}

	/*
	 * (non-Javadoc)
	 * 
	 * @see de.ugoe.cs.eventbench.models.IDotCompatible#getDotRepresentation()
	 */
	public String getDotRepresentation() {
		StringBuilder stringBuilder = new StringBuilder();
		stringBuilder.append("digraph model {" + StringTools.ENDLINE);
		rootNode.appendDotRepresentation(stringBuilder);
		stringBuilder.append('}' + StringTools.ENDLINE);
		return stringBuilder.toString();
	}

	/**
	 * <p>
	 * Returns the string representation of the root node.
	 * </p>
	 * 
	 * @see TrieNode#toString()
	 * @see java.lang.Object#toString()
	 */
	@Override
	public String toString() {
		return rootNode.toString();
	}

	/**
	 * <p>
	 * Returns the number of symbols contained in the trie.
	 * </p>
	 * 
	 * @return number of symbols contained in the trie
	 */
	public int getNumSymbols() {
		return knownSymbols.size();
	}

	/**
	 * <p>
	 * Returns the number of trie nodes that are ancestors of a leaf. This is
	 * the equivalent to the number of states a first-order markov model would
	 * have.
	 * <p>
	 * 
	 * @return number of trie nodes that are ancestors of leafs.
	 */
	public int getNumLeafAncestors() {
		Set<TrieNode<T>> ancestors = new HashSet<TrieNode<T>>();
		rootNode.getLeafAncestors(ancestors);
		return ancestors.size();
	}

	/**
	 * <p>
	 * Returns the number of trie nodes that are leafs.
	 * </p>
	 * 
	 * @return number of leafs in the trie
	 */
	public int getNumLeafs() {
		return rootNode.getNumLeafs();
	}

	/**
	 * <p>
	 * Updates the list of known symbols by replacing it with all symbols that
	 * are found in the child nodes of the root node. This should be the same as
	 * all symbols that are contained in the trie.
	 * </p>
	 */
	public void updateKnownSymbols() {
		knownSymbols = new HashSet<T>();
		for (TrieNode<T> node : rootNode.getChildren()) {
			knownSymbols.add(node.getSymbol());
		}
	}
}
