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


def setup_parser():
    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=100, type=int, help='Maximum stack size (default=100)')
    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()
    return opts


def extract_english_recursive(h):
    return '' if h.predecessor is None else '%s%s ' % (extract_english_recursive(h.predecessor), h.phrase.english)

def print_stack(stack):
    for hyp in stack:
        print ""
        print hyp
        print ""

def main():
    opts = setup_parser()

    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]]

    # A hypothesis tuple object.
    #   logprob: the current probability of the hypothesis in log space
    #   logfuture: the estimated future cost of the rest of the translation in
    #              log space
    #   lm_state: the language model state, necessary to interact with LM class
    #   trans_array: a boolean array indicating which words have been translated
    #   predecessor: the previous hypothesis we had, which is used to
    #                reconstruct phrase sequence as a string
    #   src_ps: the index of the starting word of the source phrase
    #   src_pe: the index of the ending word of the source phrase
    #   phrase: the target phrase that we are adding in this hypothesis
    hypothesis = namedtuple('hypothesis',
                            'logprob, logfuture, lm_state, trans_array, predecessor, src_ps, src_pe, phrase')

    for f in input_sents:
        trans_array = [False] * len(f)

        future_table = compute_future_table(tm, lm, f)

        # The following code implements a decoding algorithm that can permute
        # the target phrases.
        # Hence all hypotheses in stacks[i] represent translations of where i
        # words have been translated.
        initial_hypothesis = hypothesis(0.0, 0.0, lm.begin(),
                                        trans_array, None,
                                        None, None, None)

        stacks = [[] for _ in f] + [[]]
        stacks[0].append(initial_hypothesis)
        # iterate over every stack except for the last one, since the last stack has
        # the fully decoded sentences, and thus we cannot do any more work on it
        for i, stack in enumerate(stacks[:-1]):
            # extend the top s hypotheses in the current stack
            # For some reasOn, if we define our key function to be:
            #   key = lambda h: h.logprob + h.logfuture
            # we get worse results. I am pretty sure that my computation of
            # future probability is right, so I don't really know what is going
            # on.
            for h in heapq.nlargest(opts.s, stack, key=lambda h: h.logprob): # prune
                trans_array = h.trans_array
                prev_start = h.src_ps
                prev_end = h.src_pe
                src_phrases = source_phrase_options(f, trans_array, prev_start, prev_end)

                for (s, e, src_phrase) in src_phrases:
                    if src_phrase in tm:
                        for tgt_phrase in tm[src_phrase]:
                            # translation model cost: phi(f_i | e_i)
                            logprob = h.logprob + tgt_phrase.logprob
                            lm_state = h.lm_state
                            # language model cost: p_{LM}(e)
                            for word in tgt_phrase.english.split():
                                (lm_state, word_logprob) = lm.score(lm_state, word)
                                logprob += word_logprob

                            # reordering cost
                            if prev_start is None:
                                prev_start = -1
                            if prev_end is None:
                                prev_end = -1
                            alpha = math.log(0.3)
                            logprob += abs(float(prev_end) + 1.0 - float(s)) * alpha

                            # update our coverage vector
                            new_trans = update_trans(trans_array, s, e)
                            nt = num_trans(new_trans)

                            # calculate future cost probability
                            futureprob = future_prob(f, new_trans, future_table)

                            if nt == len(new_trans):
                                logprob += lm.end(lm_state)
                            new_hypothesis = hypothesis(logprob, futureprob, lm_state,
                                                        new_trans, h, s, e, tgt_phrase)
                            # Adding the hypothesis to the stack if there is not
                            # another equal hypothesis with a better score
                            stacks[nt] = update_stack(stacks[nt], new_hypothesis)

        # find best translation by looking at the best scoring hypothesis
        # on the last stack.
        winner = max(stacks[-1], key=lambda h: h.logprob)
        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))


# Adds the hypothesis to the stack
def update_stack(stack, hyp):
    for i in xrange(0, len(stack)):
        other_hyp = stack[i]
        if hyp_eq_no_alpha(hyp, other_hyp):
            other_score = other_hyp.logprob + other_hyp.logfuture
            new_score = hyp.logprob + hyp.logfuture
            if (other_score < new_score):
                # recombination
                stack[i] = hyp
            # We found the equal stack and we either updated its score or there
            # is nothing to change in the stack
            return stack
    # We have not found any equal hypothesis, so we add our current one
    stack.append(hyp)
    return stack

# Tests for equality of hypotheses. Specifically,
#   1) do the words in the ngram model match
#   2) conversion vectors are equal
#   3) end of last translated phrases are equal
def hyp_eq_no_alpha(h1, h2):
    lm_model_eq = h1.lm_state == h2.lm_state
    conv_vec_eq = h1.trans_array == h2.trans_array
    prev_end_eq = h1.src_pe == h2.src_pe

    return (lm_model_eq and conv_vec_eq and prev_end_eq)


def update_trans(trans_array, start, end):
    trans_array = list(trans_array)
    for i in xrange(start, end+1):
        if (trans_array[i]):
            sys.stderr.write("Tried to translate something already translated\n")
        trans_array[i] = True
    return trans_array

# Counts the number of translated words in the translation array
def num_trans(trans_array):
    i = 0
    for word_trans in trans_array:
        if word_trans:
            i += 1

    return i


# Given a precomputed future cost table, we calculate the future probability
# cost based on phrases for all of the untranslated words in sentence f.
def future_prob(f, trans_array, cost_table):
    total_prob = 0.0

    sub_start = 0

    # Each iteration of this loop of this loop calculates the maximal
    # probability combination for decoding a given phrase
    while sub_start < len(f):
        # Fast-track the substring start until we hit an untranslated word
        while sub_start < len(f) and trans_array[sub_start]:
            sub_start += 1

        # At this point, sub_start has either passed the end of the string or it
        # points to an untranslated word. We can only proceed to find a phrase
        # if we are still within the string.
        if sub_start < len(f):
            # At this point, sub_start points to an untranslated word
            sub_end = sub_start+1
            while sub_end < len(f) and not trans_array[sub_end]:
                sub_end += 1

            # At this point, sub_end has either hit the end of the string or
            # found the first translated word after sub_start. We can now take
            # this phrase, bookended by sub_start and sub_end, and calculate its
            # probability.
            # cost_table[x][y] = cost for phrase f[x:x+y+1]
            # ==> cost for phrase f[x:y] = cost_table[x][y - x - 1]
            total_prob += cost_table[sub_start][sub_end - sub_start - 1]
            # update and reset our start and end word indices
            sub_start = sub_end

    return total_prob


# Computes a future cost table that takes into account the translation model
# and the language model, but not any reordering cost.
# The table that we return is not an actual rectangle. It follows the format:
#   - there are len(f) rows
#   - for row i, there len(f) - i columns
#   - row[i], col[j] represents the maximum probability to decode the phrase
#     f[i:i+j+1]
# Note that the i and j values used below are different than the i and j values
# referenced in the comments. That is, i and j are indexed from 0 when we
# interact with the future table. j is not indexed from i.
def compute_future_table(tm, lm, f):
    future_table = [None] * len(f)
    for i in xrange(0, len(f)):
        part_table = []
        for j in xrange(i, len(f)):
            src_phrase = f[i:j+1]
            temp_prob = tm_lm_prob(tm, lm, src_phrase)
            if (j > i):
                # TODO fix this.
                part_prob = tm_lm_prob(tm, lm, f[j:j+1])
                trans_prob = max(part_table[j-i-1] + part_prob, temp_prob)
            else:
                trans_prob = temp_prob
            part_table.append(trans_prob)
        future_table[i] = part_table
    return future_table


# Computes the maximal probability of decoding a source phrase given the
# translation model and language model.
def tm_lm_prob(tm, lm, src_phrase):
    if src_phrase in tm:
        max_logprob = float("-inf")
        for tgt_phrase in tm[src_phrase]:
            logprob = tgt_phrase.logprob
            lm_state = lm.begin()
            for word in tgt_phrase.english.split():
                (lm_state, word_logprob) = lm.score(lm_state, word)
                logprob += word_logprob
            logprob += lm.end(lm_state)
            max_logprob = max(logprob, max_logprob)
        return max_logprob
    # source phrase does not exist in our translation model
    else:
        return float("-inf")


# f: the foreign source sentence we are trying to find phrases in
# tm: the translation model that contains possible translation phrases from
#     source phrase
# trans_array: which words in the source have been translated already
# prev_start: the start index of the previous phrase that we translated
# prev_end: the end index of the previous phrase that we translated.
#
# If prev_start and prev_end are None (because we don't yet have a previous
# hypothesis), then we just make sure that the distance from the start of the
# sentence is not too great.
def source_phrase_options(f, trans_array, prev_start, prev_end):
    dist_limit = 5

    i = 0
    earliest_trans = None
    while i < len(trans_array) and trans_array[i]:
        i += 1
    # Either everything is translated and we went over the whole list or we
    # found the first untranslated word / start of a phrase
    if i < len(trans_array):
        earliest_trans = i
    else:
        return []

    if prev_end is None:
        # The earliest index that our phrase can start at
        i = 0
        # The latest index that our phrase can start at
        j = dist_limit - 1
    else:
        # The earliest index that our phrase can start at
        i = max(0, prev_end + 1 - dist_limit)
        # The latest index that our phrase can start at
        j = min(prev_end + 1 + dist_limit, len(f)-1)

    # The last argument restricts the max phrase size to 3.
    # Just remove the argument or set it to None if you don't want any limit on
    # phrase size.
    vrs = valid_ranges(trans_array, i, j, 3)

    # A list of tuples, where each element is a possible phrase.
    # Each tuple has: (phrase_start, phrase_end, phrase)
    phrase_opts = []

    for (s, e) in vrs:
        # We only allow phrases that will not push us far enough away that we
        # can never back track to the earliest untranslated phrase.
        if (e + 1 - earliest_trans) <= dist_limit:
            phrase_opts.append((s, e, f[s:e+1]))
    return phrase_opts



# Gets all valid ranges of phrases, restricting the start of the phrase to not
# go past the start and end indices, and also to not including any previously
# translated words. start and end are inclusive indices.
def valid_ranges(trans_array, start, end, max_len=None):
    valid_ranges = []

    vr_start = None
    vr_end = None
    for i in xrange(start, end+1):
        # the i'th word is not yet translated
        if not trans_array[i]:
            if max_len is None:
                j_end = len(trans_array)
            else:
                j_end = min(i + max_len, len(trans_array))
            for j in xrange(i, j_end):
                if not trans_array[j]:
                    valid_ranges.append((i, j))
                # If we hit an already translated word, it cannot be part of a
                # phrase, so we must break out from finding any longer phrase
                # end points given our current start index
                else:
                    break
        # We don't need any else statement here. If the i'th index is already
        # translated, we simply don't do anything and skip to the next part

    return valid_ranges


if (__name__ == "__main__"):
    main()
