// fstext/prune-special-inl.h
  
  // Copyright 2014  Johns Hopkins University (Author: Daniel Povey)
  //                 Guoguo Chen
  
  // 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_PRUNE_SPECIAL_INL_H_
  #define KALDI_FSTEXT_PRUNE_SPECIAL_INL_H_
  // Do not include this file directly.  It is included by prune-special.h
  
  #include "fstext/prune-special.h"
  #include "base/kaldi-error.h"
  
  namespace fst {
  
  
  /// This class is used to implement the function PruneSpecial.
  template<class Arc> class PruneSpecialClass {
   public:
    typedef typename Arc::StateId InputStateId;
    typedef typename Arc::StateId OutputStateId;
    typedef typename Arc::Weight Weight;
    typedef typename Arc::Label Label;
    
    PruneSpecialClass(const Fst<Arc> &ifst,
                      VectorFst<Arc> *ofst,
                      Weight beam,
                      size_t max_states):
        ifst_(ifst), ofst_(ofst), beam_(beam), max_states_(max_states),
        best_weight_(Weight::Zero()) {
      KALDI_ASSERT(beam != Weight::One());
      KALDI_ASSERT(queue_.size() == 0);
      ofst_->DeleteStates(); // make sure it's empty.
      if (ifst_.Start() == kNoStateId)
        return;
      ofst_->SetStart(ProcessState(ifst_.Start(), Weight::One()));
  
      while (!queue_.empty()) {
        Task task = queue_.top();
        queue_.pop();
        if (Done(task)) break;
        else ProcessTask(task);
      }
      Connect(ofst);
      if (beam_ != Weight::One())
        Prune(ofst, beam_);
    }
    
    struct Task {
      InputStateId istate;
      OutputStateId ostate; // could be looked up; this is for speed.
      size_t position; // arc position, or -1 if final-prob.
      Weight weight;
      
      Task(InputStateId istate, OutputStateId ostate, size_t position,
           Weight weight): istate(istate), ostate(ostate), position(position),
                           weight(weight) { }
      bool operator < (const Task &other) const {
        return Compare(weight, other.weight) < 0;
      }
    };
  
    bool Done(const Task &task) {
      if (beam_ != Weight::One() && best_weight_ != Weight::Zero() &&
          Compare(task.weight, Times(best_weight_, beam_)) < 0)
        return true;
      if (max_states_ > 0 &&
          static_cast<size_t>(ofst_->NumStates()) > max_states_)
        return true;
      return false;
    }
    
  
    // This function assumes "state" has not been seen before, so we need to
    // create a new output-state for it and add tasks.  It returns the
    // output-state id.  "weight" is the best cost from the start-state to this
    // state.
    inline OutputStateId ProcessState(InputStateId istate, const Weight &weight) {
      OutputStateId ostate = ofst_->AddState();
      state_map_[istate] = ostate;
      for (ArcIterator<Fst<Arc> > aiter(ifst_, istate); !aiter.Done();
           aiter.Next()) {
        const Arc &arc = aiter.Value();
        Task new_task(istate, ostate, aiter.Position(),
                      Times(weight, arc.weight));
        KALDI_ASSERT(Compare(arc.weight, Weight::One()) != 1);
        queue_.push(new_task);
      }
      Weight final = ifst_.Final(istate);
      if (final != Weight::Zero()) {
        Task final_task(istate, ostate, static_cast<size_t>(-1),
                        Times(weight, final));
        KALDI_ASSERT(Compare(final, Weight::One()) != 1);
        queue_.push(final_task);
      }
      return ostate;
    }
  
    // Returns the output-state id corresponding to "istate".  This assumes we are
    // processing a task corresponding to an arc to "istate", and the cost from
    // the start-state to this state is "weight".  Since we process tasks in
    // order, if this is the first time we see this istate, then this is the best
    // cost from the start-state to this state, and it can be used in setting the
    // priority costs in ProcessState().
    inline OutputStateId GetOutputStateId(InputStateId istate,
                                          const Weight &weight) {
      typedef typename unordered_map<InputStateId, OutputStateId>::iterator IterType;
      IterType iter = state_map_.find(istate);
      if (iter == state_map_.end())
        return ProcessState(istate, weight);
      else 
        return iter->second;
    }
    
    void ProcessTask(const Task &task) {
      if (task.position == static_cast<size_t>(-1)) {
        ofst_->SetFinal(task.ostate, ifst_.Final(task.istate));
        if (best_weight_ == Weight::Zero())
          best_weight_ = task.weight; // best-path cost through FST, used for
                                      // beam-pruning.
      } else {
        ArcIterator<Fst<Arc> > aiter(ifst_, task.istate);
        aiter.Seek(task.position); // if we spend most of our time here, we may
                                   // need to store the arc in the Task.
        const Arc &arc = aiter.Value();
        InputStateId next_istate = arc.nextstate;
        OutputStateId next_ostate = GetOutputStateId(next_istate, task.weight);
        Arc oarc(arc.ilabel, arc.olabel, arc.weight, next_ostate);
        ofst_->AddArc(task.ostate, oarc);
      }
    }
    
   private:
    const Fst<Arc> &ifst_;
    VectorFst<Arc> *ofst_;
    Weight beam_;
    size_t max_states_;
  
    unordered_map<InputStateId, OutputStateId> state_map_;
    std::priority_queue<Task> queue_;
    Weight best_weight_; // if not Zero(), then we have now processed a successful path
                         // through ifst_, and this is the weight.
    
  };
  
  template<class Arc>
  void PruneSpecial(const Fst<Arc> &ifst,
                    VectorFst<Arc> *ofst,
                    typename Arc::Weight beam,
                    size_t max_states) {
    PruneSpecialClass<Arc> c(ifst, ofst, beam, max_states);
  }
  
  
  
  }
  
  
  #endif