#!/usr/bin/env python3

# Copyright 2016  Johns Hopkins University (Author: Daniel Povey)
# Apache 2.0.

from __future__ import print_function
from __future__ import division
import sys
import argparse
import math
from collections import defaultdict

import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer,encoding="utf8")
sys.stderr = io.TextIOWrapper(sys.stderr.buffer,encoding="utf8")
sys.stdin = io.TextIOWrapper(sys.stdin.buffer,encoding="utf8")

parser = argparse.ArgumentParser(description="""
This script creates a biased language model suitable for alignment and
data-cleanup purposes.   It reads (possibly multiple) lines of integerized text
from the input and writes a text-form FST of a backoff language model to
the standard output, to be piped into fstcompile.""")

parser.add_argument("--word-disambig-symbol", type = int, required = True,
                    help = "Integer corresponding to the disambiguation "
                    "symbol (normally #0) for backoff arcs")
parser.add_argument("--ngram-order", type = int, default = 4,
                    choices = [2,3,4,5,6,7],
                    help = "Maximum order of n-gram to use (but see also "
                    "--min-lm-state-count; the effective order may be less.")
parser.add_argument("--min-lm-state-count", type = int, default = 10,
                    help = "Minimum count below which we will completely "
                    "discount an LM-state (if it is of order > 2, i.e. "
                    "history-length > 1).")
parser.add_argument("--top-words", type = str,
                    help = "File containing frequent words and probabilities to be added into "
                    "the language model, with lines in the format '<integer-id-of-word> <prob>'. "
                    "These probabilities will be added to the probabilities in the unigram "
                    "backoff state and then renormalized; this option allows you to introduce "
                    "common words to the LM with specified probabilities.")
parser.add_argument("--discounting-constant", type = float, default = 0.3,
                    help = "Discounting constant D for standard (unmodified) Kneser-Ney; "
                    "must be strictly between 0 and 1.  A value closer to 0 will give "
                    "you a more-strongly-biased LM.")
parser.add_argument("--verbose", type = int, default = 0,
                    choices=[0,1,2,3,4,5], help = "Verbose level")

args = parser.parse_args()

if args.verbose >= 1:
    print(' '.join(sys.argv), file = sys.stderr)




class NgramCounts(object):
    ## A note on data-structure.
    ## Firstly, all words are represented as integers.
    ## We store n-gram counts as an array, indexed by (history-length == n-gram order minus one)
    ## (note: python calls arrays "lists")  of dicts from histories to counts, where
    ## histories are arrays of integers and "counts" are dicts from integer to float.
    ## For instance, when accumulating the 4-gram count for the '8' in the sequence '5 6 7 8',
    ## we'd do as follows:
    ##  self.counts[3][[5,6,7]][8] += 1.0
    ## where the [3] indexes an array, the [[5,6,7]] indexes a dict, and
    ## the [8] indexes a dict.
    def __init__(self, ngram_order):
        self.ngram_order = ngram_order
        # Integerized counts will never contain negative numbers, so
        # inside this program, we use -3 and -2 for the BOS and EOS symbols
        # respectively.
        # Note: it's actually important that the bos-symbol is the most negative;
        # it helps ensure that we print the state with left-context <s> first
        # when we print the FST, and this means that the start-state will have
        # the correct value.
        self.bos_symbol = -3
        self.eos_symbol = -2
        # backoff_symbol is kind of a pseudo-word, it's used in keeping track of
        # the backoff counts in each state.
        self.backoff_symbol = -1
        self.counts = []
        for n in range(ngram_order):
            # The 'lambda: defaultdict(float)' is an anonymous function taking
            # no arguments that returns a new defaultdict(float).
            # If we index self.counts[n][history] for a history-length n < ngram_order
            # and a previously unseen history, it will create a new defaultdict
            # that defaults to 0.0 [since the function float() will return 0.0].
            # This means that we can index self.counts without worrying about
            # undefined values.
            self.counts.append(defaultdict(lambda: defaultdict(float)))

    # adds a raw count (called while processing input data).
    # Suppose we see the sequence '6 7 8 9' and ngram_order=4, 'history'
    # would be (6,7,8) and 'predicted_word' would be 9; 'count' would be
    # 1.0.
    def AddCount(self, history, predicted_word, count):
        self.counts[len(history)][history][predicted_word] += count

    # 'line' is a string containing a sequence of integer word-ids.
    # This function adds the un-smoothed counts from this line of text.
    def AddRawCountsFromLine(self, line):
        try:
            words = [self.bos_symbol] + [ int(x) for x in line.split() ] + [self.eos_symbol]
        except:
            sys.exit("make_one_biased_lm.py: bad input line {0} (expected a sequence "
                     "of integers)".format(line))

        for n in range(1, len(words)):
            predicted_word = words[n]
            history_start = max(0, n + 1 - self.ngram_order)
            history = tuple(words[history_start:n])
            self.AddCount(history, predicted_word, 1.0)

    def AddRawCountsFromStandardInput(self):
        lines_processed = 0
        while True:
            line = sys.stdin.readline()
            if line == '':
                break
            self.AddRawCountsFromLine(line)
            lines_processed += 1
        if lines_processed == 0 or args.verbose > 0:
            print("make_one_biased_lm.py: processed {0} lines of input".format(
                    lines_processed), file = sys.stderr)


    # This function returns a dict from history (as a tuple of integers of
    # length > 1, ignoring lower-order histories), to the total count of this
    # history state plus all history-states which back off to this history state.
    # It's used inside CompletelyDiscountLowCountStates().
    def GetHistToTotalCount(self):
        ans = defaultdict(float)
        for n in range(2, self.ngram_order):
            for hist, word_to_count in self.counts[n].items():
                total_count = sum(word_to_count.values())
                while len(hist) >= 2:
                    ans[hist] += total_count
                    hist = hist[1:]
        return ans


    # This function will completely discount the counts in any LM-states of
    # order > 2 (i.e. history-length > 1) that have total count below
    # 'min_count'; when computing the total counts, we include higher-order
    # LM-states that would back off to 'this' lm-state, in the total.
    def CompletelyDiscountLowCountStates(self, min_count):
        hist_to_total_count = self.GetHistToTotalCount()
        for n in reversed(list(range(2, self.ngram_order))):
            this_order_counts = self.counts[n]
            to_delete = []
            for hist in this_order_counts.keys():
                if hist_to_total_count[hist] < min_count:
                    # we need to completely back off this count.
                    word_to_count = this_order_counts[hist]
                    # mark this key for deleting
                    to_delete.append(hist)
                    backoff_hist = hist[1:]  # this will be a tuple not a list.
                    for word, count in word_to_count.items():
                        self.AddCount(backoff_hist, word, count)
            for hist in to_delete:
                del this_order_counts[hist]

    # This backs off the counts according to Kneser-Ney (unmodified,
    # with interpolation).
    def ApplyBackoff(self, D):
        assert D > 0.0 and D < 1.0
        for n in reversed(list(range(1, self.ngram_order))):
            this_order_counts = self.counts[n]
            for hist, word_to_count in this_order_counts.items():
                backoff_hist = hist[1:]
                backoff_word_to_count = self.counts[n-1][backoff_hist]
                this_discount_total = 0.0
                for word in word_to_count:
                    assert word_to_count[word] >= 1.0
                    word_to_count[word] -= D
                    this_discount_total += D
                    # Interpret the following line as incrementing the
                    # count-of-counts for the next-lower order.
                    backoff_word_to_count[word] += 1.0
                word_to_count[self.backoff_symbol] += this_discount_total


    # This function prints out to stderr the n-gram counts stored in this
    # object; it's used for debugging.
    def Print(self, info_string):
        print(info_string, file=sys.stderr)
        # these are useful for debug.
        total = 0.0
        total_excluding_backoff = 0.0
        for this_order_counts in self.counts:
            for hist, word_to_count in this_order_counts.items():
                this_total_count = sum(word_to_count.values())
                print('{0}: total={1} '.format(hist, this_total_count),
                      end='', file=sys.stderr)
                print(' '.join(['{0} -> {1} '.format(word, count)
                                for word, count in word_to_count.items() ]),
                      file = sys.stderr)
                total += this_total_count
                total_excluding_backoff += this_total_count
                if self.backoff_symbol in word_to_count:
                    total_excluding_backoff -= word_to_count[self.backoff_symbol]
        print('total count = {0}, excluding discount = {1}'.format(
                total, total_excluding_backoff), file = sys.stderr)

    def AddTopWords(self, top_words_file):
        empty_history = ()
        word_to_count = self.counts[0][empty_history]
        total = sum(word_to_count.values())
        try:
            f = open(top_words_file, mode='r', encoding='utf-8')
        except:
            sys.exit("make_one_biased_lm.py: error opening top-words file: "
                     "--top-words=" + top_words_file)
        while True:
            line = f.readline()
            if line == '':
                break
            try:
                [ word_index, prob ] = line.split()
                word_index = int(word_index)
                prob = float(prob)
                assert word_index > 0 and prob > 0.0
                word_to_count[word_index] += prob * total
            except Exception as e:
                sys.exit("make_one_biased_lm.py: could not make sense of the "
                         "line '{0}' in op-words file: {1} ".format(line, str(e)))
        f.close()


    def GetTotalCountMap(self):
        # This function, called from PrintAsFst, returns a map from
        # history to the total-count for that state.
        total_count_map = dict()
        for n in range(0, self.ngram_order):
            for hist, word_to_count in self.counts[n].items():
                total_count_map[hist] = sum(word_to_count.values())
        return total_count_map

    def GetHistToStateMap(self):
        # This function, called from PrintAsFst, returns a map from
        # history to integer FST-state.
        hist_to_state = dict()
        fst_state_counter = 0
        for n in range(0, self.ngram_order):
            for hist in self.counts[n].keys():
                hist_to_state[hist] = fst_state_counter
                fst_state_counter += 1
        return hist_to_state

    def GetProb(self, hist, word, total_count_map):
        total_count = total_count_map[hist]
        word_to_count = self.counts[len(hist)][hist]
        prob = float(word_to_count[word]) / total_count
        if len(hist) > 0 and word != self.backoff_symbol:
            prob_in_backoff = self.GetProb(hist[1:], word, total_count_map)
            backoff_prob = float(word_to_count[self.backoff_symbol]) / total_count
            prob += backoff_prob * prob_in_backoff
        return prob

    # This function prints the estimated language model as an FST.
    def PrintAsFst(self, word_disambig_symbol):
        # n is the history-length (== order + 1).  We iterate over the
        # history-length in the order 1, 0, 2, 3, and then iterate over the
        # histories of each order in sorted order.  Putting order 1 first
        # and sorting on the histories
        # ensures that the bigram state with <s> as the left context comes first.
        # (note: self.bos_symbol is the most negative symbol)

        # History will map from history (as a tuple) to integer FST-state.
        hist_to_state = self.GetHistToStateMap()
        total_count_map = self.GetTotalCountMap()

        for n in [ 1, 0 ] + list(range(2, self.ngram_order)):
            this_order_counts = self.counts[n]
            # For order 1, make sure the keys are sorted.
            keys = this_order_counts.keys() if n != 1 else sorted(this_order_counts.keys())
            for hist in keys:
                word_to_count = this_order_counts[hist]
                this_fst_state = hist_to_state[hist]

                for word in word_to_count.keys():
                    # work out this_cost.  Costs in OpenFst are negative logs.
                    this_cost = -math.log(self.GetProb(hist, word, total_count_map))

                    if word > 0: # a real word.
                        next_hist = hist + (word,)  # appending tuples
                        while not next_hist in hist_to_state:
                            next_hist = next_hist[1:]
                        next_fst_state = hist_to_state[next_hist]
                        print(this_fst_state, next_fst_state, word, word,
                              this_cost)
                    elif word == self.eos_symbol:
                        # print final-prob for this state.
                        print(this_fst_state, this_cost)
                    else:
                        assert word == self.backoff_symbol
                        backoff_fst_state = hist_to_state[hist[1:len(hist)]]
                        print(this_fst_state, backoff_fst_state,
                              word_disambig_symbol, 0, this_cost)


ngram_counts = NgramCounts(args.ngram_order)
ngram_counts.AddRawCountsFromStandardInput()

if args.verbose >= 3:
    ngram_counts.Print("Raw counts:")
ngram_counts.CompletelyDiscountLowCountStates(args.min_lm_state_count)
if args.verbose >= 3:
    ngram_counts.Print("Counts after discounting low-count states:")
ngram_counts.ApplyBackoff(args.discounting_constant)
if args.verbose >= 3:
    ngram_counts.Print("Counts after applying Kneser-Ney discounting:")
if args.top_words != None:
    ngram_counts.AddTopWords(args.top_words)
    if args.verbose >= 3:
        ngram_counts.Print("Counts after applying top-n-words")
ngram_counts.PrintAsFst(args.word_disambig_symbol)


# test comand:
# (echo 6 7 8 4; echo 7 8 9; echo 7 8) | ./make_one_biased_lm.py --word-disambig-symbol=1000 --min-lm-state-count=2 --verbose=3 --top-words=<(echo 1 0.5; echo 2 0.25)