#!/usr/bin/env python
# coding=utf-8
import argparse
import sys
import models
import heapq
from collections import namedtuple
from nltk.corpus import cess_esp as cess
from nltk import UnigramTagger as ut

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, tm_score, lm_score, lm_state, translation')

cess_tags = cess.tagged_sents()
uni = ut(cess_tags)

def score_lm(translation,isEnd):
    lm_score = 0.0
    lm_state = lm.begin()
    for word in ' '.join(translation).split():
        (lm_state, word_logprob) = lm.score(lm_state, word)
        lm_score += word_logprob
    lm_score += lm.end(lm_state) if isEnd is True else 0.0
    return lm_score,lm_state

def isNoun(word, t):
    if t:
       if "n" in t[0]:
          return True
    elif word.endswith('ción') or word.endswith('sión') or word.endswith('dad') or word.endswith('es') or word.endswith('tad') or word.endswith('tado'):
       return True
    else:
       return False

def isAdj(word, t):
    if t:
        if "a" in t[0]:
            return True
    elif word.endswith("ado") or word.endswith("ada") or word.endswith("ido") or word.endswith("ida"):
        return True
    else:
        return False

def reorder(f):
    old_spanish_uni = uni.tag(f)
    spanish_uni = []
    index = 0
    added = False
    while index < len(old_spanish_uni):
        (word, tag) = old_spanish_uni[index]
        if isNoun(word, tag):
            if index < len(old_spanish_uni)-2:
               nextword = old_spanish_uni[index+1]
               if isAdj(nextword[0], nextword[1]):
                   spanish_uni.append(nextword[0])
                   spanish_uni.append(word)
                   index+=1
                   added = True
            if not added:
               spanish_uni.append(word)
        elif index > 0:
            oldword = old_spanish_uni[index-1]
            if oldword[1]:
               if "d" in oldword[1][0]: #determiner
                    if index < len(old_spanish_uni)-2:
                       nextword = old_spanish_uni[index+1]
                       if isAdj(nextword[0], nextword[1]):
                           spanish_uni.append(nextword[0])
                           spanish_uni.append(word)
                           index+=1
                           added = True
            if not added:
                spanish_uni.append(word)
        else:              
           spanish_uni.append(word)
        added = False
        index+=1
    return tuple(spanish_uni)
 
for f in input_sents:
    f = reorder(f)
    # 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, 0.0, 0.0, lm.begin(), [])

    stacks = [{} for _ in f] + [{}]
    stacks[0][lm.begin()] = initial_hypothesis
    for i, stack in enumerate(stacks[:-1]):
        # extend the top s hypotheses in the current stack
        for h in heapq.nlargest(opts.s, stack.itervalues(), key=lambda h: h.logprob): # prune
            #starting at index i, look one additional consecutive work until the end of the sentence
            for j in xrange(i+1,len(f)+1):
		#for each existing phrase, figure out the prob of the hypothesis, the Spanish phrase, the individual English words, & it ending there
                if f[i:j] in tm:
                    for phrase in tm[f[i:j]]:
                        tm_score = h.tm_score + phrase.logprob
                        for k in xrange( max(0, len(h.translation)-2), len(h.translation) +1):
                            new_translation = h.translation[:k] + [phrase.english] + h.translation[k:]
                            lm_score = 0.0
                            lm_state = None
                            if j!=len(f):
                                lm_score,lm_state = score_lm(new_translation, False)
                            else:
                                lm_score,lm_state = score_lm(new_translation, True)
                            new_hypothesis = hypothesis(tm_score + lm_score, tm_score, lm_score, lm_state, new_translation)
                            if new_hypothesis.lm_state not in stacks[j] or stacks[j][new_hypothesis.lm_state].logprob < new_hypothesis.logprob: # second case is recombination
                                stacks[j][new_hypothesis.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 ' '.join(h.translation)
    print extract_english_recursive(winner)
