// lm/arpa-lm-compiler.cc
  
  // Copyright 2009-2011 Gilles Boulianne
  // Copyright 2016 Smart Action LLC (kkm)
  // Copyright 2017 Xiaohui Zhang
  
  // See ../../COPYING for clarification regarding multiple authors
  //
  // Licensed under the Apache License, Version 2.0 (the "License");
  // you may not use this file except in compliance with the License.
  // You may obtain a copy of the License at
  //
  //  http://www.apache.org/licenses/LICENSE-2.0
  //
  // THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  // KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
  // WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
  // MERCHANTABLITY OR NON-INFRINGEMENT.
  // See the Apache 2 License for the specific language governing permissions and
  // limitations under the License.
  
  #include <algorithm>
  #include <limits>
  #include <sstream>
  #include <utility>
  
  #include "base/kaldi-math.h"
  #include "lm/arpa-lm-compiler.h"
  #include "util/stl-utils.h"
  #include "util/text-utils.h"
  #include "fstext/remove-eps-local.h"
  
  namespace kaldi {
  
  class ArpaLmCompilerImplInterface {
   public:
    virtual ~ArpaLmCompilerImplInterface() { }
    virtual void ConsumeNGram(const NGram& ngram, bool is_highest) = 0;
  };
  
  namespace {
  
  typedef int32 StateId;
  typedef int32 Symbol;
  
  // GeneralHistKey can represent state history in an arbitrarily large n
  // n-gram model with symbol ids fitting int32.
  class GeneralHistKey {
   public:
    // Construct key from being and end iterators.
    template<class InputIt>
    GeneralHistKey(InputIt begin, InputIt end) : vector_(begin, end) { }
    // Construct empty history key.
    GeneralHistKey() : vector_() { }
    // Return tails of the key as a GeneralHistKey. The tails of an n-gram
    // w[1..n] is the sequence w[2..n] (and the heads is w[1..n-1], but the
    // key class does not need this operartion).
    GeneralHistKey Tails() const {
      return GeneralHistKey(vector_.begin() + 1, vector_.end());
    }
    // Keys are equal if represent same state.
    friend bool operator==(const GeneralHistKey& a, const GeneralHistKey& b) {
      return a.vector_ == b.vector_;
    }
    // Public typename HashType for hashing.
    struct HashType : public std::unary_function<GeneralHistKey, size_t> {
      size_t operator()(const GeneralHistKey& key) const {
        return VectorHasher<Symbol>().operator()(key.vector_);
      }
    };
  
   private:
    std::vector<Symbol> vector_;
  };
  
  // OptimizedHistKey combines 3 21-bit symbol ID values into one 64-bit
  // machine word. allowing significant memory reduction and some runtime
  // benefit over GeneralHistKey. Since 3 symbols are enough to track history
  // in a 4-gram model, this optimized key is used for smaller models with up
  // to 4-gram and symbol values up to 2^21-1.
  //
  // See GeneralHistKey for interface requrements of a key class.
  class OptimizedHistKey {
   public:
    enum {
      kShift = 21,  // 21 * 3 = 63 bits for data.
      kMaxData = (1 << kShift) - 1
    };
    template<class InputIt>
    OptimizedHistKey(InputIt begin, InputIt end) : data_(0) {
      for (uint32 shift = 0; begin != end; ++begin, shift += kShift) {
        data_ |= static_cast<uint64>(*begin) << shift;
      }
    }
    OptimizedHistKey() : data_(0) { }
    OptimizedHistKey Tails() const {
      return OptimizedHistKey(data_ >> kShift);
    }
    friend bool operator==(const OptimizedHistKey& a, const OptimizedHistKey& b) {
      return a.data_ == b.data_;
    }
    struct HashType : public std::unary_function<OptimizedHistKey, size_t> {
      size_t operator()(const OptimizedHistKey& key) const { return key.data_; }
    };
  
   private:
    explicit OptimizedHistKey(uint64 data) : data_(data) { }
    uint64 data_;
  };
  
  }  // namespace
  
  template <class HistKey>
  class ArpaLmCompilerImpl : public ArpaLmCompilerImplInterface {
   public:
    ArpaLmCompilerImpl(ArpaLmCompiler* parent, fst::StdVectorFst* fst,
                       Symbol sub_eps);
  
    virtual void ConsumeNGram(const NGram &ngram, bool is_highest);
  
   private:
    StateId AddStateWithBackoff(HistKey key, float backoff);
    void CreateBackoff(HistKey key, StateId state, float weight);
  
    ArpaLmCompiler *parent_;  // Not owned.
    fst::StdVectorFst* fst_;  // Not owned.
    Symbol bos_symbol_;
    Symbol eos_symbol_;
    Symbol sub_eps_;
  
    StateId eos_state_;
    typedef unordered_map<HistKey, StateId,
                          typename HistKey::HashType> HistoryMap;
    HistoryMap history_;
  };
  
  template <class HistKey>
  ArpaLmCompilerImpl<HistKey>::ArpaLmCompilerImpl(
      ArpaLmCompiler* parent, fst::StdVectorFst* fst, Symbol sub_eps)
      : parent_(parent), fst_(fst), bos_symbol_(parent->Options().bos_symbol),
        eos_symbol_(parent->Options().eos_symbol), sub_eps_(sub_eps) {
    // The algorithm maintains state per history. The 0-gram is a special state
    // for emptry history. All unigrams (including BOS) backoff into this state.
    StateId zerogram = fst_->AddState();
    history_[HistKey()] = zerogram;
  
    // Also, if </s> is not treated as epsilon, create a common end state for
    // all transitions acepting the </s>, since they do not back off. This small
    // optimization saves about 2% states in an average grammar.
    if (sub_eps_ == 0) {
      eos_state_ = fst_->AddState();
      fst_->SetFinal(eos_state_, 0);
    }
  }
  
  template <class HistKey>
  void ArpaLmCompilerImpl<HistKey>::ConsumeNGram(const NGram &ngram,
                                                 bool is_highest) {
    // Generally, we do the following. Suppose we are adding an n-gram "A B
    // C". Then find the node for "A B", add a new node for "A B C", and connect
    // them with the arc accepting "C" with the specified weight. Also, add a
    // backoff arc from the new "A B C" node to its backoff state "B C".
    //
    // Two notable exceptions are the highest order n-grams, and final n-grams.
    //
    // When adding a highest order n-gram (e. g., our "A B C" is in a 3-gram LM),
    // the following optimization is performed. There is no point adding a node
    // for "A B C" with a "C" arc from "A B", since there will be no other
    // arcs ingoing to this node, and an epsilon backoff arc into the backoff
    // model "B C", with the weight of \bar{1}. To save a node, create an arc
    // accepting "C" directly from "A B" to "B C". This saves as many nodes
    // as there are the highest order n-grams, which is typically about half
    // the size of a large 3-gram model.
    //
    // Indeed, this does not apply to n-grams ending in EOS, since they do not
    // back off. These are special, as they do not have a back-off state, and
    // the node for "(..anything..) </s>" is always final. These are handled
    // in one of the two possible ways, If symbols <s> and </s> are being
    // replaced by epsilons, neither node nor arc is created, and the logprob
    // of the n-gram is applied to its source node as final weight. If <s> and
    // </s> are preserved, then a special final node for </s> is allocated and
    // used as the destination of the "</s>" acceptor arc.
    HistKey heads(ngram.words.begin(), ngram.words.end() - 1);
    typename HistoryMap::iterator source_it = history_.find(heads);
    if (source_it == history_.end()) {
      // There was no "A B", therefore the probability of "A B C" is zero.
      // Print a warning and discard current n-gram.
      if (parent_->ShouldWarn())
        KALDI_WARN << parent_->LineReference()
                   << " skipped: no parent (n-1)-gram exists";
      return;
    }
  
    StateId source = source_it->second;
    StateId dest;
    Symbol sym = ngram.words.back();
    float weight = -ngram.logprob;
    if (sym == sub_eps_ || sym == 0) {
      KALDI_ERR << " <eps> or disambiguation symbol " << sym << "found in the ARPA file. ";
    }
    if (sym == eos_symbol_) {
      if (sub_eps_ == 0) {
        // Keep </s> as a real symbol when not substituting.
        dest = eos_state_;
      } else {
        // Treat </s> as if it was epsilon: mark source final, with the weight
        // of the n-gram.
        fst_->SetFinal(source, weight);
        return;
      }
    } else {
      // For the highest order n-gram, this may find an existing state, for
      // non-highest, will create one (unless there are duplicate n-grams
      // in the grammar, which cannot be reliably detected if highest order,
      // so we better do not do that at all).
      dest = AddStateWithBackoff(
          HistKey(ngram.words.begin() + (is_highest ? 1 : 0),
                  ngram.words.end()),
          -ngram.backoff);
    }
  
    if (sym == bos_symbol_) {
      weight = 0;  // Accepting <s> is always free.
      if (sub_eps_ == 0) {
        // <s> is as a real symbol, only accepted in the start state.
        source = fst_->AddState();
        fst_->SetStart(source);
      } else {
        // The new state for <s> unigram history *is* the start state.
        fst_->SetStart(dest);
        return;
      }
    }
  
    // Add arc from source to dest, whichever way it was found.
    fst_->AddArc(source, fst::StdArc(sym, sym, weight, dest));
    return;
  }
  
  // Find or create a new state for n-gram defined by key, and ensure it has a
  // backoff transition.  The key is either the current n-gram for all but
  // highest orders, or the tails of the n-gram for the highest order. The
  // latter arises from the chain-collapsing optimization described above.
  template <class HistKey>
  StateId ArpaLmCompilerImpl<HistKey>::AddStateWithBackoff(HistKey key,
                                                           float backoff) {
    typename HistoryMap::iterator dest_it = history_.find(key);
    if (dest_it != history_.end()) {
      // Found an existing state in the history map. Invariant: if the state in
      // the map, then its backoff arc is in the FST. We are done.
      return dest_it->second;
    }
    // Otherwise create a new state and its backoff arc, and register in the map.
    StateId dest = fst_->AddState();
    history_[key] = dest;
    CreateBackoff(key.Tails(), dest, backoff);
    return dest;
  }
  
  // Create a backoff arc for a state. Key is a backoff destination that may or
  // may not exist. When the destination is not found, naturally fall back to
  // the lower order model, and all the way down until one is found (since the
  // 0-gram model is always present, the search is guaranteed to terminate).
  template <class HistKey>
  inline void ArpaLmCompilerImpl<HistKey>::CreateBackoff(
      HistKey key, StateId state, float weight) {
    typename HistoryMap::iterator dest_it = history_.find(key);
    while (dest_it == history_.end()) {
      key = key.Tails();
      dest_it = history_.find(key);
    }
  
    // The arc should transduce either <eos> or #0 to <eps>, depending on the
    // epsilon substitution mode. This is the only case when input and output
    // label may differ.
    fst_->AddArc(state, fst::StdArc(sub_eps_, 0, weight, dest_it->second));
  }
  
  ArpaLmCompiler::~ArpaLmCompiler() {
    if (impl_ != NULL)
      delete impl_;
  }
  
  void ArpaLmCompiler::HeaderAvailable() {
    KALDI_ASSERT(impl_ == NULL);
    // Use optimized implementation if the grammar is 4-gram or less, and the
    // maximum attained symbol id will fit into the optimized range.
    int64 max_symbol = 0;
    if (Symbols() != NULL)
      max_symbol = Symbols()->AvailableKey() - 1;
    // If augmenting the symbol table, assume the wors case when all words in
    // the model being read are novel.
    if (Options().oov_handling == ArpaParseOptions::kAddToSymbols)
      max_symbol += NgramCounts()[0];
  
    if (NgramCounts().size() <= 4 && max_symbol < OptimizedHistKey::kMaxData) {
      impl_ = new ArpaLmCompilerImpl<OptimizedHistKey>(this, &fst_, sub_eps_);
    } else {
      impl_ = new ArpaLmCompilerImpl<GeneralHistKey>(this, &fst_, sub_eps_);
      KALDI_LOG << "Reverting to slower state tracking because model is large: "
                << NgramCounts().size() << "-gram with symbols up to "
                << max_symbol;
    }
  }
  
  void ArpaLmCompiler::ConsumeNGram(const NGram &ngram) {
    // <s> is invalid in tails, </s> in heads of an n-gram.
    for (int i = 0; i < ngram.words.size(); ++i) {
      if ((i > 0 && ngram.words[i] == Options().bos_symbol) ||
          (i + 1 < ngram.words.size()
           && ngram.words[i] == Options().eos_symbol)) {
        if (ShouldWarn())
          KALDI_WARN << LineReference()
                     << " skipped: n-gram has invalid BOS/EOS placement";
        return;
      }
    }
  
    bool is_highest = ngram.words.size() == NgramCounts().size();
    impl_->ConsumeNGram(ngram, is_highest);
  }
  
  void ArpaLmCompiler::RemoveRedundantStates() {
    fst::StdArc::Label backoff_symbol = sub_eps_;
    if (backoff_symbol == 0) {
      // The method of removing redundant states implemented in this function
      // leads to slow determinization of L o G when people use the older style of
      // usage of arpa2fst where the --disambig-symbol option was not specified.
      // The issue seems to be that it creates a non-deterministic FST, while G is
      // supposed to be deterministic.  By 'return'ing below, we just disable this
      // method if people were using an older script.  This method isn't really
      // that consequential anyway, and people will move to the newer-style
      // scripts (see current utils/format_lm.sh), so this isn't much of a
      // problem.
      return;
    }
  
    fst::StdArc::StateId num_states = fst_.NumStates();
  
  
    // replace the #0 symbols on the input of arcs out of redundant states (states
    // that are not final and have only a backoff arc leaving them), with <eps>.
    for (fst::StdArc::StateId state = 0; state < num_states; state++) {
      if (fst_.NumArcs(state) == 1 && fst_.Final(state) == fst::TropicalWeight::Zero()) {
        fst::MutableArcIterator<fst::StdVectorFst> iter(&fst_, state);
        fst::StdArc arc = iter.Value();
        if (arc.ilabel == backoff_symbol) {
          arc.ilabel = 0;
          iter.SetValue(arc);
        }
      }
    }
  
    // we could call fst::RemoveEps, and it would have the same effect in normal
    // cases, where backoff_symbol != 0 and there are no epsilons in unexpected
    // places, but RemoveEpsLocal is a bit safer in case something weird is going
    // on; it guarantees not to blow up the FST.
    fst::RemoveEpsLocal(&fst_);
    KALDI_LOG << "Reduced num-states from " << num_states << " to "
              << fst_.NumStates();
  }
  
  void ArpaLmCompiler::Check() const {
    if (fst_.Start() == fst::kNoStateId) {
      KALDI_ERR << "Arpa file did not contain the beginning-of-sentence symbol "
                << Symbols()->Find(Options().bos_symbol) << ".";
    }
  }
  
  void ArpaLmCompiler::ReadComplete() {
    fst_.SetInputSymbols(Symbols());
    fst_.SetOutputSymbols(Symbols());
    RemoveRedundantStates();
    Check();
  }
  
  }  // namespace kaldi