// chain/language-model.cc

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

#include <algorithm>
#include <numeric>
#include "chain/language-model.h"
#include "util/simple-io-funcs.h"


namespace kaldi {
namespace chain {

void LanguageModelEstimator::AddCounts(const std::vector<int32> &sentence) {
  KALDI_ASSERT(opts_.ngram_order >= 2 && "--ngram-order must be >= 2");
  KALDI_ASSERT(opts_.ngram_order >= opts_.no_prune_ngram_order);
  int32 order = opts_.ngram_order;
  // 0 is used for left-context at the beginning of the file.. treat it as BOS.
  std::vector<int32> history(1, 0);
  std::vector<int32>::const_iterator iter = sentence.begin(),
      end = sentence.end();
  for (; iter != end; ++iter) {
    KALDI_ASSERT(*iter != 0);
    IncrementCount(history, *iter);
    history.push_back(*iter);
    if (history.size() >= order)
      history.erase(history.begin());
  }
  // Probability of end of sentence.  This will end up getting ignored later, but
  // it still makes a difference for probability-normalization reasons.
  IncrementCount(history, 0);
}

void LanguageModelEstimator::IncrementCount(const std::vector<int32> &history,
                                            int32 next_phone) {
  int32 lm_state_index = FindOrCreateLmStateIndexForHistory(history);
  if (lm_states_[lm_state_index].tot_count == 0) {
    num_active_lm_states_++;
  }
  lm_states_[lm_state_index].AddCount(next_phone, 1);
}

void LanguageModelEstimator::SetParentCounts() {
  int32 num_lm_states = lm_states_.size();
  for (int32 l = 0; l < num_lm_states; l++) {
    int32 this_count = lm_states_[l].tot_count;
    int32 l_iter = l;
    while (l_iter != -1) {
      lm_states_[l_iter].tot_count_with_parents += this_count;
      l_iter = lm_states_[l_iter].backoff_lmstate_index;
    }
  }
  for (int32 l = 0; l < num_lm_states; l++) {
    KALDI_ASSERT(lm_states_[l].tot_count_with_parents >=
                 lm_states_[l].tot_count);
  }
}

int32 LanguageModelEstimator::CheckActiveStates() const {
  int32 num_active_states = 0,
      num_lm_states = lm_states_.size(),
      num_basic_lm_states = 0;
  for (int32 l = 0; l < num_lm_states; l++) {
    if (lm_states_[l].tot_count != 0)
      num_active_states++;
    if (lm_states_[l].history.size() == opts_.no_prune_ngram_order - 1)
      num_basic_lm_states++;
  }
  KALDI_ASSERT(num_active_states == num_active_lm_states_);
  return num_basic_lm_states;
}

int32 LanguageModelEstimator::FindLmStateIndexForHistory(
    const std::vector<int32> &hist) const {
  MapType::const_iterator iter = hist_to_lmstate_index_.find(hist);
  if (iter == hist_to_lmstate_index_.end())
    return -1;
  else
    return iter->second;
}

int32 LanguageModelEstimator::FindNonzeroLmStateIndexForHistory(
    std::vector<int32> hist) const {
  while (1) {
    int32 l = FindLmStateIndexForHistory(hist);
    if (l == -1 || lm_states_[l].tot_count == 0) {
      // no such state or state has zero count.
      if (hist.empty())
        KALDI_ERR << "Error looking up LM state index for history "
                  << "(likely code bug)";
      hist.erase(hist.begin());  // back off.
    } else {
      return l;
    }
  }
}

int32 LanguageModelEstimator::FindOrCreateLmStateIndexForHistory(
    const std::vector<int32> &hist) {
  MapType::const_iterator iter = hist_to_lmstate_index_.find(hist);
  if (iter != hist_to_lmstate_index_.end())
    return iter->second;
  int32 ans = lm_states_.size();  // index of next element
  // next statement relies on default construct of LmState.
  lm_states_.resize(lm_states_.size() + 1);
  lm_states_.back().history = hist;
  hist_to_lmstate_index_[hist] = ans;
  // make sure backoff_lmstate_index is set, if needed.
  if (hist.size() >= opts_.no_prune_ngram_order) {
    // we need a backoff state to exist- create one if needed.
    std::vector<int32> backoff_hist(hist.begin() + 1,
                                    hist.end());

    int32 backoff_lm_state = FindOrCreateLmStateIndexForHistory(
        backoff_hist);
    lm_states_[ans].backoff_lmstate_index = backoff_lm_state;
  }
  return ans;
}

void LanguageModelEstimator::LmState::AddCount(int32 phone, int32 count) {
  std::map<int32, int32>::iterator iter = phone_to_count.find(phone);
  if (iter == phone_to_count.end())
    phone_to_count[phone] = count;
  else
    iter->second += count;
  tot_count += count;
}

void LanguageModelEstimator::LmState::Add(const LmState &other) {
  KALDI_ASSERT(&other != this);
  std::map<int32, int32>::const_iterator iter = other.phone_to_count.begin(),
      end = other.phone_to_count.end();
  for (; iter != end; ++iter)
    AddCount(iter->first, iter->second);
}

void LanguageModelEstimator::LmState::Clear() {
  phone_to_count.clear();
  tot_count = 0;
  tot_count_with_parents = false;
  backoff_allowed = false;
}

BaseFloat LanguageModelEstimator::LmState::LogLike() const {
  double ans = 0.0;
  int32 tot_count_check = 0;
  std::map<int32, int32>::const_iterator iter = phone_to_count.begin(),
      end = phone_to_count.end();
  for (; iter != end; ++iter) {
    int32 count = iter->second;
    tot_count_check += count;
    double prob = count * 1.0 / tot_count;
    ans += log(prob) * count;
  }
  KALDI_ASSERT(tot_count_check == tot_count);
  return ans;
}

void LanguageModelEstimator::InitializeQueue() {
  int32 num_lm_states = lm_states_.size();
  while (!queue_.empty()) queue_.pop();
  for (int32 l = 0; l < num_lm_states; l++) {
    lm_states_[l].backoff_allowed = BackoffAllowed(l);
    if (lm_states_[l].backoff_allowed) {
      BaseFloat like_change = BackoffLogLikelihoodChange(l);
      queue_.push(std::pair<BaseFloat,int32>(like_change, l));
    }
  }
}

BaseFloat LanguageModelEstimator::BackoffLogLikelihoodChange(
    int32 l) const {
  const LmState &lm_state = lm_states_.at(l);
  KALDI_ASSERT(lm_state.backoff_allowed && lm_state.backoff_lmstate_index >= 0);
  const LmState &backoff_lm_state = lm_states_.at(
      lm_state.backoff_lmstate_index);
  KALDI_ASSERT(lm_state.tot_count != 0);
  // if the backoff state has zero count, there would naturally be a zero
  // cost, but return -1e15 * (count of this lm state)... this encourages the
  // lowest-count state to be backed off first.
  if (backoff_lm_state.tot_count == 0)
    return -1.0e-15 * lm_state.tot_count;
  LmState sum_state(backoff_lm_state);
  sum_state.Add(lm_state);
  BaseFloat log_like_change =
      sum_state.LogLike() -
      lm_state.LogLike() -
      backoff_lm_state.LogLike();
  // log-like change should not be positive... give it a margin for round-off
  // error.
  KALDI_ASSERT(log_like_change < 0.1);
  if (log_like_change > 0.0)
    log_like_change = 0.0;
  return log_like_change;
}


void LanguageModelEstimator::DoBackoff() {
  int32 initial_active_states = num_active_lm_states_,
      target_num_lm_states = num_basic_lm_states_ + opts_.num_extra_lm_states;

  // create 3 intermediate targets and the final target.  Between each phase we'll
  // do InitializeQueue(), which will get us more exact values.
  int32 num_targets = 4;
  std::vector<int32> targets(num_targets);
  for (int32 t = 0; t < num_targets; t++) {
    // the targets get progressively closer to target_num_lm_states;
    targets[t] = initial_active_states +
        ((target_num_lm_states - initial_active_states) * (t + 1)) / num_targets;
  }
  KALDI_ASSERT(targets.back() == target_num_lm_states);

  for (int32 t = 0; t < num_targets; t++) {
    KALDI_VLOG(2) << "Backing off states, stage " << t;
    InitializeQueue();
    int32 this_target = targets[t];
    while (num_active_lm_states_ > this_target && !queue_.empty()) {
      BaseFloat like_change = queue_.top().first;
      int32 lm_state = queue_.top().second;
      queue_.pop();
      BaseFloat recomputed_like_change = BackoffLogLikelihoodChange(lm_state);
      if (!ApproxEqual(like_change, recomputed_like_change)) {
        // If it changed (i.e. we had a stale likelihood-change on the queue),
        // just put back the recomputed like-change on the queue and make no other
        // changes.
        KALDI_VLOG(2) << "Not backing off state, since like-change changed from "
                      << like_change << " to " << recomputed_like_change;
        queue_.push(std::pair<BaseFloat,int32>(recomputed_like_change, lm_state));
      } else {
        KALDI_VLOG(2) << "Backing off state with like-change = "
                      << recomputed_like_change;
        BackOffState(lm_state);
      }
    }
  }
  KALDI_LOG << "In LM [hard] backoff, target num states was "
            << num_basic_lm_states_ << " + --num-extra-lm-states="
            << opts_.num_extra_lm_states << " = " << target_num_lm_states
            << ", pruned from " << initial_active_states << " to "
            << num_active_lm_states_;
}

void LanguageModelEstimator::BackOffState(int32 l) {
  LmState &lm_state = lm_states_.at(l);
  KALDI_ASSERT(lm_state.backoff_allowed);
  KALDI_ASSERT(lm_state.backoff_lmstate_index >= 0);
  KALDI_ASSERT(lm_state.tot_count > 0);  // or shouldn't be backing it off.
  LmState &backoff_lm_state = lm_states_.at(lm_state.backoff_lmstate_index);
  bool backoff_state_had_backoff_allowed = backoff_lm_state.backoff_allowed;
  if (backoff_lm_state.tot_count != 0)
    num_active_lm_states_--;
  // add the counts of lm_state to backoff_lm_state.
  backoff_lm_state.Add(lm_state);
  // zero the counts in this lm_state.
  lm_state.Clear();
  backoff_lm_state.backoff_allowed = BackoffAllowed(
      lm_state.backoff_lmstate_index);

  if (!backoff_state_had_backoff_allowed &&
      backoff_lm_state.backoff_allowed) {
    // the backoff state would not have been in the queue, but is now allowed in
    // the queue.
    BaseFloat backoff_like_change = BackoffLogLikelihoodChange(
        lm_state.backoff_lmstate_index);
    queue_.push(std::pair<BaseFloat,int32>(backoff_like_change,
                                           lm_state.backoff_lmstate_index));
  }
}

int32 LanguageModelEstimator::AssignFstStates() {
  CheckActiveStates();
  int32 num_lm_states = lm_states_.size();
  int32 current_fst_state = 0;
  for (int32 l = 0; l < num_lm_states; l++)
    if (lm_states_[l].tot_count != 0)
      lm_states_[l].fst_state = current_fst_state++;
  KALDI_ASSERT(current_fst_state == num_active_lm_states_);
  return current_fst_state;
}

void LanguageModelEstimator::Estimate(fst::StdVectorFst *fst) {
  KALDI_LOG << "Estimating language model with --no-prune-ngram-order="
            << opts_.no_prune_ngram_order << ", --ngram-order="
            << opts_.ngram_order << ", --num-extra-lm-states="
            << opts_.num_extra_lm_states;
  SetParentCounts();
  num_basic_lm_states_ = CheckActiveStates();
  DoBackoff();
  int32 num_fst_states = AssignFstStates();
  OutputToFst(num_fst_states, fst);
}

int32 LanguageModelEstimator::FindInitialFstState() const {
  std::vector<int32> history(1, 0);
  int32 l = FindNonzeroLmStateIndexForHistory(history);
  KALDI_ASSERT(l != -1 && lm_states_[l].fst_state != -1);
  return lm_states_[l].fst_state;
}


bool LanguageModelEstimator::BackoffAllowed(int32 l) const {
  const LmState &lm_state = lm_states_.at(l);
  if (lm_state.history.size() < opts_.no_prune_ngram_order)
    return false;
  KALDI_ASSERT(lm_state.tot_count <= lm_state.tot_count_with_parents);
  if (lm_state.tot_count != lm_state.tot_count_with_parents)
    return false;
  if (lm_state.tot_count == 0)
    return false;
  // the next if-statement is an optimization where we skip the
  // following test if we know that it must always be true.
  if (lm_state.history.size() == opts_.ngram_order - 1)
    return true;
  std::map<int32, int32>::const_iterator
      iter = lm_state.phone_to_count.begin(),
      end = lm_state.phone_to_count.end();
  for (; iter != end; ++iter) {
    int32 phone = iter->first;
    if (phone != 0) {
      std::vector<int32> next_hist(lm_state.history);
      next_hist.push_back(phone);
      int32 next_lmstate = FindLmStateIndexForHistory(next_hist);
      if (next_lmstate != -1 &&
          lm_states_[next_lmstate].tot_count_with_parents != 0) {
        // backoff is not allowed because we need all the context we have
        // in order to make this transition; we can't afford to discard
        // the leftmost phone.
        return false;
      }
    }
  }
  return true;
}

void LanguageModelEstimator::OutputToFst(
    int32 num_states,
    fst::StdVectorFst *fst) const {
  KALDI_ASSERT(num_states == num_active_lm_states_);
  fst->DeleteStates();
  for (int32 i = 0; i < num_states; i++)
    fst->AddState();
  fst->SetStart(FindInitialFstState());

  int64 tot_count = 0;
  double tot_logprob = 0.0;

  int32 num_lm_states = lm_states_.size();
  // note: not all lm-states end up being 'active'.
  for (int32 l = 0; l < num_lm_states; l++) {
    const LmState &lm_state = lm_states_[l];
    if (lm_state.fst_state == -1)
      continue;
    int32 state_count = lm_state.tot_count;
    KALDI_ASSERT(state_count != 0);
    std::map<int32, int32>::const_iterator
        iter = lm_state.phone_to_count.begin(),
        end = lm_state.phone_to_count.end();
    for (; iter != end; ++iter) {
      int32 phone = iter->first, count = iter->second;
      BaseFloat logprob = log(count * 1.0 / state_count);
      tot_count += count;
      tot_logprob += logprob * count;
      if (phone == 0) {  // Interpret as final-prob.
        fst->SetFinal(lm_state.fst_state, fst::TropicalWeight(-logprob));
      } else {  // It becomes a transition.
        std::vector<int32> next_history(lm_state.history);
        next_history.push_back(phone);
        int32 dest_lm_state = FindNonzeroLmStateIndexForHistory(next_history),
            dest_fst_state = lm_states_[dest_lm_state].fst_state;
        KALDI_ASSERT(dest_fst_state != -1);
        fst->AddArc(lm_state.fst_state,
                    fst::StdArc(phone, phone, fst::TropicalWeight(-logprob),
                                dest_fst_state));
      }
    }
  }
  BaseFloat perplexity = exp(-(tot_logprob / tot_count));
  KALDI_LOG << "Total number of phone instances seen was " << tot_count;
  KALDI_LOG << "Perplexity on training data is: " << perplexity;
  KALDI_LOG << "Note: perplexity on unseen data will be infinity as there is "
            << "no smoothing.  This is by design, to reduce the number of arcs.";
  fst::Connect(fst);
  // Make sure that Connect does not delete any states.
  int32 num_states_connected = fst->NumStates();
  KALDI_ASSERT(num_states_connected == num_states);
  // arc-sort.  ilabel or olabel doesn't matter, it's an acceptor.
  fst::ArcSort(fst, fst::ILabelCompare<fst::StdArc>());
  KALDI_LOG << "Created phone language model with " << num_states
            << " states and " << fst::NumArcs(*fst) << " arcs.";
}

}  // namespace chain
}  // namespace kaldi