#!/usr/bin/python
#lexicalised type driven semantic grammar. Andreas van Cranenburgh 2010.
import nltk
types = {'A': '<<e,t>,t>', 'B': '<<e,t>,t>', 'C': '<<e,t>,t>',
	'D': '<<<e,t>,t>, <<e,t>, t>>',
	'F': '<e,e>', 'wife' : '<e,e>',
	'L': '<e,<e,t>>', 'loves' : '<e,<e,t>>',
	'P': '<e,t>', 'Q': '<e,t>',
	'girl': '<e,t>', 'book': '<e,t>', 'boy': '<e,t>', 'smokes': '<e,t>',
	'x': 'e', 'y': 'e', 'z': 'e', 'bill' : 'e', 'john' : 'e'}

lexicon = {'every': r'\P Q.all x.(P(x) -> Q(x))',
	'girl' : r'\x.girl(x)',
	'book' : r'\x.book(x)',
	'boy' : r'\x.boy(x)',
	'with' : [r'\A P B.(A(\y.(B(\x.(using(x)(y) & P(x))))))',
				r'\A B P.(A(\y.(B(\x.(has(x)(y) & P(y))))))'],
	'smokes' : r'\x.smokes(x)',
	'gives' : [r'\B C A.C(\x.A(\z.(B(\y.(give(x)(y)(z))))))',
			r'\B C x.C(\z.(B(\y.(give(x)(y)(z)))))'],
	'mary' : r'\P.P(mary)',
	#'john' : [r'\P.P(john)', 'john'],
	'john' : ['john', '\P.P(john)'],
	'bill' : r'bill',
	'a' : r'\P Q.exists x.(P(x) & Q(x))',
	'hits' : r'\A x.A(\y.(hits(x)(y)))',
	'does' : r'\A.A',
	#'too' : r'\D.D', 
	'too' : r'\E.E',
	#	[r'\D.D', r'\E.E'],
	'does_too' : r'\x P.P(x)',
	'and' : r'\P x y.(P(x) & P(y))',
	'or' : r'\P x y.(P(x) | P(y))',
	'loves' : r'\x y.loves(x,y)',
	'his' : [r'\F L x.L(x, F(x))', r'\F L x y.L(y, F(x))'],
	'wife' : r'\x.wife(x)'
	}

grammar = nltk.parse_cfg("""
S -> NP VP | SCon S
SCon -> NP VPCon
PP -> P NP
NP -> DET N | NP PP | PN
VPCon -> VP CON
VP -> TV NP | IV | VP PP | DVP NP | NP VP
DVP -> DTV NP
CON -> 'and'
IV -> 'smokes' | 'does' 'too' | 'does_too'
TV -> 'loves' | 'hits'
DTV -> 'gives'
PN -> 'mary' | 'john' | 'bill'
N -> 'girl' | 'boy' | 'wife' | 'book'
DET -> 'his' | 'every' | 'a'
P -> 'with'
""")
grammar = nltk.parse_cfg("""
    S -> NP VP
    S -> SCon S
    SCon -> NP VPCon
    PP -> 'with' NP
    NP -> 'every' 'girl'
    NP -> 'a' 'book'
    NP -> 'a' 'boy'
    NP -> 'his' 'wife'
    NPO -> 'or' NP
    NP -> NP NPO
    NP -> NP PP
    VPCon -> 'smokes' CON
    VPCon -> VP CON
    VP -> 'hits' NP
    VP -> 'loves' NP
    VP -> 'does' 'too'
    VP -> 'does_too'
    VP -> VP PP
    VP -> DVP NP
    VP -> NP VP
    DVP -> 'gives' NP
    CON -> 'and'
    NP -> 'mary'
    NP -> 'john'
    NP -> 'bill'
""")

sr = nltk.parse.EarleyChartParser(grammar)
lp = nltk.sem.logic.LogicParser(True) #typed logic
#lp = nltk.sem.logic.LogicParser(not True) #typed logic
lexicon = dict((w, (type(m) == list and m or [m])) for w, m in lexicon.items())
lexicon = dict((w, [lp.parse(a, types) for a in m]) for w, m in lexicon.items())
print "the lexicon:"
for w, m in lexicon.items():
	print w,
	for a in m: print a, a.type

def cartpi(seq):
	""" fold a sequence using cartesian product as operator """
	if len(seq) == 1:
		for a in seq[0]: yield (a,)
	else:
		for a in seq[0]:
			for b in cartpi(seq[1:]): yield (a,) + b

def applylexicon(parsetree):
	#replace terminals with their lexical semantics
	for a in cartpi([lexicon[b] for b in parsetree.leaves()]):
		copy = parsetree.copy(True)
		for b, c in zip(a, copy.treepositions('leaves')):
			#copy[c] = nltk.Tree((copy[c], b), [])
			copy[c] = (copy[c], b)
		yield copy

def applytree(parsetree):
	#apply expressions to each other working up from bottom to root node
	#direction of application is determined by type matching.
	if type(parsetree) == nltk.tree.Tree:
		if len(parsetree) == 1:
			return nltk.Tree((parsetree.node, parsetree[0][1]), [parsetree[0]])
		elif len(parsetree) == 2:
			aa, bb = applytree(parsetree[0]), applytree(parsetree[1])
			a, b = aa.node[1], bb.node[1]
			if (type(a.type) == nltk.sem.logic.ComplexType
			and a.type.first.matches(b.type)):
				node = a(b).simplify()
			elif (type(b.type) == nltk.sem.logic.ComplexType
			and b.type.first.matches(a.type)):
				node = b(a).simplify()
			else:
				raise nltk.sem.logic.TypeException("type mismatch: %s / %s" % (a.type, b.type))
			return nltk.Tree((parsetree.node, node), [aa, bb])
	return nltk.Tree(parsetree, [])

def applytypes(parsetree):
	#decorate tree with the type of the expression at each node
	def addtype(node): return node + (node[1].type,)
	if type(parsetree) == tuple:
		return addtype(parsetree)
	if len(parsetree) == 0:
		return addtype(parsetree.node)
	return nltk.Tree(addtype(parsetree.node), [applytypes(a) for a in parsetree])

def tostring(parsetree):
	# converts tuples at each node of a tree to multiline strings
	if type(parsetree) == tuple:
		return '\n'.join(str(a) for a in parsetree)
	if type(parsetree.node) == str:
		return nltk.Tree(parsetree.node, [tostring(a) for a in parsetree])
	return nltk.Tree('\n'.join(str(a) for a in parsetree.node),
						[tostring(a) for a in parsetree])

def interpret(sentence, grammar=grammar, draw=False):
	sr = nltk.parse.EarleyChartParser(grammar)
	for a in sr.nbest_parse(sentence.split()):		#syntactic ambiguities
		for b in applylexicon(a):			#lexical ambiguities
			try:
				c = applytree(b)
				yield str(c.node[1]), tostring(applytypes(c))
				#yield str(c.node[1]), tostring(applytypes(c))
				#print tostring(applytypes(c))
				if draw: tostring(applytypes(c)).draw()
			except nltk.sem.logic.TypeException, msg:
				yield msg, ""
			#print c
def main():
	#interpret("every girl gives mary a book")
	#interpret("john hits a boy with a book")
	#interpret("john smokes and bill does_too")

	#interpret("john hits a boy with a book and mary does too")
	#interpret("mary gives every girl a book")
	#interpret("every girl gives a boy a book")
	#interpret("john with a book hits a boy")
	#interpret("john with a book hits mary with bill")
	#interpret("every girl gives a boy with a book a book with a boy")
	#interpret("john loves his wife and bill does too")
	for a in interpret("john loves his wife and bill does_too", draw=True): print a[0]
	for a in interpret("john john loves his wife and bill does_too", draw=True): print a[0]

if __name__ == '__main__': main()