// fstext/factor-inl.h
  
  // Copyright 2009-2011  Microsoft Corporation
  
  // 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_FACTOR_INL_H_
  #define KALDI_FSTEXT_FACTOR_INL_H_
  
  #include "util/stl-utils.h"
  // Do not include this file directly.  It is included by factor.h.
  
  namespace fst {
  
  // GetStateProperties takes in an FST and a number "max_state" which is the
  // highest numbered state in the FST (this could be fst.NumStates()-1 for an
  // ExpandedFst, or derived from some kind of traversal).  It outputs a vector
  // numbered from 0..max_state, of type FstStateProperties which is a bitmask
  // with information about the states.
  
  // GetStateProperties has not been tested directly (only implicitly via
  // testing Factor).
  template<class Arc>
  void GetStateProperties(const Fst<Arc> &fst,
                          typename Arc::StateId max_state,
                          vector<StatePropertiesType> *props) {
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Weight Weight;
    assert(props != NULL);
    props->clear();
    if (fst.Start() < 0) return;  // Empty fst.
    props->resize(max_state+1, 0);
    assert(fst.Start() <= max_state);
    (*props)[fst.Start()] |= kStateInitial;
    for (StateId s = 0; s <= max_state; s++) {
      StatePropertiesType &s_info = (*props)[s];
      for (ArcIterator<Fst<Arc> > aiter(fst, s); !aiter.Done(); aiter.Next()) {
        const Arc &arc = aiter.Value();
        if (arc.ilabel != 0) s_info |= kStateIlabelsOut;
        if (arc.olabel != 0) s_info |= kStateOlabelsOut;
        StateId nexts = arc.nextstate;
        assert(nexts <= max_state);  // or input was invalid.
        StatePropertiesType &nexts_info = (*props)[nexts];
        if (s_info&kStateArcsOut) s_info |= kStateMultipleArcsOut;
        s_info |= kStateArcsOut;
        if (nexts_info&kStateArcsIn) nexts_info |= kStateMultipleArcsIn;
        nexts_info |= kStateArcsIn;
      }
      if (fst.Final(s) != Weight::Zero())  s_info |= kStateFinal;
    }
  }
  
  
  
  template<class Arc, class I>
  void Factor(const Fst<Arc> &fst, MutableFst<Arc> *ofst,
                 vector<vector<I> > *symbols_out) {
    KALDI_ASSERT_IS_INTEGER_TYPE(I);
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Label Label;
    typedef typename Arc::Weight Weight;
    assert(symbols_out != NULL);
    ofst->DeleteStates();
    if (fst.Start() < 0) return;  // empty FST.
    vector<StateId> order;
    DfsOrderVisitor<Arc> dfs_order_visitor(&order);
    DfsVisit(fst, &dfs_order_visitor);
    assert(order.size() > 0);
    StateId max_state = *(std::max_element(order.begin(), order.end()));
    vector<StatePropertiesType> state_properties;
    GetStateProperties(fst, max_state, &state_properties);
  
    vector<bool> remove(max_state+1);  // if true, will remove this state.
  
    // Now identify states that will be removed (made the middle of a chain).
    // The basic rule is that if the FstStateProperties equals
    // (kStateArcsIn|kStateArcsOut) or (kStateArcsIn|kStateArcsOut|kStateIlabelsOut),
    // then it is in the middle of a chain.  This eliminates state with
    // multiple input or output arcs, final states, and states with arcs out
    // that have olabels [we assume these are pushed to the left, so occur on the
    // 1st arc of a chain.
  
    for (StateId i = 0; i <= max_state; i++)
      remove[i] = (state_properties[i] == (kStateArcsIn|kStateArcsOut)
                   || state_properties[i] == (kStateArcsIn|kStateArcsOut|kStateIlabelsOut));
    vector<StateId> state_mapping(max_state+1, kNoStateId);
  
    typedef unordered_map<vector<I>, Label, kaldi::VectorHasher<I> > SymbolMapType;
    SymbolMapType symbol_mapping;
    Label symbol_counter = 0;
    {
      vector<I> eps;
      symbol_mapping[eps] = symbol_counter++;
    }
    vector<I> this_sym;  // a temporary used inside the loop.
    for (size_t i = 0; i < order.size(); i++) {
      StateId state = order[i];
      if (!remove[state]) {  // Process this state...
        StateId &new_state = state_mapping[state];
        if (new_state == kNoStateId) new_state = ofst->AddState();
        for (ArcIterator<Fst<Arc> > aiter(fst, state); !aiter.Done(); aiter.Next()) {
          Arc arc = aiter.Value();
          if (arc.ilabel == 0) this_sym.clear();
          else {
            this_sym.resize(1);
            this_sym[0] = arc.ilabel;
          }
          while (remove[arc.nextstate]) {
            ArcIterator<Fst<Arc> > aiter2(fst, arc.nextstate);
            assert(!aiter2.Done());
            const Arc &nextarc = aiter2.Value();
            arc.weight = Times(arc.weight, nextarc.weight);
            assert(nextarc.olabel == 0);
            if (nextarc.ilabel != 0) this_sym.push_back(nextarc.ilabel);
            assert(static_cast<Label>(static_cast<I>(nextarc.ilabel))
                   == nextarc.ilabel); // check within integer range.
            arc.nextstate = nextarc.nextstate;
          }
          StateId &new_nextstate = state_mapping[arc.nextstate];
          if (new_nextstate == kNoStateId) new_nextstate = ofst->AddState();
          arc.nextstate = new_nextstate;
          if (symbol_mapping.count(this_sym) != 0) arc.ilabel = symbol_mapping[this_sym];
          else arc.ilabel = symbol_mapping[this_sym] = symbol_counter++;
          ofst->AddArc(new_state, arc);
        }
        if (fst.Final(state) != Weight::Zero())
          ofst->SetFinal(new_state, fst.Final(state));
      }
    }
    ofst->SetStart(state_mapping[fst.Start()]);
  
    // Now output the symbol sequences.
    symbols_out->resize(symbol_counter);
    for (typename SymbolMapType::const_iterator iter = symbol_mapping.begin();
        iter != symbol_mapping.end(); ++iter) {
      (*symbols_out)[iter->second] = iter->first;
    }
  }
  
  template<class Arc>
  void Factor(const Fst<Arc> &fst, MutableFst<Arc> *ofst1,
              MutableFst<Arc> *ofst2) {
    typedef typename Arc::Label Label;
    vector<vector<Label> > symbols;
    Factor(fst, ofst2, &symbols);
    CreateFactorFst(symbols, ofst1);
  }
  
  template<class Arc, class I>
  void ExpandInputSequences(const vector<vector<I> > &sequences,
                            MutableFst<Arc> *fst) {
    KALDI_ASSERT_IS_INTEGER_TYPE(I);
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Label Label;
    typedef typename Arc::Weight Weight;
    fst->SetInputSymbols(NULL);
    size_t size = sequences.size();
    if (sequences.size() > 0) assert(sequences[0].size() == 0);  // should be eps.
    StateId num_states_at_start = fst->NumStates();
    for (StateId s = 0; s < num_states_at_start; s++) {
      StateId num_arcs = fst->NumArcs(s);
      for (StateId aidx = 0; aidx < num_arcs; aidx++) {
        ArcIterator<MutableFst<Arc> > aiter(*fst, s);
        aiter.Seek(aidx);
        Arc arc = aiter.Value();
  
        Label ilabel = arc.ilabel;
        Label dest_state = arc.nextstate;
        if (ilabel != 0) {  // non-eps [nothing to do if eps]...
          assert(ilabel < static_cast<Label>(size));
          size_t len = sequences[ilabel].size();
          if (len <= 1) {
            if (len == 0) arc.ilabel = 0;
            else arc.ilabel = sequences[ilabel][0];
            MutableArcIterator<MutableFst<Arc> > mut_aiter(fst, s);
            mut_aiter.Seek(aidx);
            mut_aiter.SetValue(arc);
          } else {  // len>=2.  Must create new states...
            StateId curstate = -1;  // keep compiler happy: this value never used.
            for (size_t n = 0; n < len; n++) {  // adding/modifying "len" arcs.
              StateId nextstate;
              if (n < len-1) {
                nextstate = fst->AddState();
                assert(nextstate >= num_states_at_start);
              } else nextstate = dest_state;  // going back to original arc's
              // destination.
              if (n == 0) {
                arc.ilabel = sequences[ilabel][0];
                arc.nextstate = nextstate;
                MutableArcIterator<MutableFst<Arc> > mut_aiter(fst, s);
                mut_aiter.Seek(aidx);
                mut_aiter.SetValue(arc);
              } else {
                arc.ilabel = sequences[ilabel][n];
                arc.olabel = 0;
                arc.weight = Weight::One();
                arc.nextstate = nextstate;
                fst->AddArc(curstate, arc);
              }
              curstate = nextstate;
            }
          }
        }
      }
    }
  }
  
  
  template<class Arc, class I>
  class RemoveSomeInputSymbolsMapper {
  public:
    Arc operator ()(const Arc &arc_in) {
      Arc ans = arc_in;
      if (to_remove_set_.count(ans.ilabel) != 0) ans.ilabel = 0;  // remove this symbol
      return ans;
    }
    MapFinalAction FinalAction() { return MAP_NO_SUPERFINAL; }
    MapSymbolsAction InputSymbolsAction() { return MAP_CLEAR_SYMBOLS; }
    MapSymbolsAction OutputSymbolsAction() { return MAP_COPY_SYMBOLS; }
    uint64 Properties(uint64 props) const {
      // remove the following as we don't know now if any of them are true.
      uint64 to_remove = kAcceptor|kNotAcceptor|kIDeterministic|kNonIDeterministic|
          kNoEpsilons|kNoIEpsilons|kILabelSorted|kNotILabelSorted;
      return props & ~to_remove;
    }
    RemoveSomeInputSymbolsMapper(const vector<I> &to_remove):
        to_remove_set_(to_remove) {
      KALDI_ASSERT_IS_INTEGER_TYPE(I);
           assert(to_remove_set_.count(0) == 0);  // makes no sense to remove epsilon.
         }
  private:
    kaldi::ConstIntegerSet<I> to_remove_set_;
  };
  
  
  template<class Arc, class I>
  void CreateFactorFst(const vector<vector<I> > &sequences,
                       MutableFst<Arc> *fst) {
    KALDI_ASSERT_IS_INTEGER_TYPE(I);
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Label Label;
    typedef typename Arc::Weight Weight;
  
    assert(fst != NULL);
    fst->DeleteStates();
    StateId loopstate = fst->AddState();
    assert(loopstate == 0);
    fst->SetStart(0);
    fst->SetFinal(0, Weight::One());
    if (sequences.size() != 0) assert(sequences[0].size() == 0);  // can't replace epsilon...
  
    for (Label olabel = 1; olabel < static_cast<Label>(sequences.size()); olabel++) {
      size_t len = sequences[olabel].size();
      if (len == 0) {
        Arc arc(0, olabel, Weight::One(), loopstate);
        fst->AddArc(loopstate, arc);
      } else {
        StateId curstate = loopstate;
        for (size_t i = 0; i < len; i++) {
          StateId nextstate = (i == len-1 ? loopstate : fst->AddState());
          Arc arc(sequences[olabel][i], (i == 0 ? olabel : 0), Weight::One(), nextstate);
          fst->AddArc(curstate, arc);
          curstate = nextstate;
        }
      }
    }
    fst->SetProperties(kOLabelSorted, kOLabelSorted);
  }
  
  
  template<class Arc, class I>
  void CreateMapFst(const vector<I> &symbol_map,
                    MutableFst<Arc> *fst) {
    KALDI_ASSERT_IS_INTEGER_TYPE(I);
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Label Label;
    typedef typename Arc::Weight Weight;
  
    assert(fst != NULL);
    fst->DeleteStates();
    StateId loopstate = fst->AddState();
    assert(loopstate == 0);
    fst->SetStart(0);
    fst->SetFinal(0, Weight::One());
    assert(symbol_map.empty() || symbol_map[0] == 0);  // FST cannot map epsilon to something else.
    for (Label olabel = 1; olabel < static_cast<Label>(symbol_map.size()); olabel++) {
      Arc arc(symbol_map[olabel], olabel, Weight::One(), loopstate);
      fst->AddArc(loopstate, arc);
    }
  }
  
  
  
  
  } // end namespace fst.
  
  #endif