package dsa;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

public class Trie {

	private TrieNode root;

	public Trie() {
		root = new TrieNode(' ', null);
	}


	/**
	 * Inserts a word into the Trie.
	 *
	 * @param word
	 */
	public void insert(String word) {
		Map<Character, TrieNode> children = root.childern;

		// Keep track of the parent node while we traverse through the trie.
		// Parent node info is required while creating a new TrieNode.
		TrieNode parentNode = root;
		for(int i=0; i< word.length(); i++) {
			char currentChar = word.charAt(i);

			// Traverse through the Trie to add missing nodes for the word
			// that is getting inserted.
			TrieNode currentNode;
			if (children.containsKey(currentChar)) {
				currentNode = children.get(currentChar);
			} else {
				currentNode = new TrieNode(currentChar, parentNode);
				children.put(currentChar, currentNode);
			}

			// Make the current node as parent node and process it's children.
			parentNode = currentNode;
			children = currentNode.childern;

			// If the current character is the last character,
			// mark the current node as a valid Word in the Trie
			if (i == word.length()-1) {
				currentNode.isEndOfWord = true;
			}
		}
	}

	/**
	 * Displays all words present in a Trie.
	 */
	public void displayWords(TrieNode currentNode, String str) {
		if (currentNode.isEndOfWord) {
			System.out.println(str);
		}

		Map<Character, TrieNode> children = currentNode.childern;

		for(Entry<Character, TrieNode> entry : children.entrySet()) {
			displayWords(entry.getValue(), str + entry.getKey());
		}
	}

	/**
	 *  Returns a List of suggestions based on the prefix.
	 *
	 * @param prefix
	 * @return List containing words that starts with prefix.
	 */
	public List<String> suggest(String prefix) {
		List<String> words = new ArrayList<String>();

		// Locate the prefix in the Trie
		TrieNode t = root;
		boolean prefixFound = true;
		for(int i=0; i< prefix.length(); i++) {
			char c = prefix.charAt(i);

			t = t.childern.get(c);
			if (t == null) {
				prefixFound = false;
				break;
			}
		}

		// Collect all child words for the prefix.
		if (prefixFound) {
			collectChildWords(t, words);
			Collections.sort(words);
		}

		return words;
	}

	private void collectChildWords(TrieNode currentNode, List<String> words) {
		if (currentNode.isEndOfWord) {
			words.add(currentNode.toString());
		}

		Map<Character, TrieNode> children = currentNode.childern;
		for(Entry<Character, TrieNode> entry : children.entrySet()) {
			collectChildWords(entry.getValue(), words);
		}
	}


	public static void main(String[] args) {
		Trie trie = new Trie();
		trie.insert("dog");
		trie.insert("damp");
		trie.insert("dance");
		trie.insert("damage");

		trie.insert("cat");
		trie.insert("camera");
		trie.insert("catch");
		trie.insert("camp");

		trie.displayWords(trie.root, "");

		List<String> suggestions = trie.suggest("ca");

		System.out.println(suggestions);
	}


	static class TrieNode {
		char c;
		Map<Character, TrieNode> childern = new HashMap<>();
		boolean isEndOfWord;
		// Having a parent pointer helps us to implement toString()
		TrieNode parent;

		TrieNode(char c, TrieNode parent) {
			this.c = c;
			this.parent = parent;
		}

		public String toString() {
			if (parent == null) {
				return "";
			}

			return parent.toString() + c;
		}
	}
}