// fstext/deterministic-fst-inl.h

// Copyright 2011-2012 Gilles Boulianne
//                2014 Telepoint Global Hosting Service, LLC. (Author: David Snyder)
//           2012-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.

#ifndef KALDI_FSTEXT_DETERMINISTIC_FST_INL_H_
#define KALDI_FSTEXT_DETERMINISTIC_FST_INL_H_
#include "base/kaldi-common.h"
#include "fstext/fstext-utils.h"


namespace fst {
// Do not include this file directly.  It is included by deterministic-fst.h.

template<class Arc>
typename Arc::StateId
BackoffDeterministicOnDemandFst<Arc>::GetBackoffState(StateId s,
                                                      Weight *w) {
  ArcIterator<Fst<Arc> > aiter(fst_, s);
  if (aiter.Done()) // no arcs.
    return kNoStateId;
  const Arc &arc = aiter.Value();
  if (arc.ilabel == 0) {
    *w = arc.weight;
    return arc.nextstate;
  } else {
    return kNoStateId;
  }
}

template<class Arc>
typename Arc::Weight BackoffDeterministicOnDemandFst<Arc>::Final(StateId state) {
  Weight w = fst_.Final(state);
  if (w != Weight::Zero()) return w;
  Weight backoff_w;
  StateId backoff_state = GetBackoffState(state, &backoff_w);
  if (backoff_state == kNoStateId) return Weight::Zero();
  else return Times(backoff_w, this->Final(backoff_state));
}

template<class Arc>
BackoffDeterministicOnDemandFst<Arc>::BackoffDeterministicOnDemandFst(
    const Fst<Arc> &fst): fst_(fst) {
#ifdef KALDI_PARANOID
  KALDI_ASSERT(fst_.Properties(kILabelSorted|kIDeterministic, true) ==
               (kILabelSorted|kIDeterministic) &&
               "Input FST is not i-label sorted and deterministic.");
#endif
}

template<class Arc>
bool BackoffDeterministicOnDemandFst<Arc>::GetArc(
    StateId s, Label ilabel, Arc *oarc) {
  KALDI_ASSERT(ilabel != 0); //  We don't allow GetArc for epsilon.

  SortedMatcher<Fst<Arc> > sm(fst_, MATCH_INPUT, 1);
  sm.SetState(s);
  if (sm.Find(ilabel)) {
    const Arc &arc = sm.Value();
    *oarc = arc;
    return true;
  } else {
    Weight backoff_w;
    StateId backoff_state = GetBackoffState(s, &backoff_w);
    if (backoff_state == kNoStateId) return false;
    if (!this->GetArc(backoff_state, ilabel, oarc)) return false;
    oarc->weight = Times(oarc->weight, backoff_w);
    return true;
  }
}

template<class Arc>
UnweightedNgramFst<Arc>::UnweightedNgramFst(int n): n_(n) {
  // Starting state is an empty vector
  std::vector<Label> start_state;
  state_vec_.push_back(start_state);
  start_state_ = 0;
  state_map_[start_state] = 0;
}

template<class Arc>
bool UnweightedNgramFst<Arc>::GetArc(
  StateId s, Label ilabel, Arc *oarc) {

  // The state ids increment with each state we encounter.
  // if the assert fails, then we are trying to access
  // unseen states that are not immediately traversable.
  KALDI_ASSERT(static_cast<size_t>(s) < state_vec_.size());
  std::vector<Label> seq = state_vec_[s];
  // Update state info.
  seq.push_back(ilabel);
  if (seq.size() > n_-1) {
    // Remove oldest word in the history.
    seq.erase(seq.begin());
  }
  std::pair<const std::vector<Label>, StateId> new_state(
    seq,
    static_cast<Label>(state_vec_.size()));
  // Now get state id for destination state.
  typedef typename MapType::iterator IterType;
  std::pair<IterType, bool> result = state_map_.insert(new_state);
  if (result.second == true) {
    state_vec_.push_back(seq);
  }
  oarc->weight = Weight::One(); // Because the FST is unweightd.
  oarc->ilabel = ilabel;
  oarc->olabel = ilabel;
  oarc->nextstate = result.first->second; // The next state id.
  // All arcs can be matched.
  return true;
}

template<class Arc>
typename Arc::Weight UnweightedNgramFst<Arc>::Final(StateId state) {
  KALDI_ASSERT(state < static_cast<StateId>(state_vec_.size()));
  return Weight::One();
}

template<class Arc>
ComposeDeterministicOnDemandFst<Arc>::ComposeDeterministicOnDemandFst(
    DeterministicOnDemandFst<Arc> *fst1,
    DeterministicOnDemandFst<Arc> *fst2): fst1_(fst1), fst2_(fst2) {
  KALDI_ASSERT(fst1 != NULL && fst2 != NULL);
  if (fst1_->Start() == -1 || fst2_->Start() == -1) {
    start_state_ = -1;
    next_state_ = 0; // actually we don't care about this value.
  } else {
    start_state_ = 0;
    std::pair<StateId,StateId> start_pair(fst1_->Start(), fst2_->Start());
    state_map_[start_pair] = start_state_;
    state_vec_.push_back(start_pair);
    next_state_ = 1;
  }
}

template<class Arc>
typename Arc::Weight ComposeDeterministicOnDemandFst<Arc>::Final(StateId s) {
  KALDI_ASSERT(s < static_cast<StateId>(state_vec_.size()));
  const std::pair<StateId, StateId> &pr (state_vec_[s]);
  return Times(fst1_->Final(pr.first), fst2_->Final(pr.second));
}

template<class Arc>
bool ComposeDeterministicOnDemandFst<Arc>::GetArc(StateId s, Label ilabel,
                                                  Arc *oarc) {
  typedef typename MapType::iterator IterType;
  KALDI_ASSERT(ilabel != 0 &&
         "This program expects epsilon-free compact lattices as input");
  KALDI_ASSERT(s < static_cast<StateId>(state_vec_.size()));
  const std::pair<StateId, StateId> pr (state_vec_[s]);

  Arc arc1;
  if (!fst1_->GetArc(pr.first, ilabel, &arc1)) return false;
  if (arc1.olabel == 0) { // There is no output label on the
    // arc, so only the first state changes.
    std::pair<const std::pair<StateId, StateId>, StateId> new_value(
        std::pair<StateId, StateId>(arc1.nextstate, pr.second),
        next_state_);

    std::pair<IterType, bool> result = state_map_.insert(new_value);
    oarc->ilabel = ilabel;
    oarc->olabel = 0;
    oarc->nextstate = result.first->second;
    oarc->weight = arc1.weight;
    if (result.second == true) { // was inserted
      next_state_++;
      const std::pair<StateId, StateId> &new_pair (new_value.first);
      state_vec_.push_back(new_pair);
    }
    return true;
  }
  // There is an output label, so we need to traverse an arc on the
  // second fst also.
  Arc arc2;
  if (!fst2_->GetArc(pr.second, arc1.olabel, &arc2)) return false;
  std::pair<const std::pair<StateId, StateId>, StateId> new_value(
      std::pair<StateId, StateId>(arc1.nextstate, arc2.nextstate),
      next_state_);
  std::pair<IterType, bool> result =
      state_map_.insert(new_value);
  oarc->ilabel = ilabel;
  oarc->olabel = arc2.olabel;
  oarc->nextstate = result.first->second;
  oarc->weight = Times(arc1.weight, arc2.weight);
  if (result.second == true) { // was inserted
    next_state_++;
    const std::pair<StateId, StateId> &new_pair (new_value.first);
    state_vec_.push_back(new_pair);
  }
  return true;
}

template<class Arc>
inline size_t CacheDeterministicOnDemandFst<Arc>::GetIndex(
    StateId src_state, Label ilabel) {
  const StateId p1 = 26597, p2 = 50329; // these are two
  // values that I drew at random from a table of primes.
  // note: num_cached_arcs_ > 0.

  // We cast to size_t before the modulus, to ensure the
  // result is positive.
  return static_cast<size_t>(src_state * p1 + ilabel * p2) %
      static_cast<size_t>(num_cached_arcs_);
}

template<class Arc>
CacheDeterministicOnDemandFst<Arc>::CacheDeterministicOnDemandFst(
    DeterministicOnDemandFst<Arc> *fst,
    StateId num_cached_arcs): fst_(fst),
                              num_cached_arcs_(num_cached_arcs),
                              cached_arcs_(num_cached_arcs) {
  KALDI_ASSERT(num_cached_arcs > 0);
  for (StateId i = 0; i < num_cached_arcs; i++)
    cached_arcs_[i].first = kNoStateId; // Invalidate all elements of the cache.
}

template<class Arc>
bool CacheDeterministicOnDemandFst<Arc>::GetArc(StateId s, Label ilabel,
                                                Arc *oarc) {
  // Note: we don't cache anything in case a requested arc does not exist.
  // In the uses that we imagine this will be put to, essentially all the
  // requested arcs will exist.  This only affects efficiency.
  KALDI_ASSERT(s >= 0 && ilabel != 0);
  size_t index = this->GetIndex(s, ilabel);
  if (cached_arcs_[index].first == s &&
      cached_arcs_[index].second.ilabel == ilabel) {
    *oarc = cached_arcs_[index].second;
    return true;
  } else {
    Arc arc;
    if (fst_->GetArc(s, ilabel, &arc)) {
      cached_arcs_[index].first = s;
      cached_arcs_[index].second = arc;
      *oarc = arc;
      return true;
    } else {
      return false;
    }
  }
}

template<class Arc>
LmExampleDeterministicOnDemandFst<Arc>::LmExampleDeterministicOnDemandFst(
    void *lm, Label bos_symbol, Label eos_symbol):
    lm_(lm), bos_symbol_(bos_symbol), eos_symbol_(eos_symbol) {
  std::vector<Label> begin_state; // history state corresponding to beginning of sentence
  begin_state.push_back(bos_symbol); // Depending how your LM is set up, you might
  // want to have a history vector with more than one bos_symbol on it.

  state_vec_.push_back(begin_state);
  start_state_ = 0;
  state_map_[begin_state] = 0;
}

template<class Arc>
typename Arc::Weight LmExampleDeterministicOnDemandFst<Arc>::Final(StateId s) {
  KALDI_ASSERT(static_cast<size_t>(s) < state_vec_.size());
  // In a real version you would probably use the following variable somehow
  // (commenting it because it's generating warnings).
  // const std::vector<Label> &wseq = state_vec_[s];
  float log_prob = -0.5; // e.g. log_prob = lm->GetLogProb(wseq, eos_symbol_);
  return Weight(-log_prob); // assuming weight is FloatWeight.
}

template<class Arc>
bool LmExampleDeterministicOnDemandFst<Arc>::GetArc(
    StateId s, Label ilabel, Arc *oarc) {
  KALDI_ASSERT(static_cast<size_t>(s) < state_vec_.size());
  std::vector<Label> wseq = state_vec_[s];
  float log_prob = -0.25; // e.g. log_prob = lm->GetLogProb(wseq, ilabel);
  wseq.push_back(ilabel); // the code might be different if your histories are the
  // other way around.

  while (0) { // e.g. while !lm->HistoryStateExists(wseq)
    wseq.erase(wseq.begin(), wseq.begin() + 1); // remove most distant element of history.
    // note: if your histories are the other way round, you might just do
    // wseq.pop() here.
  }
  if (log_prob == -numeric_limits<float>::infinity()) { // assume this
    // is what happens if prob of the word is zero.  Some LMs will never
    // return zero.
    return false; // no arc.
  }
  std::pair<const std::vector<Label>, StateId> new_value(
      wseq,
      static_cast<Label>(state_vec_.size()));

  // Now get state id for destination state.
  typedef typename MapType::iterator IterType;
  std::pair<IterType, bool> result = state_map_.insert(new_value);
  if (result.second == true) // was inserted
    state_vec_.push_back(wseq);
  oarc->ilabel = ilabel;
  oarc->olabel = ilabel;
  oarc->nextstate = result.first->second; // the next-state id.
  oarc->weight = Weight(-log_prob);
  return true;
}


template<class Arc>
void ComposeDeterministicOnDemand(const Fst<Arc> &fst1,
                                  DeterministicOnDemandFst<Arc> *fst2,
                                  MutableFst<Arc> *fst_composed) {
  typedef typename Arc::Weight Weight;
  typedef typename Arc::StateId StateId;
  typedef std::pair<StateId, StateId> StatePair;
  typedef unordered_map<StatePair, StateId,
    kaldi::PairHasher<StateId> > MapType;
  typedef typename MapType::iterator IterType;

  fst_composed->DeleteStates();

  MapType state_map;
  std::queue<StatePair> state_queue;

  // Set start state in fst_composed.
  StateId s1 = fst1.Start(),
          s2 = fst2->Start(),
          start_state = fst_composed->AddState();
  StatePair start_pair(s1, s2);
  state_queue.push(start_pair);
  fst_composed->SetStart(start_state);
  // A mapping between pairs of states in fst1 and fst2 and the corresponding
  // state in fst_composed.
  std::pair<const StatePair, StateId> start_map(start_pair, start_state);
  std::pair<IterType, bool> result = state_map.insert(start_map);
  KALDI_ASSERT(result.second == true);

  while (!state_queue.empty()) {
    StatePair q = state_queue.front();
    StateId q1 = q.first,
            q2 = q.second;
    state_queue.pop();
    // If the product of the final weights of the two fsts is non-zero then
    // we can set a final-prob in fst_composed
    Weight final_weight = Times(fst1.Final(q1), fst2->Final(q2));
    if (final_weight != Weight::Zero()) {
      KALDI_ASSERT(state_map.find(q) != state_map.end());
      fst_composed->SetFinal(state_map[q], final_weight);
    }

    // for each pair of edges from fst1 and fst2 at q1 and q2.
    for (ArcIterator<Fst<Arc> > aiter(fst1, q1); !aiter.Done(); aiter.Next()) {
      const Arc &arc1 = aiter.Value();
      Arc arc2;
      StatePair next_pair;
      StateId next_state1 = arc1.nextstate,
              next_state2,
              next_state;
      // If there is an epsilon on the arc of fst1 we transition to the next
      // state but keep fst2 at the current state.
      if (arc1.olabel == 0) {
        next_state2 = q2;
      } else {
        bool match = fst2->GetArc(q2, arc1.olabel, &arc2);
        if (!match)  // There is no matching arc -> nothing to do.
          continue;
        next_state2 = arc2.nextstate;
      }
      next_pair = StatePair(next_state1, next_state2);
      IterType sitr = state_map.find(next_pair);
      // If sitr == state_map.end() then the state isn't in fst_composed yet.
      if (sitr == state_map.end()) {
        next_state = fst_composed->AddState();
        std::pair<const StatePair, StateId> new_state(
          next_pair, next_state);
        std::pair<IterType, bool> result = state_map.insert(new_state);
        // Since we already checked if state_map contained new_state,
        // it should always be added if we reach here.
        KALDI_ASSERT(result.second == true);
        state_queue.push(next_pair);
      // If sitr != state_map.end() then the next state is already in
      // the state_map.
      } else {
        next_state = sitr->second;
      }
      if (arc1.olabel == 0) {
        fst_composed->AddArc(state_map[q], Arc(arc1.ilabel, 0, arc1.weight,
                                               next_state));
      } else {
        fst_composed->AddArc(state_map[q], Arc(arc1.ilabel, arc2.olabel,
          Times(arc1.weight, arc2.weight), next_state));
      }
    }
  }
}


// we are doing *fst_composed = Compose(Inverse(*left), right).
template<class Arc>
void ComposeDeterministicOnDemandInverse(const Fst<Arc> &right,
                                         DeterministicOnDemandFst<Arc> *left,
                                         MutableFst<Arc> *fst_composed) {
  typedef typename Arc::Weight Weight;
  typedef typename Arc::StateId StateId;
  typedef std::pair<StateId, StateId> StatePair;
  typedef unordered_map<StatePair, StateId,
    kaldi::PairHasher<StateId> > MapType;
  typedef typename MapType::iterator IterType;

  fst_composed->DeleteStates();

  // the queue and map contain pairs (state-in-left, state-in-right)
  MapType state_map;
  std::queue<StatePair> state_queue;

  // Set start state in fst_composed.
  StateId s_left = left->Start(),
      s_right = right.Start();
  if (s_left == kNoStateId || s_right == kNoStateId)
    return;  // Empty result.
  StatePair start_pair(s_left, s_right);
  StateId start_state = fst_composed->AddState();
  state_queue.push(start_pair);
  fst_composed->SetStart(start_state);
  // A mapping between pairs of states in *left and right, and the corresponding
  // state in fst_composed.
  std::pair<const StatePair, StateId> start_map(start_pair, start_state);
  std::pair<IterType, bool> result = state_map.insert(start_map);
  KALDI_ASSERT(result.second == true);

  while (!state_queue.empty()) {
    StatePair q = state_queue.front();
    StateId q_left = q.first,
            q_right = q.second;
    state_queue.pop();
    // If the product of the final weights of the two fsts is non-zero then
    // we can set a final-prob in fst_composed
    Weight final_weight = Times(left->Final(q_left), right.Final(q_right));
    if (final_weight != Weight::Zero()) {
      KALDI_ASSERT(state_map.find(q) != state_map.end());
      fst_composed->SetFinal(state_map[q], final_weight);
    }

    for (ArcIterator<Fst<Arc> > aiter(right, q_right); !aiter.Done(); aiter.Next()) {
      const Arc &arc_right = aiter.Value();
      Arc arc_left;
      StatePair next_pair;
      StateId next_state_right = arc_right.nextstate,
              next_state_left,
              next_state;
      // If there is an epsilon on the input side of the rigth arc, we
      // transition to the next state of the output but keep 'left' at the
      // current state.
      if (arc_right.ilabel == 0) {
        next_state_left = q_left;
      } else {
        bool match = left->GetArc(q_left, arc_right.ilabel, &arc_left);
        if (!match)  // There is no matching arc -> nothing to do.
          continue;
        // the next 'swap' is because we are composing with the inverse of
        // *left.  Just removing the swap statement wouldn't let us compose
        // with non-inverted *left though, because the GetArc function call
        // above interprets the second argument as an ilabel not an olabel.
        std::swap(arc_left.ilabel, arc_left.olabel);
        next_state_left = arc_left.nextstate;
      }
      next_pair = StatePair(next_state_left, next_state_right);
      IterType sitr = state_map.find(next_pair);
      // If sitr == state_map.end() then the state isn't in fst_composed yet.
      if (sitr == state_map.end()) {
        next_state = fst_composed->AddState();
        std::pair<const StatePair, StateId> new_state(
          next_pair, next_state);
        std::pair<IterType, bool> result = state_map.insert(new_state);
        // Since we already checked if state_map contained new_state,
        // it should always be added if we reach here.
        KALDI_ASSERT(result.second == true);
        state_queue.push(next_pair);
      // If sitr != state_map.end() then the next state is already in
      // the state_map.
      } else {
        next_state = sitr->second;
      }
      if (arc_right.ilabel == 0) {
        // we didn't get an actual arc from the left FST.
        fst_composed->AddArc(state_map[q], Arc(0, arc_right.olabel,
                                               arc_right.weight,
                                               next_state));
      } else {
        fst_composed->AddArc(state_map[q],
                             Arc(arc_left.ilabel, arc_right.olabel,
                                 Times(arc_left.weight, arc_right.weight),
                                 next_state));
      }
    }
  }
}



} // end namespace fst


#endif