""" Web interface to the disco-dop parser. Requires Flask.
Expects a series of grammars in subdirectories of grammar/ """
# Wishlist:
# - shortest derivation, SL-DOP, MPSD, &c.
# - arbitrary configuration of CTF stages;
#   should become class also used by runexp.py & parser.py
import os, re, cgi, sys, glob, gzip, heapq, time, string, random, codecs, logging
from functools import wraps
from operator import itemgetter
from flask import Flask, Markup, request, render_template, send_from_directory
from werkzeug.contrib.cache import SimpleCache

from discodop import treetransforms, disambiguation, coarsetofine,\
		lexicon, pcfg, plcfrs
from discodop.tree import Tree
from discodop.treedraw import DrawTree
from discodop.containers import Grammar

app = Flask(__name__)
morphtags = re.compile(
		r'\(([_*A-Z0-9]+)(?:\[[^ ]*\][0-9]?)?((?:-[_A-Z0-9]+)?(?:\*[0-9]+)? )')
limit = 40 # maximum sentence length
prunek = 5000 # number of PLCFRS derivations to use for DOP parsing
grammars = {}
backtransforms = {}
knownwords = {}

@app.route('/')
def main():
	""" Serve the main form. """
	return render_template('form.html', result=Markup(parse()))

@app.route('/parse')
def parse():
	""" Parse sentence and return a textual representation of a parse tree,
	in a HTML fragment. To be invoked by an AJAX call."""
	sent = request.args.get('sent', None)
	objfun = request.args.get('objfun', 'mpp')
	marg = request.args.get('marg', 'nbest')
	if not sent:
		return ''
	frags = nbest = None
	senttok = tokenize(sent)
	if not senttok or not 1 <= len(senttok) <= limit:
		return 'Sentence too long: %d words, maximum %d' % (len(senttok), limit)
	result = getparse(senttok, objfun, marg)
	if result == 'no parse':
		return "no parse!\n"
	(parsetrees, fragments, elapsed, msg1, msg2, msg3) = result
	tree, prob = parsetrees[0]
	app.logger.info('[%g] %s' % (prob, tree))
	tree = morphtags.sub(r'(\1\2', tree)
	tree = Tree.parse(tree, parse_leaf=int)
	treetransforms.unbinarize(tree)
	treetransforms.removefanoutmarkers(tree)
	result = Markup(DrawTree(tree, senttok).text(
			unicodelines=True, html=True))
	frags = Markup('\n\n'.join(
			DrawTree(Tree.parse(frag, parse_leaf=int), terminals
				).text(unicodelines=True, html=True)
			for frag, terminals in fragments))
	elapsed = 'CPU time elapsed: %s => %gs' % (
			' '.join('%gs' % a for a in elapsed), sum(elapsed))
	nbest = Markup('\n\n'.join('%d. [p=%g]\n%s' % (n + 1, prob,
				DrawTree(treetransforms.removefanoutmarkers(
					treetransforms.unbinarize(Tree.parse(morphtags.sub(
						r'(\1\2', tree), parse_leaf=int))),
					senttok).text(unicodelines=True, html=True))
				for n, (tree, prob) in enumerate(parsetrees)))
	info = Markup('\n'.join(('sentence length: %d; objfun=%s; marg=%s' % (
			len(senttok), objfun, marg), msg1, msg2, msg3, elapsed,
			'10 most probable parse trees:',
			'\n'.join('%d. [p=%g] %s' % (n + 1, prob, cgi.escape(tree))
				for n, (tree, prob) in enumerate(parsetrees)) + '\n')))
	return render_template('parse.html', sent=sent, result=result,
			frags=frags, nbest=nbest, info=info, randid=randid())

@app.route('/favicon.ico')
def favicon():
	""" Serve the favicon. """
	return send_from_directory(os.path.join(app.root_path, 'static'),
			'favicon.ico', mimetype='image/vnd.microsoft.icon')

def loadgrammars():
	""" Load grammars if necessary. """
	if grammars != {}:
		return
	for folder in glob.glob('grammars/*/'):
		_, lang = os.path.split(os.path.dirname(folder))
		app.logger.info('Loading grammar %r', lang)
		grammarlist = []
		for stagename in ('pcfg', 'plcfrs', 'dop'):
			rules = gzip.open("%s/%s.rules.gz" % (folder, stagename)).read()
			lexical = codecs.getreader('utf-8')(gzip.open("%s/%s.lex.gz" % (
					folder, stagename))).read()
			if stagename == 'pcfg':
				grammarlist.append(Grammar(rules, lexical,
						logprob=False, bitpar=True))
			else:
				grammarlist.append(Grammar(rules, lexical))
			assert grammarlist[-1].testgrammar(), stagename
			if stagename == 'plcfrs':
				_ = grammarlist[-1].getmapping(grammarlist[-2],
						striplabelre=re.compile(b'@.+$'),
						neverblockre=re.compile(b'.+}<'),
						splitprune=True, markorigin=True)
			elif stagename == 'dop':
				_ = grammarlist[-1].getmapping(grammarlist[-2],
						striplabelre=re.compile(b'@.+$'),
						neverblockre=re.compile(b'.+}<'),
						splitprune=False, markorigin=False)
				backtransforms[lang] = dict(enumerate(gzip.open(
						"%s/dop.backtransform.gz" % folder).read().splitlines()))
		grammars[lang] = grammarlist
		knownwords[lang] = {w for w in grammars[lang][0].lexical
				if not w.startswith("UNK")}
		app.logger.info('Grammar for %s loaded.' % lang)

def cached(timeout=3600):
	def decorator(f):
		f.cache = SimpleCache()
		@wraps(f)
		def decorated_function(*args, **kwargs):
			cache_key = args
			rv = f.cache.get(cache_key)
			if rv is not None:
				return rv
			rv = f(*args, **kwargs)
			f.cache.set(cache_key, rv, timeout=timeout)
			return rv
		return decorated_function
	return decorator

@cached(timeout=24 * 3600)
def getparse(senttok, objfun, marg):
	""" Do the actual parsing. """
	elapsed = []
	begin = time.clock()
	lang = guesslang(senttok)
	grammar = grammars[lang]
	knownword = knownwords[lang]
	backtransform = backtransforms[lang]
	unksent = [w if w in knownword
			else lexicon.unknownword4(w, n, knownword)
			for n, w in enumerate(senttok)]
	inside, outside, start, _ = pcfg.doinsideoutside(
			unksent, grammar[0], tags=None)
	elapsed.append(time.clock() - begin)
	begin = time.clock()
	if start:
		(whitelist, _, _, _, _) = coarsetofine.whitelistfromposteriors(
				inside, outside, start,
				grammar[0], grammar[1], 1e-5, True, True)
		elapsed.append(time.clock() - begin)
		begin = time.clock()
		chart, start, _ = plcfrs.parse(unksent, grammar[1],
				exhaustive=True, whitelist=whitelist,
				splitprune=True, markorigin=True)
		elapsed.append(time.clock() - begin)
		begin = time.clock()
	else:
		app.logger.warning('stage 1 fail')
	if start:
		whitelist, items = coarsetofine.prunechart(
				chart, start, grammar[1], grammar[2],
				prunek, False, False, False)
		elapsed.append(time.clock() - begin)
		msg1 = "PLCFRS items: %d; In %d-best derivations: %d" % (
				prunek, len(chart), items)
		begin = time.clock()
		chart, start, msg2 = plcfrs.parse(unksent, grammar[2],
				exhaustive=objfun != 'mpd', whitelist=whitelist,
				splitprune=False, markorigin=False)
		elapsed.append(time.clock() - begin)
		begin = time.clock()
	else:
		app.logger.warning('stage 2 fail')
	if start:
		parsetrees, msg3 = disambiguation.marginalize(objfun, chart,
				start, grammar[2], 10000,
				kbest=marg in ('nbest', 'both'),
				sample=marg in ('sample', 'both'),
				sent=unksent, tags=None, backtransform=backtransform)
		fragments = disambiguation.extractfragments(
				chart, start, grammar[2], backtransform)
		elapsed.append(time.clock() - begin)
		return (heapq.nlargest(10, parsetrees.items(), key=itemgetter(1)),
				fragments, elapsed, msg1, msg2, msg3)
	else:
		app.logger.warning('stage 3 fail')
		return 'no parse'

def randid():
	""" return a string with 6 random letters. """
	return ''.join(random.choice(string.ascii_letters)
		for _ in range(6))

# List of contractions adapted from Robert MacIntyre's tokenizer.
CONTRACTIONS2 = re.compile(
		"(?i)(?:%s)\b" % "|".join([
		r"(.)('ll|'re|'ve|n't|'s|'m|'d)",
		r"\b(can)(not)",
		r"\b(D)('ye)",
		r"\b(Gim)(me)",
		r"\b(Gon)(na)",
		r"\b(Got)(ta)",
		r"\b(Lem)(me)",
		r"\b(Mor)('n)",
		r"\b(T)(is)",
		r"\b(T)(was)",
		r"\b(Wan)(na)"]))
CONTRACTIONS3 = re.compile(r"(?i)\b(?:(Whad)(dd)(ya)|(Wha)(t)(cha))\b")

def tokenize(text):
	""" Adapted from nltk.tokenize.TreebankTokenizer. """
	text = CONTRACTIONS2.sub(r'\1 \2', text)
	text = CONTRACTIONS3.sub(r'\1 \2 \3', text)
	# Separate most punctuation
	text = re.sub(r"([^\w\.\'\-\/,&])", r' \1 ', text, flags=re.UNICODE)
	# Separate commas if they're followed by space; e.g., don't separate 2,500
	# Separate single quotes if they're followed by a space.
	text = re.sub(r"([,']\s)", r' \1', text)
	# Separate periods that come before newline or end of string.
	text = re.sub(r'\. *(\n|$)', ' . ', text)
	return tuple(text.split())

def guesslang(sent):
	""" simple heuristic: language that contains most words from input. """
	lang = max(knownwords, key=lambda x: len(knownwords[x] & set(sent)))
	app.logger.info('Lang: %s; Sent: %s' % (lang, ' '.join(sent)))
	return lang

if __name__ == '__main__':
	logging.basicConfig()
	for log in (logging.getLogger(), app.logger):
		log.setLevel(logging.DEBUG)
		log.handlers[0].setFormatter(logging.Formatter(
				fmt='%(asctime)s %(message)s', datefmt='%Y-%m-%d %H:%M:%S'))
	loadgrammars()
	app.run(debug=False, host='0.0.0.0')
