// fstext/remove-eps-local-inl.h
  
  // Copyright 2009-2011  Microsoft Corporation
  //                2014  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_REMOVE_EPS_LOCAL_INL_H_
  #define KALDI_FSTEXT_REMOVE_EPS_LOCAL_INL_H_
  
  
  namespace fst {
  
  
  template<class Weight>
  struct ReweightPlusDefault {
    inline Weight operator () (const Weight &a, const Weight &b) {
      return Plus(a, b);
    }
  };
  
  struct ReweightPlusLogArc {
    inline TropicalWeight operator () (const TropicalWeight &a,
                                       const TropicalWeight &b) {
      LogWeight a_log(a.Value()), b_log(b.Value());
      return TropicalWeight(Plus(a_log, b_log).Value());
    }
  };
  
  
  
  template<class Arc, class ReweightPlus = ReweightPlusDefault<typename Arc::Weight> >
  class RemoveEpsLocalClass {
    typedef typename Arc::StateId StateId;
    typedef typename Arc::Label Label;
    typedef typename Arc::Weight Weight;
  
   public:
    RemoveEpsLocalClass(MutableFst<Arc> *fst):
        fst_(fst) {
      if (fst_->Start() == kNoStateId) return;  // empty.
      non_coacc_state_ = fst_->AddState();
      InitNumArcs();
      StateId num_states = fst_->NumStates();
      for (StateId s = 0; s < num_states; s++)
        for (size_t pos = 0; pos < fst_->NumArcs(s); pos++)
          RemoveEps(s, pos);
      assert(CheckNumArcs());
      Connect(fst);  // remove inaccessible states.
    }
   private:
    MutableFst<Arc> *fst_;
    StateId non_coacc_state_;  //  use this to delete arcs: make it nextstate
    vector<StateId> num_arcs_in_;   // The number of arcs into the state, plus one
                                    // if it's the start state.
    vector<StateId> num_arcs_out_;  // The number of arcs out of the state, plus
                                    // one if it's a final state.
    ReweightPlus reweight_plus_;
  
    bool CanCombineArcs(const Arc &a, const Arc &b, Arc *c) {
      if (a.ilabel != 0 && b.ilabel != 0) return false;
      if (a.olabel != 0 && b.olabel != 0) return false;
      c->weight = Times(a.weight, b.weight);
      c->ilabel = (a.ilabel != 0 ? a.ilabel : b.ilabel);
      c->olabel = (a.olabel != 0 ? a.olabel : b.olabel);
      c->nextstate = b.nextstate;
      return true;
    }
  
    static bool CanCombineFinal(const Arc &a, Weight final_prob, Weight *final_prob_out) {
      if (a.ilabel != 0 || a.olabel != 0) return false;
      else {
        *final_prob_out = Times(a.weight, final_prob);
        return true;
      }
    }
  
    void InitNumArcs() {  // init num transitions in/out of each state.
      StateId num_states = fst_->NumStates();
      num_arcs_in_.resize(num_states);
      num_arcs_out_.resize(num_states);
      num_arcs_in_[fst_->Start()]++;  // count start as trans in.
      for (StateId s = 0; s < num_states; s++) {
        if (fst_->Final(s) != Weight::Zero())
        num_arcs_out_[s]++;  // count final as transition.
        for (ArcIterator<MutableFst<Arc> > aiter(*fst_, s); !aiter.Done(); aiter.Next()) {
          num_arcs_in_[aiter.Value().nextstate]++;
          num_arcs_out_[s]++;
        }
      }
    }
  
    bool CheckNumArcs() {  // check num arcs in/out of each state, at end.  Debug.
      num_arcs_in_[fst_->Start()]--;  // count start as trans in.
      StateId num_states =   fst_->NumStates();
      for (StateId s = 0; s < num_states; s++) {
        if (s == non_coacc_state_) continue;
        if (fst_->Final(s) != Weight::Zero())
          num_arcs_out_[s]--;  // count final as transition.
        for (ArcIterator<MutableFst<Arc> > aiter(*fst_, s); !aiter.Done(); aiter.Next()) {
          if (aiter.Value().nextstate == non_coacc_state_) continue;
          num_arcs_in_[aiter.Value().nextstate]--;
          num_arcs_out_[s]--;
        }
      }
      for (StateId s = 0; s < num_states; s++) {
        assert(num_arcs_in_[s] == 0);
        assert(num_arcs_out_[s] == 0);
      }
      return true;  // always does this.  so we can assert it w/o warnings.
    }
  
    inline void GetArc(StateId s, size_t pos, Arc *arc) const {
      ArcIterator<MutableFst<Arc> > aiter(*fst_, s);
      aiter.Seek(pos);
      *arc = aiter.Value();
    }
  
    inline void SetArc(StateId s, size_t pos, const Arc &arc) {
      MutableArcIterator<MutableFst<Arc> > aiter(fst_, s);
      aiter.Seek(pos);
      aiter.SetValue(arc);
    }
  
  
    void Reweight(StateId s, size_t pos, Weight reweight) {
      // Reweight is called from RemoveEpsPattern1; it is a step we
      // do to preserve stochasticity.  This function multiplies the
      // arc at (s, pos) by reweight and divides all the arcs [+final-prob]
      // out of the next state by the same.  This is only valid if
      // the next state has only one arc in and is not the start state.
      assert(reweight != Weight::Zero());
      MutableArcIterator<MutableFst<Arc> > aiter(fst_, s);
      aiter.Seek(pos);
      Arc arc = aiter.Value();
      assert(num_arcs_in_[arc.nextstate] == 1);
      arc.weight = Times(arc.weight, reweight);
      aiter.SetValue(arc);
  
      for (MutableArcIterator<MutableFst<Arc> > aiter_next(fst_, arc.nextstate);
           !aiter_next.Done();
           aiter_next.Next()) {
        Arc nextarc = aiter_next.Value();
        if (nextarc.nextstate != non_coacc_state_) {
          nextarc.weight = Divide(nextarc.weight, reweight, DIVIDE_LEFT);
          aiter_next.SetValue(nextarc);
        }
      }
      Weight final = fst_->Final(arc.nextstate);
      if (final != Weight::Zero()) {
        fst_->SetFinal(arc.nextstate, Divide(final, reweight, DIVIDE_LEFT));
      }
    }
  
    // RemoveEpsPattern1 applies where this arc, which is not a
    // self-loop, enters a state which has only one input transition
    // [and is not the start state], and has multiple output
    // transitions [counting being the final-state as a final-transition].
  
    void RemoveEpsPattern1(StateId s, size_t pos, Arc arc) {
      const StateId nextstate = arc.nextstate;
      Weight total_removed = Weight::Zero(),
          total_kept = Weight::Zero();  // totals out of nextstate.
      vector<Arc> arcs_to_add;  // to add to state s.
      for (MutableArcIterator<MutableFst<Arc> > aiter_next(fst_, nextstate);
          !aiter_next.Done();
          aiter_next.Next()) {
        Arc nextarc = aiter_next.Value();
        if (nextarc.nextstate == non_coacc_state_) continue;  // deleted.
        Arc combined;
        if (CanCombineArcs(arc, nextarc, &combined)) {
          total_removed = reweight_plus_(total_removed, nextarc.weight);
          num_arcs_out_[nextstate]--;
          num_arcs_in_[nextarc.nextstate]--;
          nextarc.nextstate = non_coacc_state_;
          aiter_next.SetValue(nextarc);
          arcs_to_add.push_back(combined);
        } else {
          total_kept = reweight_plus_(total_kept, nextarc.weight);
        }
      }
  
      {  // now final-state.
        Weight next_final = fst_->Final(nextstate);
        if (next_final != Weight::Zero()) {
          Weight new_final;
          if (CanCombineFinal(arc, next_final, &new_final)) {
            total_removed = reweight_plus_(total_removed, next_final);
            if (fst_->Final(s) == Weight::Zero())
              num_arcs_out_[s]++;  // final is counted as arc.
            fst_->SetFinal(s, Plus(fst_->Final(s), new_final));
            num_arcs_out_[nextstate]--;
            fst_->SetFinal(nextstate, Weight::Zero());
          } else {
            total_kept = reweight_plus_(total_kept, next_final);
          }
        }
      }
  
      if (total_removed != Weight::Zero()) {  // did something...
        if (total_kept == Weight::Zero()) {  // removed everything: remove arc.
          num_arcs_out_[s]--;
          num_arcs_in_[arc.nextstate]--;
          arc.nextstate = non_coacc_state_;
          SetArc(s, pos, arc);
        } else {
          // Have to reweight.
          Weight total = reweight_plus_(total_removed, total_kept);
          Weight reweight = Divide(total_kept, total, DIVIDE_LEFT);  // <=1
          Reweight(s, pos, reweight);
        }
      }
      // Now add the arcs we were going to add.
      for (size_t i = 0; i < arcs_to_add.size(); i++) {
        num_arcs_out_[s]++;
        num_arcs_in_[arcs_to_add[i].nextstate]++;
        fst_->AddArc(s, arcs_to_add[i]);
      }
    }
  
    void RemoveEpsPattern2(StateId s, size_t pos, Arc arc) {
  
      // Pattern 2 is where "nextstate" has only one arc out, counting
      // being-the-final-state as an arc, but possibly multiple arcs in.
      // Also, nextstate != s.
  
      const StateId nextstate = arc.nextstate;
      bool can_delete_next = (num_arcs_in_[nextstate] == 1);  // if
      // we combine, can delete the corresponding out-arc/final-prob
      // of nextstate.
      bool delete_arc = false;  // set to true if this arc to be deleted.
  
      Weight next_final = fst_->Final(arc.nextstate);
      if (next_final != Weight::Zero()) {  // nextstate has no actual arcs out, only final-prob.
        Weight new_final;
        if (CanCombineFinal(arc, next_final, &new_final)) {
          if (fst_->Final(s) == Weight::Zero())
            num_arcs_out_[s]++;  // final is counted as arc.
          fst_->SetFinal(s, Plus(fst_->Final(s), new_final));
          delete_arc = true;  // will delete "arc".
          if (can_delete_next) {
            num_arcs_out_[nextstate]--;
            fst_->SetFinal(nextstate, Weight::Zero());
          }
        }
      } else {  // has an arc but no final prob.
        MutableArcIterator<MutableFst<Arc> > aiter_next(fst_, nextstate);
        assert(!aiter_next.Done());
        while (aiter_next.Value().nextstate == non_coacc_state_) {
          aiter_next.Next();
          assert(!aiter_next.Done());
        }
        // now aiter_next points to a real arc out of nextstate.
        Arc nextarc = aiter_next.Value();
        Arc combined;
        if (CanCombineArcs(arc, nextarc, &combined)) {
          delete_arc = true;
          if (can_delete_next) {  // do it before we invalidate iterators
            num_arcs_out_[nextstate]--;
            num_arcs_in_[nextarc.nextstate]--;
            nextarc.nextstate = non_coacc_state_;
            aiter_next.SetValue(nextarc);
          }
          num_arcs_out_[s]++;
          num_arcs_in_[combined.nextstate]++;
          fst_->AddArc(s, combined);
        }
      }
      if (delete_arc) {
        num_arcs_out_[s]--;
        num_arcs_in_[nextstate]--;
        arc.nextstate = non_coacc_state_;
        SetArc(s, pos, arc);
      }
    }
  
    void RemoveEps(StateId s, size_t pos) {
      // Tries to do local epsilon-removal for arc sequences starting with this arc
      Arc arc;
      GetArc(s, pos, &arc);
      StateId nextstate = arc.nextstate;
      if (nextstate == non_coacc_state_) return;  // deleted arc.
      if (nextstate == s) return;  // don't handle self-loops: too complex.
  
      if (num_arcs_in_[nextstate] == 1 && num_arcs_out_[nextstate] > 1) {
        RemoveEpsPattern1(s, pos, arc);
      } else if (num_arcs_out_[nextstate] == 1) {
        RemoveEpsPattern2(s, pos, arc);
      }
    }
  
  };
  
  
  template<class Arc>
  void RemoveEpsLocal(MutableFst<Arc> *fst) {
    RemoveEpsLocalClass<Arc> c(fst);  // work gets done in initializer.
  }
  
  
  void RemoveEpsLocalSpecial(MutableFst<StdArc> *fst) {
    // work gets done in initializer.
    RemoveEpsLocalClass<StdArc, ReweightPlusLogArc> c(fst);
  }
  
  } // end namespace fst.
  
  #endif