# -*- coding: utf-8 -*-
"""
	Aho-Corasick string search algorithm.

	Author    : Wojciech Muła, wojciech_mula@poczta.onet.pl
	WWW       : http://0x80.pl
	License   : public domain
"""

from collections import deque

nil = object()	# used to distinguish from None

class TrieNode(object):
	"""
	Node of trie/Aho-Corasick automaton
	"""

	__slots__ = ['char', 'output', 'fail', 'children']

	def __init__(self, char):
		"""
		Constructs an empty node
		"""

		self.char = char		# character
		self.output = nil		# an output function for this node
		self.fail = nil			# fail link used by Aho-Corasick automaton
		self.children = {}		# children

	def __repr__(self):
		"""
		Textual representation of node.
		"""

		if self.output is not nil:
			return "<TrieNode '%s' '%s'>" % (self.char, self.output)
		else:
			return "<TrieNode '%s'>" % self.char


class Trie(object):
	"""
	Trie/Aho-Corasick automaton.
	"""

	def __init__(self):
		"""
		Construct an empty trie
		"""

		self.root = TrieNode('')


	def __get_node(self, word):
		"""
		Private function retrieving a final node of trie
		for given word

		Returns node or None, if the trie doesn't contain the word.
		"""

		node = self.root
		for c in word:
			try:
				node = node.children[c]
			except KeyError:
				return None

		return node


	def get(self, word, default=nil):
		"""
		Retrieves output value associated with word.

		If there is no word returns default value,
		and if default is not given rises KeyError.
		"""

		node = self.__get_node(word)
		output = nil
		if node:
			output = node.output

		if output is nil:
			if default is nil:
				raise KeyError("no key '%s'" % word)
			else:
				return default
		else:
			return output


	def keys(self):
		"""
		Generator returning all keys (i.e. word) stored in trie
		"""

		for key, _ in self.items():
			yield key


	def values(self):
		"""
		Generator returning all values associated with words stored in a trie.
		"""

		for _, value in self.items():
			yield value


	def items(self):
		"""
		Generator returning all keys and values stored in a trie.
		"""

		L = []
		def aux(node, s):
			s = s + node.char
			if node.output is not nil:
				L.append((s, node.output))

			for child in node.children.values():
				if child is not node:
					aux(child, s)

		aux(self.root, '')
		return iter(L)


	def __len__(self):
		"""
		Calculates number of words in a trie.
		"""

		stack = deque()
		stack.append(self.root)
		n = 0
		while stack:
			node = stack.pop()
			if node.output is not nil:
				n += 1

			for child in node.children.values():
				stack.append(child)

		return n


	def add_word(self, word, value):
		"""
		Adds word and associated value.

		If word already exists, its value is replaced.
		"""
		if not word:
			return

		node = self.root
		for c in word:
			try:
				node = node.children[c]
			except KeyError:
				n = TrieNode(c)
				node.children[c] = n
				node = n

		node.output = value


	def clear(self):
		"""
		Clears trie.
		"""

		self.root = TrieNode('')


	def exists(self, word):
		"""
		Checks if whole word is present in the trie.
		"""

		node = self.__get_node(word)
		if node:
			return bool(node.output != nil)
		else:
			return False


	def match(self, word):
		"""
		Checks if word is a prefix of any existing word in the trie.
		"""

		return (self.__get_node(word) is not None)


	def make_automaton(self):
		"""
		Converts trie to Aho-Corasick automaton.
		"""

		queue = deque()

		# 1.
		for i in range(256):
			c = chr(i)
			if c in self.root.children:
				node = self.root.children[c]
				node.fail = self.root	# f(s) = 0
				queue.append(node)
			else:
				self.root.children[c] = self.root

		# 2.
		while queue:
			r = queue.popleft()
			for node in r.children.values():
				queue.append(node)
				state = r.fail
				while node.char not in state.children:
						state = state.fail

				node.fail = state.children.get(node.char, self.root)


	def iter(self, string):
		"""
		Generator performs Aho-Corasick search string algorithm, yielding
		tuples containing two values:
		- position in string
		- outputs associated with matched strings
		"""
		state = self.root
		for index, c in enumerate(string):
			while c not in state.children:
				state = state.fail

			state = state.children.get(c, self.root)

			tmp = state
			output = []
			while tmp is not nil:
				if tmp.output is not nil:
					output.append(tmp.output)

				tmp = tmp.fail

			if output:
				yield (index, output)

	def iter_long(self, string):
		"""
		Generator performs a modified Aho-Corasick search string algorithm,
		which maches only the longest word.

		"""
		state = self.root
		last  = None

		index = 0
		while index < len(string):
			c = string[index]

			if c in state.children:
				state = state.children[c]

				if state.output is not nil:
					# save the last node on the path
					last = (state.output, index)

				index += 1
			else:
				if last:
					# return the saved match
					yield last

					# and start over, as we don't want overlapped results
					# Note: this leads to quadratic complexity in the worst case
					index = last[1] + 1
					state = self.root
					last  = None
				else:
					# if no output, perform classic Aho-Corasick algorithm
					while c not in state.children:
						state = state.fail

		# corner case
		if last:
			yield last

	def find_all(self, string, callback):
		"""
		Wrapper on iter method, callback gets an iterator result
		"""
		for index, output in self.iter(string):
			callback(index, output)



if __name__ == '__main__':

	def demo():
		words = "he hers his she hi him man".split()

		t = Trie();
		for w in words:
			t.add_word(w, w)

		s = "he rshershidamanza "

		t.make_automaton()
		for res in t.items():
			print(res)

		for res in t.iter(s):
			print
			print('%s' % s)
			pos, matches = res
			for fragment in matches:
				print('%s%s' % ((pos - len(fragment) + 1)*' ', fragment))

	demo()


	def bug():
		patterns = ['GT-C3303','SAMSUNG-GT-C3303K/']
		text = 'SAMSUNG-GT-C3303i/1.0 NetFront/3.5 Profile/MIDP-2.0 Configuration/CLDC-1.1'

		t = Trie()
		for pattern in patterns:
			ret = t.add_word(pattern, (0, pattern))

		t.make_automaton()

		res = list(t.iter(text))

		assert len(res) == 1, 'failed'

	bug()

# vim: ts=4 sw=4 nowrap