#!/usr/bin/env python
import argparse
import sys
import models
import heapq
from collections import namedtuple

parser = argparse.ArgumentParser(description='Simple phrase based decoder.')
parser.add_argument('-i', '--input', dest='input', default='data/input', help='File containing sentences to translate (default=data/input)')
parser.add_argument('-t', '--translation-model', dest='tm', default='data/tm', help='File containing translation model (default=data/tm)')
parser.add_argument('-s', '--stack-size', dest='s', default=1, type=int, help='Maximum stack size (default=1)')
parser.add_argument('-n', '--num_sentences', dest='num_sents', default=sys.maxint, type=int, help='Number of sentences to decode (default=no limit)')
parser.add_argument('-l', '--language-model', dest='lm', default='data/lm', help='File containing ARPA-format language model (default=data/lm)')
parser.add_argument('-v', '--verbose', dest='verbose', action='store_true', default=False,  help='Verbose mode (default=off)')
opts = parser.parse_args()

tm = models.TM(opts.tm, sys.maxint)
lm = models.LM(opts.lm)
sys.stderr.write('Decoding %s...\n' % (opts.input,))
input_sents = [tuple(line.strip().split()) for line in open(opts.input).readlines()[:opts.num_sents]]

hypothesis = namedtuple('hypothesis', 'logprob, lm_state, predecessor, phrase')
for sent in input_sents:
    # The following code implements a DP monotone decoding
    # algorithm (one that doesn't permute the target phrases).
    # Hence all hypotheses in stacks[i] represent translations of 
    # the first i words of the input sentence.
    # HINT: Generalize this so that stacks[i] contains translations
    # of any i words (remember to keep track of which words those
    # are, and to estimate future costs)
    initial_hypothesis = hypothesis(0.0, lm.begin(), None, None)
    stacks = [{} for _ in sent] + [{}]
    stacks[0][lm.begin()] = initial_hypothesis
    for position, stack in enumerate(stacks[:-1]):
        for h in heapq.nlargest(opts.s, stack.itervalues(), key=lambda h: h.logprob): # prune
            for end_span in xrange(position+1,len(sent)+1):
                for other_span in range(end_span+1,len(sent)+1):
                    if sent[end_span:other_span] in tm and sent[position:end_span] in  tm:
                        switch_poss = []
                        for phrase in tm[sent[end_span:other_span]]:
                            logprob = h.logprob + phrase.logprob
                            lm_state = h.lm_state
                            for word in phrase.english.split():
                                (lm_state, word_logprob) = lm.score(lm_state, word)
                                logprob += word_logprob
                            logprob += lm.end(lm_state) if end_span == len(sent) else 0.0
                            new_hypothesis = hypothesis(logprob, lm_state, h, phrase)
                            switch_poss.append(new_hypothesis)
                        #print "newstack", switch_poss
                        for poss in switch_poss:
                            for phrase in tm[sent[position:end_span]]:
                                logprob = poss.logprob + phrase.logprob
                                lm_state = poss.lm_state
                                for word in phrase.english.split():
                                    (lm_state, word_logprob) = lm.score(lm_state, word)
                                    logprob += word_logprob
                                logprob += lm.end(lm_state) if other_span == len(sent) else 0.0
                                new_hypothesis = hypothesis(logprob, lm_state, poss, phrase)
                                if lm_state not in stacks[other_span] or stacks[other_span][lm_state].logprob < logprob: # second case is recombination
                                    stacks[other_span][lm_state] = new_hypothesis
        for h in heapq.nlargest(opts.s, stack.itervalues(), key=lambda h: h.logprob): # prune
            for end_span in xrange(position+1,len(sent)+1):
                if sent[position:end_span] in tm:
                    for phrase in tm[sent[position:end_span]]:
                        logprob = h.logprob + phrase.logprob
                        lm_state = h.lm_state
                        for word in phrase.english.split():
                            (lm_state, word_logprob) = lm.score(lm_state, word)
                            logprob += word_logprob
                        logprob += lm.end(lm_state) if end_span == len(sent) else 0.0
                        new_hypothesis = hypothesis(logprob, lm_state, h, phrase)
                        if lm_state not in stacks[end_span] or stacks[end_span][lm_state].logprob < logprob: # second case is recombination
                            stacks[end_span][lm_state] = new_hypothesis

    # find best translation by looking at the best scoring hypothesis
    # on the last stack
    winner = max(stacks[-1].itervalues(), key=lambda h: h.logprob)
    def extract_english_recursive(h):
        return '' if h.predecessor is None else '%s%s ' % (extract_english_recursive(h.predecessor), h.phrase.english)
    print extract_english_recursive(winner)

    if opts.verbose:
        def extract_tm_logprob(h):
            return 0.0 if h.predecessor is None else h.phrase.logprob + extract_tm_logprob(h.predecessor)
        tm_logprob = extract_tm_logprob(winner)
        sys.stderr.write('LM = %f, TM = %f, Total = %f\n' % 
            (winner.logprob - tm_logprob, tm_logprob, winner.logprob))
