// lm/const-arpa-lm.cc

// Copyright 2014  Guoguo Chen

// 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-file-parser.h"
#include "lm/const-arpa-lm.h"
#include "util/stl-utils.h"
#include "util/text-utils.h"


namespace kaldi {

// Auxiliary struct for converting ConstArpaLm format langugae model to Arpa
// format.
struct ArpaLine {
  std::vector<int32> words;  // Sequence of words to be printed.
  float logprob;             // Logprob corresponds to word sequence.
  float backoff_logprob;     // Backoff_logprob corresponds to word sequence.
  // Comparison function for sorting.
  bool operator < (const ArpaLine &other) const {
    if (words.size() < other.words.size()) {
      return true;
    } else if (words.size() > other.words.size()) {
      return false;
    } else {
      return words < other.words;
    }
  }
};

// Auxiliary class to build ConstArpaLm. We first use this class to figure out
// the relative address of different LmStates, and then put everything into one
// block in memory.
class LmState {
 public:
  union ChildType {
    // If child is not the final order, we keep the pointer to its LmState.
    LmState* state;

    // If child is the final order, we only keep the log probability for it.
    float prob;
  };

  struct ChildrenVectorLessThan {
    bool operator()(
        const std::pair<int32, union ChildType>& lhs,
        const std::pair<int32, union ChildType>& rhs) const {
      return lhs.first < rhs.first;
    }
  };

  LmState(const bool is_unigram, const bool is_child_final_order,
          const float logprob, const float backoff_logprob) :
      is_unigram_(is_unigram), is_child_final_order_(is_child_final_order),
      logprob_(logprob), backoff_logprob_(backoff_logprob) {}

  void SetMyAddress(const int64 address) {
    my_address_ = address;
  }

  void AddChild(const int32 word, LmState* child_state) {
    KALDI_ASSERT(!is_child_final_order_);
    ChildType child;
    child.state = child_state;
    children_.push_back(std::make_pair(word, child));
  }

  void AddChild(const int32 word, const float child_prob) {
    KALDI_ASSERT(is_child_final_order_);
    ChildType child;
    child.prob = child_prob;
    children_.push_back(std::make_pair(word, child));
  }

  int64 MyAddress() const {
    return my_address_;
  }

  bool IsUnigram() const {
    return is_unigram_;
  }

  bool IsChildFinalOrder() const {
    return is_child_final_order_;
  }

  float Logprob() const {
    return logprob_;
  }

  float BackoffLogprob() const {
    return backoff_logprob_;
  }

  int32 NumChildren() const {
    return children_.size();
  }

  std::pair<int32, union ChildType> GetChild(const int32 index) {
    KALDI_ASSERT(index < children_.size());
    KALDI_ASSERT(index >= 0);
    return children_[index];
  }

  void SortChildren() {
    std::sort(children_.begin(), children_.end(), ChildrenVectorLessThan());
  }

  // Checks if the current LmState is a leaf.
  bool IsLeaf() const {
    return (backoff_logprob_ == 0.0 && children_.empty());
  }

  // Computes the size of the memory that the current LmState would take in
  // <lm_states> array. It's the number of 4-byte chunks.
  int32 MemSize() const {
    if (IsLeaf() && !is_unigram_) {
      // We don't create an entry in this case; the logprob will be stored in
      // the same int32 that we would normally store the pointer in.
      return 0;
    } else {
      // We store the following information:
      // logprob, backoff_logprob, children.size() and children data.
      return (3 + 2 * children_.size());
    }
  }

 private:
  // Unigram states will have LmStates even if they are leaves, therefore we
  // need to note when this is a unigram or not.
  bool is_unigram_;

  // If the current LmState has an order of (final_order - 1), then its child
  // must be the final order. We only keep the log probability for its child.
  bool is_child_final_order_;

  // When we compute the addresses of the LmStates as offsets into <lm_states_>
  // pointer, and put the offsets here. Note that this is just offset, not
  // actual pointer.
  int64 my_address_;

  // Language model log probability of the current sequence. For example, if
  // this state is "A B", then it would be the logprob of "A -> B".
  float logprob_;

  // Language model backoff log probability of the current sequence, e.g., state
  // "A B -> X" backing off to "B -> X".
  float backoff_logprob_;

  // List of children.
  std::vector<std::pair<int32, union ChildType> > children_;
};

// Class to build ConstArpaLm from Arpa format language model. It relies on the
// auxiliary class LmState above.
class ConstArpaLmBuilder : public ArpaFileParser {
 public:
  explicit ConstArpaLmBuilder(ArpaParseOptions options)
      : ArpaFileParser(options, NULL) {
    ngram_order_ = 0;
    num_words_ = 0;
    overflow_buffer_size_ = 0;
    lm_states_size_ = 0;
    max_address_offset_ = pow(2, 30) - 1;
    is_built_ = false;
    lm_states_ = NULL;
    unigram_states_ = NULL;
    overflow_buffer_ = NULL;
  }

  ~ConstArpaLmBuilder() {
    unordered_map<std::vector<int32>,
                  LmState*, VectorHasher<int32> >::iterator iter;
    for (iter = seq_to_state_.begin(); iter != seq_to_state_.end(); ++iter) {
      delete iter->second;
    }
    if (is_built_) {
      delete[] lm_states_;
      delete[] unigram_states_;
      delete[] overflow_buffer_;
    }
  }

  // Writes ConstArpaLm.
  void Write(std::ostream &os, bool binary) const;

  void SetMaxAddressOffset(const int32 max_address_offset) {
    KALDI_WARN << "You are changing <max_address_offset_>; the default should "
        << "not be changed unless you are in testing mode.";
    max_address_offset_ = max_address_offset;
  }

 protected:
  // ArpaFileParser overrides.
  virtual void HeaderAvailable();
  virtual void ConsumeNGram(const NGram& ngram);
  virtual void ReadComplete();

 private:
  struct WordsAndLmStatePairLessThan {
    bool operator()(
        const std::pair<std::vector<int32>*, LmState*>& lhs,
        const std::pair<std::vector<int32>*, LmState*>& rhs) const {
      return *(lhs.first) < *(rhs.first);
    }
  };

 private:
  // Indicating if ConstArpaLm has been built or not.
  bool is_built_;

  // Maximum relative address for the child. We put it here just for testing.
  // The default value is 30-bits and should not be changed except for testing.
  int32 max_address_offset_;

  // N-gram order of language model. This can be figured out from "/data/"
  // section in Arpa format language model.
  int32 ngram_order_;

  // Index of largest word-id plus one. It defines the end of <unigram_states_>
  // array.
  int32 num_words_;

  // Number of entries in the overflow buffer for pointers that couldn't be
  // represented as a 30-bit relative index.
  int32 overflow_buffer_size_;

  // Size of the <lm_states_> array, which will be needed by I/O.
  int64 lm_states_size_;

  // Memory blcok for storing LmStates.
  int32* lm_states_;

  // Memory block for storing pointers of unigram LmStates.
  int32** unigram_states_;

  // Memory block for storing pointers of the LmStates that have large relative
  // address to their parents.
  int32** overflow_buffer_;

  // Hash table from word sequences to LmStates.
  unordered_map<std::vector<int32>,
                LmState*, VectorHasher<int32> > seq_to_state_;
};

void ConstArpaLmBuilder::HeaderAvailable() {
  ngram_order_ = NgramCounts().size();
}

void ConstArpaLmBuilder::ConsumeNGram(const NGram &ngram) {
  int32 cur_order = ngram.words.size();
  // If <ngram_order_> is larger than 1, then we do not create LmState for
  // the final order entry. We only keep the log probability for it.
  LmState *lm_state = NULL;
  if (cur_order != ngram_order_ || ngram_order_ == 1) {
    lm_state = new LmState(cur_order == 1,
                           cur_order == ngram_order_ - 1,
                           ngram.logprob, ngram.backoff);

    if (seq_to_state_.find(ngram.words) != seq_to_state_.end()) {
      std::ostringstream os;
      os << "[ ";
      for (size_t i = 0; i < ngram.words.size(); i++) {
        os << ngram.words[i] << " ";
      }
      os <<"]";

      KALDI_ERR << "N-gram " << os.str() << " appears twice in the arpa file";
    }
    seq_to_state_[ngram.words] = lm_state;
  }

  // If n-gram order is larger than 1, we have to add possible child to
  // existing LmStates. We have the following two assumptions:
  // 1. N-grams are processed from small order to larger ones, i.e., from
  //    1, 2, ... to the highest order.
  // 2. If a n-gram exists in the Arpa format language model, then the
  //    "history" n-gram also exists. For example, if "A B C" is a valid
  //    n-gram, then "A B" is also a valid n-gram.
  int32 last_word = ngram.words[cur_order - 1];
  if (cur_order > 1) {
    std::vector<int32> hist(ngram.words.begin(), ngram.words.end() - 1);
    unordered_map<std::vector<int32>,
                  LmState*, VectorHasher<int32> >::iterator hist_iter;
    hist_iter = seq_to_state_.find(hist);
    if (hist_iter == seq_to_state_.end()) {
      std::ostringstream ss;
      for (int i = 0; i < cur_order; ++i)
        ss << (i == 0 ? '[' : ' ') << ngram.words[i];
      KALDI_ERR << "In line " << LineNumber() << ": "
                << cur_order << "-gram " << ss.str() << "] does not have "
                << "a parent model " << cur_order << "-gram.";
    }
    if (cur_order != ngram_order_ || ngram_order_ == 1) {
      KALDI_ASSERT(lm_state != NULL);
      KALDI_ASSERT(!hist_iter->second->IsChildFinalOrder());
      hist_iter->second->AddChild(last_word, lm_state);
    } else {
      KALDI_ASSERT(lm_state == NULL);
      KALDI_ASSERT(hist_iter->second->IsChildFinalOrder());
      hist_iter->second->AddChild(last_word, ngram.logprob);
    }
  } else {
    // Figures out <max_word_id>.
    num_words_ = std::max(num_words_, last_word + 1);
  }
}

// ConstArpaLm can be built in the following steps, assuming we have already
// created LmStates <seq_to_state_>:
// 1. Sort LmStates lexicographically.
//    This enables us to compute relative address. When we say lexicographic, we
//    treat the word-ids as letters. After sorting, the LmStates are in the
//    following order:
//    ...
//    A B
//    A B A
//    A B B
//    A B C
//    ...
//    where each line represents a LmState.
// 2. Update <my_address> in LmState, which is relative to the first element in
//    <sorted_vec>.
// 3. Put the following structure into the memory block
//    struct LmState {
//      float logprob;
//      float backoff_logprob;
//      int32 num_children;
//      std::pair<int32, int32> [] children;
//    }
//
//    At the same time, we will also create two special buffers:
//    <unigram_states_>
//    <overflow_buffer_>
void ConstArpaLmBuilder::ReadComplete() {
  // STEP 1: sorting LmStates lexicographically.
  // Vector for holding the sorted LmStates.
  std::vector<std::pair<std::vector<int32>*, LmState*> > sorted_vec;
  unordered_map<std::vector<int32>,
                LmState*, VectorHasher<int32> >::iterator iter;
  for (iter = seq_to_state_.begin(); iter != seq_to_state_.end(); ++iter) {
    if (iter->second->MemSize() > 0) {
      sorted_vec.push_back(
          std::make_pair(const_cast<std::vector<int32>*>(&(iter->first)),
                         iter->second));
    }
  }

  std::sort(sorted_vec.begin(), sorted_vec.end(),
            WordsAndLmStatePairLessThan());

  // STEP 2: updating <my_address> in LmState.
  for (int32 i = 0; i < sorted_vec.size(); ++i) {
    lm_states_size_ += sorted_vec[i].second->MemSize();
    if (i == 0) {
      sorted_vec[i].second->SetMyAddress(0);
    } else {
      sorted_vec[i].second->SetMyAddress(sorted_vec[i - 1].second->MyAddress()
          + sorted_vec[i - 1].second->MemSize());
    }
  }

  // STEP 3: creating memory block to store LmStates.
  // Reserves a memory block for LmStates.
  int64 lm_states_index = 0;
  try {
    lm_states_ = new int32[lm_states_size_];
  } catch(const std::exception &e) {
    KALDI_ERR << e.what();
  }

  // Puts data into memory block.
  unigram_states_ = new int32*[num_words_];
  std::vector<int32*> overflow_buffer_vec;
  for (int32 i = 0; i < num_words_; ++i) {
    unigram_states_[i] = NULL;
  }
  for (int32 i = 0; i < sorted_vec.size(); ++i) {
    // Current address.
    int32* parent_address = lm_states_ + lm_states_index;

    // Adds logprob.
    Int32AndFloat logprob_f(sorted_vec[i].second->Logprob());
    lm_states_[lm_states_index++] = logprob_f.i;

    // Adds backoff_logprob.
    Int32AndFloat backoff_logprob_f(sorted_vec[i].second->BackoffLogprob());
    lm_states_[lm_states_index++] = backoff_logprob_f.i;

    // Adds num_children.
    lm_states_[lm_states_index++] = sorted_vec[i].second->NumChildren();

    // Adds children, there are 3 cases:
    // 1. Child is a leaf and not unigram
    // 2. Child is not a leaf or is unigram
    //    2.1 Relative address can be represented by 30 bits
    //    2.2 Relative address cannot be represented by 30 bits
    sorted_vec[i].second->SortChildren();
    for (int32 j = 0; j < sorted_vec[i].second->NumChildren(); ++j) {
      int32 child_info;
      if (sorted_vec[i].second->IsChildFinalOrder() ||
          sorted_vec[i].second->GetChild(j).second.state->MemSize() == 0) {
        // Child is a leaf and not unigram. In this case we will not create an
        // entry in <lm_states_>; instead, we put the logprob in the place where
        // we normally store the poitner.
        Int32AndFloat child_logprob_f;
        if (sorted_vec[i].second->IsChildFinalOrder()) {
          child_logprob_f.f = sorted_vec[i].second->GetChild(j).second.prob;
        } else {
          child_logprob_f.f =
              sorted_vec[i].second->GetChild(j).second.state->Logprob();
        }
        child_info = child_logprob_f.i;
        child_info &= ~1;   // Sets the last bit to 0 so <child_info> is even.
      } else {
        // Child is not a leaf or is unigram.
        int64 offset =
            sorted_vec[i].second->GetChild(j).second.state->MyAddress()
            - sorted_vec[i].second->MyAddress();
        KALDI_ASSERT(offset > 0);
        if (offset <= max_address_offset_) {
          // Relative address can be represented by 30 bits.
          child_info = offset * 2;
          child_info |= 1;
        } else {
          // Relative address cannot be represented by 30 bits, we have to put
          // the child address into <overflow_buffer_>.
          int32* abs_address = parent_address + offset;
          overflow_buffer_vec.push_back(abs_address);
          int32 overflow_buffer_index = overflow_buffer_vec.size() - 1;
          child_info = overflow_buffer_index * 2;
          child_info |= 1;
          child_info *= -1;
        }
      }
      // Child word.
      lm_states_[lm_states_index++] = sorted_vec[i].second->GetChild(j).first;
      // Child info.
      lm_states_[lm_states_index++] = child_info;
    }

    // If the current state corresponds to an unigram, then create a separate
    // loop up table to improve efficiency, since those will be looked up pretty
    // frequently.
    if (sorted_vec[i].second->IsUnigram()) {
      KALDI_ASSERT(sorted_vec[i].first->size() == 1);
      unigram_states_[(*sorted_vec[i].first)[0]] = parent_address;
    }
  }
  KALDI_ASSERT(lm_states_size_ == lm_states_index);

  // Move <overflow_buffer_> from vector holder to array.
  overflow_buffer_size_ = overflow_buffer_vec.size();
  overflow_buffer_ = new int32*[overflow_buffer_size_];
  for (int32 i = 0; i < overflow_buffer_size_; ++i) {
    overflow_buffer_[i] = overflow_buffer_vec[i];
  }

  is_built_ = true;
}

void ConstArpaLmBuilder::Write(std::ostream &os, bool binary) const {
  if (!binary) {
    KALDI_ERR << "text-mode writing is not implemented for ConstArpaLmBuilder.";
  }
  KALDI_ASSERT(is_built_);

  // Creates ConstArpaLm.
  ConstArpaLm const_arpa_lm(
      Options().bos_symbol, Options().eos_symbol, Options().unk_symbol,
      ngram_order_, num_words_, overflow_buffer_size_, lm_states_size_,
      unigram_states_, overflow_buffer_, lm_states_);
  const_arpa_lm.Write(os, binary);
}

void ConstArpaLm::Write(std::ostream &os, bool binary) const {
  KALDI_ASSERT(initialized_);
  if (!binary) {
    KALDI_ERR << "text-mode writing is not implemented for ConstArpaLm.";
  }

  WriteToken(os, binary, "<ConstArpaLm>");

  // Misc info.
  WriteToken(os, binary, "<LmInfo>");
  WriteBasicType(os, binary, bos_symbol_);
  WriteBasicType(os, binary, eos_symbol_);
  WriteBasicType(os, binary, unk_symbol_);
  WriteBasicType(os, binary, ngram_order_);
  WriteToken(os, binary, "</LmInfo>");

  // LmStates section.
  WriteToken(os, binary, "<LmStates>");
  WriteBasicType(os, binary, lm_states_size_);
  os.write(reinterpret_cast<char *>(lm_states_),
           sizeof(int32) * lm_states_size_);
  if (!os.good()) {
    KALDI_ERR << "ConstArpaLm <LmStates> section writing failed.";
  }
  WriteToken(os, binary, "</LmStates>");

  // Unigram section. We write memory offset to disk instead of the absolute
  // pointers.
  WriteToken(os, binary, "<LmUnigram>");
  WriteBasicType(os, binary, num_words_);
  int64* tmp_unigram_address = new int64[num_words_];
  for (int32 i = 0; i < num_words_; ++i) {
    // The relative address here is a little bit tricky:
    // 1. If the original address is NULL, then we set the relative address to
    //    zero.
    // 2. If the original address is not NULL, we set it to the following:
    //      unigram_states_[i] - lm_states_ + 1
    //    we plus 1 to ensure that the above value is positive.
    tmp_unigram_address[i] = (unigram_states_[i] == NULL) ? 0 :
        unigram_states_[i] - lm_states_ + 1;
  }
  os.write(reinterpret_cast<char *>(tmp_unigram_address),
           sizeof(int64) * num_words_);
  if (!os.good()) {
    KALDI_ERR << "ConstArpaLm <LmUnigram> section writing failed.";
  }
  delete[] tmp_unigram_address;   // Releases the memory.
  tmp_unigram_address = NULL;
  WriteToken(os, binary, "</LmUnigram>");

  // Overflow section. We write memory offset to disk instead of the absolute
  // pointers.
  WriteToken(os, binary, "<LmOverflow>");
  WriteBasicType(os, binary, overflow_buffer_size_);
  int64* tmp_overflow_address = new int64[overflow_buffer_size_];
  for (int32 i = 0; i < overflow_buffer_size_; ++i) {
    // The relative address here is a little bit tricky:
    // 1. If the original address is NULL, then we set the relative address to
    //    zero.
    // 2. If the original address is not NULL, we set it to the following:
    //      overflow_buffer_[i] - lm_states_ + 1
    //    we plus 1 to ensure that the above value is positive.
    tmp_overflow_address[i] = (overflow_buffer_[i] == NULL) ? 0 :
        overflow_buffer_[i] - lm_states_ + 1;
  }
  os.write(reinterpret_cast<char *>(tmp_overflow_address),
           sizeof(int64) * overflow_buffer_size_);
  if (!os.good()) {
    KALDI_ERR << "ConstArpaLm <LmOverflow> section writing failed.";
  }
  delete[] tmp_overflow_address;
  tmp_overflow_address = NULL;
  WriteToken(os, binary, "</LmOverflow>");
  WriteToken(os, binary, "</ConstArpaLm>");
}

void ConstArpaLm::Read(std::istream &is, bool binary) {
  KALDI_ASSERT(!initialized_);
  if (!binary) {
    KALDI_ERR << "text-mode reading is not implemented for ConstArpaLm.";
  }

  int first_char = is.peek();
  if (first_char == 4) {  // Old on-disk format starts with length of int32.
    ReadInternalOldFormat(is, binary);
  } else {                // New on-disk format starts with token <ConstArpaLm>.
    ReadInternal(is, binary);
  }
}

void ConstArpaLm::ReadInternal(std::istream &is, bool binary) {
  KALDI_ASSERT(!initialized_);
  if (!binary) {
    KALDI_ERR << "text-mode reading is not implemented for ConstArpaLm.";
  }

  ExpectToken(is, binary, "<ConstArpaLm>");

  // Misc info.
  ExpectToken(is, binary, "<LmInfo>");
  ReadBasicType(is, binary, &bos_symbol_);
  ReadBasicType(is, binary, &eos_symbol_);
  ReadBasicType(is, binary, &unk_symbol_);
  ReadBasicType(is, binary, &ngram_order_);
  ExpectToken(is, binary, "</LmInfo>");

  // LmStates section.
  ExpectToken(is, binary, "<LmStates>");
  ReadBasicType(is, binary, &lm_states_size_);
  lm_states_ = new int32[lm_states_size_];
  is.read(reinterpret_cast<char *>(lm_states_),
          sizeof(int32) * lm_states_size_);
  if (!is.good()) {
    KALDI_ERR << "ConstArpaLm <LmStates> section reading failed.";
  }
  ExpectToken(is, binary, "</LmStates>");

  // Unigram section. We write memory offset to disk instead of the absolute
  // pointers.
  ExpectToken(is, binary, "<LmUnigram>");
  ReadBasicType(is, binary, &num_words_);
  unigram_states_ = new int32*[num_words_];
  int64* tmp_unigram_address = new int64[num_words_];
  is.read(reinterpret_cast<char *>(tmp_unigram_address),
          sizeof(int64) * num_words_);
  if (!is.good()) {
    KALDI_ERR << "ConstArpaLm <LmUnigram> section reading failed.";
  }
  for (int32 i = 0; i < num_words_; ++i) {
    // Check out how we compute the relative address in ConstArpaLm::Write().
    unigram_states_[i] = (tmp_unigram_address[i] == 0) ? NULL
        : lm_states_ + tmp_unigram_address[i] - 1;
  }
  delete[] tmp_unigram_address;
  tmp_unigram_address = NULL;
  ExpectToken(is, binary, "</LmUnigram>");

  // Overflow section. We write memory offset to disk instead of the absolute
  // pointers.
  ExpectToken(is, binary, "<LmOverflow>");
  ReadBasicType(is, binary, &overflow_buffer_size_);
  overflow_buffer_ = new int32*[overflow_buffer_size_];
  int64* tmp_overflow_address = new int64[overflow_buffer_size_];
  is.read(reinterpret_cast<char *>(tmp_overflow_address),
          sizeof(int64) * overflow_buffer_size_);
  if (!is.good()) {
    KALDI_ERR << "ConstArpaLm <LmOverflow> section reading failed.";
  }
  for (int32 i = 0; i < overflow_buffer_size_; ++i) {
    // Check out how we compute the relative address in ConstArpaLm::Write().
    overflow_buffer_[i] = (tmp_overflow_address[i] == 0) ? NULL
        : lm_states_ + tmp_overflow_address[i] - 1;
  }
  delete[] tmp_overflow_address;
  tmp_overflow_address = NULL;
  ExpectToken(is, binary, "</LmOverflow>");
  ExpectToken(is, binary, "</ConstArpaLm>");

  KALDI_ASSERT(ngram_order_ > 0);
  KALDI_ASSERT(bos_symbol_ < num_words_ && bos_symbol_ > 0);
  KALDI_ASSERT(eos_symbol_ < num_words_ && eos_symbol_ > 0);
  KALDI_ASSERT(unk_symbol_ < num_words_ &&
               (unk_symbol_ > 0 || unk_symbol_ == -1));
  lm_states_end_ = lm_states_ + lm_states_size_ - 1;
  memory_assigned_ = true;
  initialized_ = true;
}

void ConstArpaLm::ReadInternalOldFormat(std::istream &is, bool binary) {
  KALDI_ASSERT(!initialized_);
  if (!binary) {
    KALDI_ERR << "text-mode reading is not implemented for ConstArpaLm.";
  }

  // Misc info.
  ReadBasicType(is, binary, &bos_symbol_);
  ReadBasicType(is, binary, &eos_symbol_);
  ReadBasicType(is, binary, &unk_symbol_);
  ReadBasicType(is, binary, &ngram_order_);

  // LmStates section.
  // In the deprecated version, <lm_states_size_> used to be type of int32,
  // which was a bug. We therefore use int32 for read for back-compatibility.
  int32 lm_states_size_int32;
  ReadBasicType(is, binary, &lm_states_size_int32);
  lm_states_size_ = static_cast<int64>(lm_states_size_int32);
  lm_states_ = new int32[lm_states_size_];
  for (int64 i = 0; i < lm_states_size_; ++i) {
    ReadBasicType(is, binary, &lm_states_[i]);
  }

  // Unigram section. We write memory offset to disk instead of the absolute
  // pointers.
  ReadBasicType(is, binary, &num_words_);
  unigram_states_ = new int32*[num_words_];
  for (int32 i = 0; i < num_words_; ++i) {
    int64 tmp_address;
    ReadBasicType(is, binary, &tmp_address);
    // Check out how we compute the relative address in ConstArpaLm::Write().
    unigram_states_[i] =
        (tmp_address == 0) ? NULL : lm_states_ + tmp_address - 1;
  }

  // Overflow section. We write memory offset to disk instead of the absolute
  // pointers.
  ReadBasicType(is, binary, &overflow_buffer_size_);
  overflow_buffer_ = new int32*[overflow_buffer_size_];
  for (int32 i = 0; i < overflow_buffer_size_; ++i) {
    int64 tmp_address;
    ReadBasicType(is, binary, &tmp_address);
    // Check out how we compute the relative address in ConstArpaLm::Write().
    overflow_buffer_[i] =
        (tmp_address == 0) ? NULL : lm_states_ + tmp_address - 1;
  }
  KALDI_ASSERT(ngram_order_ > 0);
  KALDI_ASSERT(bos_symbol_ < num_words_ && bos_symbol_ > 0);
  KALDI_ASSERT(eos_symbol_ < num_words_ && eos_symbol_ > 0);
  KALDI_ASSERT(unk_symbol_ < num_words_ &&
               (unk_symbol_ > 0 || unk_symbol_ == -1));
  lm_states_end_ = lm_states_ + lm_states_size_ - 1;
  memory_assigned_ = true;
  initialized_ = true;
}

bool ConstArpaLm::HistoryStateExists(const std::vector<int32>& hist) const {
  // We do not create LmState for empty word sequence, but technically it is the
  // history state of all unigrams.
  if (hist.size() == 0) {
    return true;
  }

  // Tries to locate the LmState of the given word sequence.
  int32* lm_state = GetLmState(hist);
  if (lm_state == NULL) {
    // <lm_state> does not exist means <hist> has no child.
    return false;
  } else {
    // Note that we always create LmState for unigrams, so even if <lm_state> is
    // not NULL, we still have to check if it has child.
    KALDI_ASSERT(lm_state >= lm_states_);
    KALDI_ASSERT(lm_state + 2 <= lm_states_end_);
    // <lm_state + 2> points to <num_children>.
    if (*(lm_state + 2) > 0) {
      return true;
    } else {
      return false;
    }
  }
  return true;
}

float ConstArpaLm::GetNgramLogprob(const int32 word,
                                   const std::vector<int32>& hist) const {
  KALDI_ASSERT(initialized_);

  // If the history size plus one is larger than <ngram_order_>, remove the old
  // words.
  std::vector<int32> mapped_hist(hist);
  while (mapped_hist.size() >= ngram_order_) {
    mapped_hist.erase(mapped_hist.begin(), mapped_hist.begin() + 1);
  }
  KALDI_ASSERT(mapped_hist.size() + 1 <= ngram_order_);

  // TODO(guoguo): check with Dan if this is reasonable.
  // Maps possible out-of-vocabulary words to <unk>. If a word does not have a
  // corresponding LmState, we treat it as <unk>. We map it to <unk> if <unk> is
  // specified.
  int32 mapped_word = word;
  if (unk_symbol_ != -1) {
    KALDI_ASSERT(mapped_word >= 0);
    if (mapped_word >= num_words_ || unigram_states_[mapped_word] == NULL) {
      mapped_word = unk_symbol_;
    }
    for (int32 i = 0; i < mapped_hist.size(); ++i) {
      KALDI_ASSERT(mapped_hist[i] >= 0);
      if (mapped_hist[i] >= num_words_ ||
          unigram_states_[mapped_hist[i]] == NULL) {
        mapped_hist[i] = unk_symbol_;
      }
    }
  }

  // Loops up n-gram probability.
  return GetNgramLogprobRecurse(mapped_word, mapped_hist);
}

float ConstArpaLm::GetNgramLogprobRecurse(
    const int32 word, const std::vector<int32>& hist) const {
  KALDI_ASSERT(initialized_);
  KALDI_ASSERT(hist.size() + 1 <= ngram_order_);

  // Unigram case.
  if (hist.size() == 0) {
    if (word >= num_words_ || unigram_states_[word] == NULL) {
      // If <unk> is defined, then the word sequence should have already been
      // mapped to <unk> is necessary; this is for the case where <unk> is not
      // defined.
      return std::numeric_limits<float>::min();
    } else {
      Int32AndFloat logprob_i(*unigram_states_[word]);
      return logprob_i.f;
    }
  }

  // High n-gram orders.
  float logprob = 0.0;
  float backoff_logprob = 0.0;
  int32* state;
  if ((state = GetLmState(hist)) != NULL) {
    int32 child_info;
    int32* child_lm_state = NULL;
    if (GetChildInfo(word, state, &child_info)) {
      DecodeChildInfo(child_info, state, &child_lm_state, &logprob);
      return logprob;
    } else {
      Int32AndFloat backoff_logprob_i(*(state + 1));
      backoff_logprob = backoff_logprob_i.f;
    }
  }
  std::vector<int32> new_hist(hist);
  new_hist.erase(new_hist.begin(), new_hist.begin() + 1);
  return backoff_logprob + GetNgramLogprobRecurse(word, new_hist);
}

int32* ConstArpaLm::GetLmState(const std::vector<int32>& seq) const {
  KALDI_ASSERT(initialized_);

  // No LmState exists for empty word sequence.
  if (seq.size() == 0) return NULL;

  // If <unk> is defined, then the word sequence should have already been mapped
  // to <unk> is necessary; this is for the case where <unk> is not defined.
  if (seq[0] >= num_words_ || unigram_states_[seq[0]] == NULL) return NULL;
  int32* parent = unigram_states_[seq[0]];

  int32 child_info;
  int32* child_lm_state = NULL;
  float logprob;
  for (int32 i = 1; i < seq.size(); ++i) {
    if (!GetChildInfo(seq[i], parent, &child_info)) {
      return NULL;
    }
    DecodeChildInfo(child_info, parent, &child_lm_state, &logprob);
    if (child_lm_state == NULL) {
      return NULL;
    } else {
      parent = child_lm_state;
    }
  }
  return parent;
}

bool ConstArpaLm::GetChildInfo(const int32 word,
                               int32* parent, int32* child_info) const {
  KALDI_ASSERT(initialized_);

  KALDI_ASSERT(parent != NULL);
  KALDI_ASSERT(parent >= lm_states_);
  KALDI_ASSERT(child_info != NULL);

  KALDI_ASSERT(parent + 2 <= lm_states_end_);
  int32 num_children = *(parent + 2);
  KALDI_ASSERT(parent + 2 + 2 * num_children <= lm_states_end_);

  if (num_children == 0) return false;

  // A binary search into the children memory block.
  int32 start_index = 1;
  int32 end_index = num_children;
  while (start_index <= end_index) {
    int32 mid_index = round((start_index + end_index) / 2);
    int32 mid_word = *(parent + 1 + 2 * mid_index);
    if (mid_word == word) {
      *child_info = *(parent + 2 + 2 * mid_index);
      return true;
    } else if (mid_word < word) {
      start_index = mid_index + 1;
    } else {
      end_index = mid_index - 1;
    }
  }

  return false;
}

void ConstArpaLm::DecodeChildInfo(const int32 child_info,
                                  int32* parent,
                                  int32** child_lm_state,
                                  float* logprob) const {
  KALDI_ASSERT(initialized_);

  KALDI_ASSERT(logprob != NULL);
  if (child_info % 2 == 0) {
    // Child is a leaf, only returns the log probability.
    *child_lm_state = NULL;
    Int32AndFloat logprob_i(child_info);
    *logprob = logprob_i.f;
  } else {
    int32 child_offset = child_info / 2;
    if (child_offset > 0) {
      *child_lm_state = parent + child_offset;
      Int32AndFloat logprob_i(**child_lm_state);
      *logprob = logprob_i.f;
    } else {
      KALDI_ASSERT(-child_offset < overflow_buffer_size_);
      *child_lm_state = overflow_buffer_[-child_offset];
      Int32AndFloat logprob_i(**child_lm_state);
      *logprob = logprob_i.f;
    }
    KALDI_ASSERT(*child_lm_state >= lm_states_);
    KALDI_ASSERT(*child_lm_state <= lm_states_end_);
  }
}

void ConstArpaLm::WriteArpaRecurse(int32* lm_state,
                                   const std::vector<int32>& seq,
                                   std::vector<ArpaLine> *output) const {
  if (lm_state == NULL) return;

  KALDI_ASSERT(lm_state >= lm_states_);
  KALDI_ASSERT(lm_state + 2 <= lm_states_end_);

  // Inserts the current LmState to <output>.
  ArpaLine arpa_line;
  arpa_line.words = seq;
  Int32AndFloat logprob_i(*lm_state);
  arpa_line.logprob = logprob_i.f;
  Int32AndFloat backoff_logprob_i(*(lm_state + 1));
  arpa_line.backoff_logprob = backoff_logprob_i.f;
  output->push_back(arpa_line);

  // Scans for possible children, and recursively adds child to <output>.
  int32 num_children = *(lm_state + 2);
  KALDI_ASSERT(lm_state + 2 + 2 * num_children <= lm_states_end_);
  for (int32 i = 0; i < num_children; ++i) {
    std::vector<int32> new_seq(seq);
    new_seq.push_back(*(lm_state + 3 + 2 * i));
    int32 child_info = *(lm_state + 4 + 2 * i);
    float logprob;
    int32* child_lm_state = NULL;
    DecodeChildInfo(child_info, lm_state, &child_lm_state, &logprob);

    if (child_lm_state == NULL) {
      // Leaf case.
      ArpaLine child_arpa_line;
      child_arpa_line.words = new_seq;
      child_arpa_line.logprob = logprob;
      child_arpa_line.backoff_logprob = 0.0;
      output->push_back(child_arpa_line);
    } else {
      WriteArpaRecurse(child_lm_state, new_seq, output);
    }
  }
}

void ConstArpaLm::WriteArpa(std::ostream &os) const {
  KALDI_ASSERT(initialized_);

  std::vector<ArpaLine> tmp_output;
  for (int32 i = 0; i < num_words_; ++i) {
    if (unigram_states_[i] != NULL) {
      std::vector<int32> seq(1, i);
      WriteArpaRecurse(unigram_states_[i], seq, &tmp_output);
    }
  }

  // Sorts ArpaLines and collects head information.
  std::sort(tmp_output.begin(), tmp_output.end());
  std::vector<int32> ngram_count(1, 0);
  for (int32 i = 0; i < tmp_output.size(); ++i) {
    if (tmp_output[i].words.size() >= ngram_count.size()) {
      ngram_count.resize(tmp_output[i].words.size() + 1);
      ngram_count[tmp_output[i].words.size()] = 1;
    } else {
      ngram_count[tmp_output[i].words.size()] += 1;
    }
  }

  // Writes the header.
  os << std::endl;
  os << "\\data\\" << std::endl;
  for (int32 i = 1; i < ngram_count.size(); ++i) {
    os << "ngram " << i << "=" << ngram_count[i] << std::endl;
  }

  // Writes n-grams.
  int32 current_order = 0;
  for (int32 i = 0; i < tmp_output.size(); ++i) {
    // Beginning of a n-gram section.
    if (tmp_output[i].words.size() != current_order) {
      current_order = tmp_output[i].words.size();
      os << std::endl;
      os << "\\" << current_order << "-grams:" << std::endl;
    }

    // Writes logprob.
    os << tmp_output[i].logprob << '\t';

    // Writes word sequence.
    for (int32 j = 0; j < tmp_output[i].words.size(); ++j) {
      os << tmp_output[i].words[j];
      if (j != tmp_output[i].words.size() - 1) {
        os << " ";
      }
    }

    // Writes backoff_logprob if it is not zero.
    if (tmp_output[i].backoff_logprob != 0.0) {
      os << '\t' << tmp_output[i].backoff_logprob;
    }
    os << std::endl;
  }

  os << std::endl << "\\end\\" << std::endl;
}

ConstArpaLmDeterministicFst::ConstArpaLmDeterministicFst(
    const ConstArpaLm& lm) : lm_(lm) {
  // Creates a history state for <s>.
  std::vector<Label> bos_state(1, lm_.BosSymbol());
  state_to_wseq_.push_back(bos_state);
  wseq_to_state_[bos_state] = 0;
  start_state_ = 0;
}

fst::StdArc::Weight ConstArpaLmDeterministicFst::Final(StateId s) {
  // At this point, we should have created the state.
  KALDI_ASSERT(static_cast<size_t>(s) < state_to_wseq_.size());
  const std::vector<Label>& wseq = state_to_wseq_[s];
  float logprob = lm_.GetNgramLogprob(lm_.EosSymbol(), wseq);
  return Weight(-logprob);
}

bool ConstArpaLmDeterministicFst::GetArc(StateId s,
                                         Label ilabel, fst::StdArc *oarc) {
  // At this point, we should have created the state.
  KALDI_ASSERT(static_cast<size_t>(s) < state_to_wseq_.size());
  std::vector<Label> wseq = state_to_wseq_[s];

  float logprob = lm_.GetNgramLogprob(ilabel, wseq);
  if (logprob == std::numeric_limits<float>::min()) {
    return false;
  }

  // Locates the next state in ConstArpaLm. Note that OOV and backoff have been
  // taken care of in ConstArpaLm.
  wseq.push_back(ilabel);
  while (wseq.size() >= lm_.NgramOrder()) {
    // History state has at most lm_.NgramOrder() -1 words in the state.
    wseq.erase(wseq.begin(), wseq.begin() + 1);
  }
  while (!lm_.HistoryStateExists(wseq)) {
    KALDI_ASSERT(wseq.size() > 0);
    wseq.erase(wseq.begin(), wseq.begin() + 1);
  }

  std::pair<const std::vector<Label>, StateId> wseq_state_pair(
      wseq, static_cast<Label>(state_to_wseq_.size()));

  // Attemps to insert the current <wseq_state_pair>. If the pair already exists
  // then it returns false.
  typedef MapType::iterator IterType;
  std::pair<IterType, bool> result = wseq_to_state_.insert(wseq_state_pair);

  // If the pair was just inserted, then also add it to <state_to_wseq_>.
  if (result.second == true)
    state_to_wseq_.push_back(wseq);

  // Creates the arc.
  oarc->ilabel = ilabel;
  oarc->olabel = ilabel;
  oarc->nextstate = result.first->second;
  oarc->weight = Weight(-logprob);

  return true;
}

bool BuildConstArpaLm(const ArpaParseOptions& options,
                      const std::string& arpa_rxfilename,
                      const std::string& const_arpa_wxfilename) {
  ConstArpaLmBuilder lm_builder(options);
  KALDI_LOG << "Reading " << arpa_rxfilename;
  Input ki(arpa_rxfilename);
  lm_builder.Read(ki.Stream());
  WriteKaldiObject(lm_builder, const_arpa_wxfilename, true);
  return true;
}

}  // namespace kaldi