// sampling-lm.cc

// Copyright 2017  Ke Li
//           2017  Johns Hopkins University (author: Daniel Povey)

// 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,
// MERCHANTABILITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.

#include "rnnlm/sampling-lm.h"

namespace kaldi {
namespace rnnlm {

// This function reads in each ngram line from an ARPA file
void SamplingLm::ConsumeNGram(const NGram& ngram) {
  int32 cur_order = ngram.words.size(),
      max_order = Order();
  int32 word = ngram.words.back();  // word is the last word in a ngram term
  KALDI_ASSERT(cur_order > 0 && word > 0);

  if (cur_order == 1) {
    // unigram
    if (unigram_probs_.size() <= static_cast<size_t>(word))
      unigram_probs_.resize(static_cast<size_t>(word + 1), 0.0);
    KALDI_ASSERT(unigram_probs_[word] == 0.0);  // or repeated unigram.
    unigram_probs_[word] = Exp(ngram.logprob);
    if (ngram.backoff != 0.0)
      higher_order_probs_[cur_order - 1][ngram.words].backoff_prob =
        Exp(ngram.backoff);
  } else {
    HistType history(ngram.words.begin(), ngram.words.end() - 1);
    // Note: we'll later on change the probability, subtracting the
    // part that is due to backoff.  This change of format is
    // convenient for our application.
    // ngram.logprob has already been converted to log-base e at
    // this point.
    higher_order_probs_[cur_order - 2][history].words_and_probs.push_back(
        std::pair<int32, BaseFloat>(word, Exp(ngram.logprob)));
    if (ngram.backoff != 0.0) {
      KALDI_ASSERT(cur_order != max_order);
      higher_order_probs_[cur_order - 1][ngram.words].backoff_prob =
          Exp(ngram.backoff);
    }
  }
}

void SamplingLm::HeaderAvailable() {
  unigram_probs_.reserve(NgramCounts()[0] + 100);
  // e.g. for a trigram LM we store bigram and trigram
  // history states in probs_, while unigram_probs_ stores
  // the unigram probabilities.
  int32 ngram_order = NgramCounts().size();
  higher_order_probs_.resize(ngram_order - 1);
}

BaseFloat SamplingLm::GetProbWithBackoff(
    const std::vector<int32> &history,
    const HistoryState *state,
    int32 word) const {
  if (state == NULL) {
    int32 order = history.size() + 1;
    if (order == 1) {
      KALDI_ASSERT(static_cast<size_t>(word) < unigram_probs_.size());
      return unigram_probs_[word];
    } else {
      std::unordered_map<HistType, HistoryState, VectorHasher<int32> >::const_iterator
          hist_iter = higher_order_probs_[order - 2].find(history);
      KALDI_ASSERT(hist_iter != higher_order_probs_[order - 2].end());
      // it's not optimally efficient to recurse here, but this is on a code
      // path that will rarely be taken in practice.
      return GetProbWithBackoff(history, &(hist_iter->second), word);
    }
  } else {
    std::pair<int32, BaseFloat> p(word, 0.0);
    std::vector<std::pair<int32, BaseFloat> >::const_iterator iter =
        std::lower_bound(state->words_and_probs.begin(),
                         state->words_and_probs.end(), p);
    if (iter != state->words_and_probs.end() && iter->first == word) {
      // the probability for this word was given in this history state.  (note:
      // we assume that at the time this function is called, the entire
      // probability is present here, as it is in the ARPA format LM.  See
      // documentation for this function for more explanation.
      return iter->second;
    } else {
      // we have to back off.
      std::vector<int32> backoff_history(history.begin() + 1,
                                         history.end());
      return state->backoff_prob *
          GetProbWithBackoff(backoff_history, NULL, word);
    }
  }
}

void SamplingLm::EnsureHistoryStatesSorted() {
  for (size_t i = 0; i < higher_order_probs_.size(); i++) {
    std::unordered_map<HistType, HistoryState, VectorHasher<int32> >::iterator
        iter = higher_order_probs_[i].begin(),
        end = higher_order_probs_[i].end();
    for (; iter != end; ++iter)
      std::sort(iter->second.words_and_probs.begin(),
                iter->second.words_and_probs.end());
  }
}

void SamplingLm::ReadComplete() {
  EnsureHistoryStatesSorted();
  int32 max_order = Order();
  for (int32 order = max_order; order >= 2; order--) {
    std::unordered_map<HistType, HistoryState, VectorHasher<int32> >
        &this_map = higher_order_probs_[order - 2];
    std::unordered_map<HistType, HistoryState,
        VectorHasher<int32> >::iterator
        hist_iter = this_map.begin(), hist_end = this_map.end();
    for (; hist_iter != hist_end; ++hist_iter) {
      const HistType &history = hist_iter->first;
      HistoryState &history_state = hist_iter->second;
      BaseFloat backoff_prob = history_state.backoff_prob;
      HistoryState *backoff_state;
      HistType backoff_history(history.begin() + 1, history.end());
      if (order == 2) backoff_state = NULL;  // unigram has different format.
      else backoff_state = &(higher_order_probs_[order - 3][backoff_history]);

      std::vector<std::pair<int32, BaseFloat> >::iterator
          word_iter = history_state.words_and_probs.begin(),
          word_end = history_state.words_and_probs.end();
      double total_prob_after_subtracting = 0.0;
      for (; word_iter != word_end; ++word_iter) {
        int32 word = word_iter->first;
        BaseFloat prob = word_iter->second;
        // OK, we want to subtract the backoff part.
        BaseFloat backoff_part_of_prob = backoff_prob *
            GetProbWithBackoff(backoff_history, backoff_state, word);
        if (backoff_part_of_prob > 1.01 * prob) {
          KALDI_WARN << "Backoff part of prob is larger than prob itself: "
                     << backoff_part_of_prob << " > " << prob
                     << ".  This may mean your language model was not "
                     << "Kneser-Ney 'with addition'.  We advise to use "
                     << "Kneser-Ney with addition or some other type of "
                     << "LM 'with addition'.";
        }
        // OK, this could now be negative.  This shouldn't matter
        BaseFloat new_prob = prob - backoff_part_of_prob;
        word_iter->second = new_prob;
        total_prob_after_subtracting += new_prob;
      }
      BaseFloat new_total = total_prob_after_subtracting + backoff_prob;
      if (fabs(new_total - 1.0) > 0.01)
        KALDI_WARN << "Expected LM-state to sum to one, got "
                   << new_total;
    }
  }
}

void SamplingLm::AddBackoffToHistoryStates(
    const WeightedHistType &histories,
    WeightedHistType *histories_closure,
    BaseFloat *total_weight_out,
    BaseFloat *unigram_weight_out) const {
  // the implementation of this function is not as efficient as it could be,
  // but it should not dominate.
  std::vector<std::pair<HistType, BaseFloat> >::const_iterator
      histories_iter = histories.begin(), histories_end = histories.end();
  int32 max_order = Order();
  std::unordered_map<HistType, BaseFloat,
      VectorHasher<int32> > hist_to_weight_map;
  double total_weight = 0.0, total_unigram_weight = 0.0;
  for (; histories_iter != histories_end; ++histories_iter) {
    std::vector<int32> history = histories_iter->first;
    int32 cur_hist_len = history.size();
    BaseFloat weight = histories_iter->second;
    total_weight += weight;
    KALDI_ASSERT(history.size() <= max_order - 1 && weight > 0);

    // back off until the history exists or until we reached the unigram state.
    while (cur_hist_len > 0 &&
           higher_order_probs_[cur_hist_len - 1].count(history) == 0) {
      history.erase(history.begin());
      cur_hist_len--;
    }
    // OK, the history-state exists.
    while (cur_hist_len > 0) {
      hist_to_weight_map[history] += weight;
      std::unordered_map<HistType, HistoryState, VectorHasher<int32> >::const_iterator
          iter = higher_order_probs_[cur_hist_len - 1].find(history);
      KALDI_ASSERT(iter != higher_order_probs_[cur_hist_len - 1].end());
      weight *= iter->second.backoff_prob;
      history.erase(history.begin());
      cur_hist_len--;
    }
    // at this point, 'history' is empty and 'weight' is the unigram
    // backoff weight for this history state.
    total_unigram_weight += weight;
  }
  histories_closure->clear();
  histories_closure->resize(hist_to_weight_map.size());
  std::unordered_map<HistType, BaseFloat, VectorHasher<int32> >::iterator
      hist_to_weight_iter = hist_to_weight_map.begin(),
      hist_to_weight_end = hist_to_weight_map.end();
  size_t pos = 0;
  for (; hist_to_weight_iter != hist_to_weight_end; ++hist_to_weight_iter) {
    (*histories_closure)[pos].first = hist_to_weight_iter->first;
    (*histories_closure)[pos].second = hist_to_weight_iter->second;
    pos++;
  }
  *total_weight_out = total_weight;
  *unigram_weight_out = total_unigram_weight;
  KALDI_ASSERT(pos == hist_to_weight_map.size());
}


BaseFloat SamplingLm::GetDistribution(
    const WeightedHistType &histories,
    std::vector<std::pair<int32, BaseFloat> > *non_unigram_probs_out) const {
  std::unordered_map<int32, BaseFloat> non_unigram_probs_temp;
  // Call the other version of GetDistribution().
  BaseFloat ans = GetDistribution(histories, &non_unigram_probs_temp);
  non_unigram_probs_out->clear();
  non_unigram_probs_out->reserve(non_unigram_probs_temp.size());
  non_unigram_probs_out->insert(non_unigram_probs_out->end(),
                                non_unigram_probs_temp.begin(),
                                non_unigram_probs_temp.end());
  std::sort(non_unigram_probs_out->begin(),
            non_unigram_probs_out->end());
  return ans;
}

BaseFloat SamplingLm::GetDistribution(
    const WeightedHistType &histories,
    std::unordered_map<int32, BaseFloat> *non_unigram_probs) const {
  WeightedHistType histories_closure;
  BaseFloat total_weight, total_unigram_weight;
  AddBackoffToHistoryStates(histories, &histories_closure,
                            &total_weight, &total_unigram_weight);
  non_unigram_probs->clear();
  double total_weight_check = total_unigram_weight;
  WeightedHistType::const_iterator iter = histories_closure.begin(),
      end = histories_closure.end();
  for (; iter != end; ++iter) {
    const HistType &history = iter->first;
    BaseFloat hist_weight = iter->second;
    int32 order = history.size() + 1;
    KALDI_ASSERT(order > 1);  // unigram history is not included at this point.
    std::unordered_map<HistType, HistoryState,
        VectorHasher<int32> >::const_iterator it_hist =
           higher_order_probs_[order - 2].find(history);
    KALDI_ASSERT(it_hist != higher_order_probs_[order - 2].end());
    std::vector<std::pair<int32, BaseFloat> >::const_iterator
        word_iter = it_hist->second.words_and_probs.begin(),
        word_end = it_hist->second.words_and_probs.end();
    for (; word_iter != word_end; ++word_iter) {
      int32 word = word_iter->first;
      BaseFloat prob = word_iter->second;
      // note: if 'word' was not in the map, it's as if it were zero, for C++
      // version >= C++11; search for unordered_map value initialization for
      // explanation
      (*non_unigram_probs)[word] += prob * hist_weight;
      total_weight_check += prob * hist_weight;
    }
  }
  // Check that 'total_weight' and 'total_weight_check' are
  // the same.  'total_weight' is the total of the of the .second
  // member of the input 'histories', and 'total_weight_check' is the
  // total weight of 'non_unigrm_probs' plus 'total_unigram_weight'.
  // Essentially this is a check that the distribution given
  // by the ARPA file (and as processed by us) sums to one for each
  // history state.  If this check fails, it could either be
  // a problem with this code, or an issue with the software that
  // created the ARPA file.
  if (fabs(total_weight - total_weight_check) >
      0.01 * total_weight) {
    static int32 num_times_warned = 0;
    if (num_times_warned < 10) {
      KALDI_WARN << "Total weight does not have expected value (problem in "
          "your ARPA file, or this code).  Won't warn >10 times.";
      num_times_warned++;
    }
  }
  KALDI_ASSERT(total_unigram_weight > 0.0);
  return total_unigram_weight;
}

SamplingLm::SamplingLm(const SamplingLmEstimator &estimator):
    ArpaFileParser(ArpaParseOptions(), NULL),
    unigram_probs_(estimator.unigram_probs_),
    higher_order_probs_(estimator.history_states_.size() - 1) {
  for (int32 o = 2;
       o <= static_cast<int32>(estimator.history_states_.size()); o++) {
    higher_order_probs_[o-2].reserve(estimator.history_states_[o-1].size());
    unordered_map<std::vector<int32>, SamplingLmEstimator::HistoryState*,
                  VectorHasher<int32> >::const_iterator
        iter = estimator.history_states_[o-1].begin(),
        end =  estimator.history_states_[o-1].end();
    for (; iter != end; ++iter) {
      const std::vector<int32> &history = iter->first;
      const SamplingLmEstimator::HistoryState &src_state = *(iter->second);
      // the next statement adds a history state to the map.
      HistoryState &dest_state = higher_order_probs_[o-2][history];
      BaseFloat inv_total_count = BaseFloat(1.0) / src_state.total_count;
      dest_state.backoff_prob = src_state.backoff_count * inv_total_count;
      dest_state.words_and_probs.resize(src_state.counts.size());
      std::vector<SamplingLmEstimator::Count>::const_iterator
          src_iter = src_state.counts.begin(),
          src_end = src_state.counts.end();
      std::vector<std::pair<int32, BaseFloat> >::iterator
          dest_iter = dest_state.words_and_probs.begin();
      for (; src_iter != src_end; ++src_iter, ++dest_iter) {
        dest_iter->first = src_iter->word;
        dest_iter->second = inv_total_count * src_iter->count;
      }
    }
  }
}

void SamplingLm::Write(std::ostream &os, bool binary) const {
  WriteToken(os, binary, "<SamplingLm>");
  WriteToken(os, binary, "<Order>");
  int32 order = higher_order_probs_.size() + 1;
  WriteBasicType(os, binary, order);
  WriteToken(os, binary, "<VocabSize>");
  int32 vocab_size = unigram_probs_.size();
  WriteBasicType(os, binary, vocab_size);
  KALDI_ASSERT(!unigram_probs_.empty());
  // we have read and write functions in class Vector, so use that.
  SubVector<BaseFloat> probs(const_cast<BaseFloat*>(&(unigram_probs_[0])),
                             static_cast<int32>(unigram_probs_.size()));
  probs.Write(os, binary);
  for (int32 o = 2; o <= order; o++) {
    WriteToken(os, binary, "<StatesOfOrder>");
    WriteBasicType(os, binary, o);
    WriteToken(os, binary, "<NumStates>");
    int32 num_states = higher_order_probs_[o-2].size();
    WriteBasicType(os, binary, num_states);

    unordered_map<std::vector<int32>, HistoryState,
                  VectorHasher<int32> >::const_iterator
        iter = higher_order_probs_[o-2].begin(),
        end = higher_order_probs_[o-2].end();
    for (; iter != end; ++iter ){
      const std::vector<int32> &history = iter->first;
      const HistoryState &state = iter->second;
      WriteIntegerVector(os, binary, history);
      WriteBasicType(os, binary, state.backoff_prob);
      int32 num_words = state.words_and_probs.size();
      WriteBasicType(os, binary, num_words);
      for (int32 i = 0; i < num_words; i++) {
        WriteBasicType(os, binary, state.words_and_probs[i].first);
        WriteBasicType(os, binary, state.words_and_probs[i].second);
      }
      if (!binary) os << std::endl;
    }
  }
  WriteToken(os, binary, "</SamplingLm>");
}


void SamplingLm::Read(std::istream &is, bool binary) {
  ExpectToken(is, binary, "<SamplingLm>");
  ExpectToken(is, binary, "<Order>");
  int32 order;
  ReadBasicType(is, binary, &order);
  KALDI_ASSERT(order >= 1 && order < 100);
  higher_order_probs_.resize(order - 1);
  ExpectToken(is, binary, "<VocabSize>");
  int32 vocab_size;
  ReadBasicType(is, binary, &vocab_size);
  unigram_probs_.resize(vocab_size);
  // we have read and write functions in class Vector, so use that.
  SubVector<BaseFloat> probs(&(unigram_probs_[0]), vocab_size);
  probs.Read(is, binary);
  for (int32 o = 2; o <= order; o++) {
    ExpectToken(is, binary, "<StatesOfOrder>");
    int32 o2;
    ReadBasicType(is, binary, &o2);
    KALDI_ASSERT(o2 == o);
    int32 num_states;
    ExpectToken(is, binary, "<NumStates>");
    ReadBasicType(is, binary, &num_states);
    higher_order_probs_[o-2].reserve(num_states);
    for  (int32 s = 0; s < num_states; s++) {
      std::vector<int32> history;
      ReadIntegerVector(is, binary, &history);
      HistoryState &state = higher_order_probs_[o-2][history];
      ReadBasicType(is, binary, &(state.backoff_prob));
      int32 num_words;
      ReadBasicType(is, binary, &num_words);
      KALDI_ASSERT(num_words >= 0);
      state.words_and_probs.resize(num_words);
      for (int32 i = 0; i < num_words; i++) {
        ReadBasicType(is, binary, &(state.words_and_probs[i].first));
        ReadBasicType(is, binary, &(state.words_and_probs[i].second));
      }
    }
  }
  ExpectToken(is, binary, "</SamplingLm>");
}

// TODO: delete if unused.
void SamplingLm::Swap(SamplingLm *other) {
  unigram_probs_.swap(other->unigram_probs_);
  higher_order_probs_.swap(other->higher_order_probs_);
}

}  // namespace rnnlm
}  // namespace kaldi