// See www.openfst.org for extensive documentation on this weighted
  // finite-state transducer library.
  //
  // Functions implementing pruning.
  
  #ifndef FST_PRUNE_H_
  #define FST_PRUNE_H_
  
  #include <type_traits>
  #include <utility>
  #include <vector>
  
  #include <fst/log.h>
  
  #include <fst/arcfilter.h>
  #include <fst/heap.h>
  #include <fst/shortest-distance.h>
  
  
  namespace fst {
  namespace internal {
  
  template <class StateId, class Weight>
  class PruneCompare {
   public:
    PruneCompare(const std::vector<Weight> &idistance,
                 const std::vector<Weight> &fdistance)
        : idistance_(idistance), fdistance_(fdistance) {}
  
    bool operator()(const StateId x, const StateId y) const {
      const auto wx = Times(IDistance(x), FDistance(x));
      const auto wy = Times(IDistance(y), FDistance(y));
      return less_(wx, wy);
    }
  
   private:
    Weight IDistance(const StateId s) const {
      return s < idistance_.size() ? idistance_[s] : Weight::Zero();
    }
  
    Weight FDistance(const StateId s) const {
      return s < fdistance_.size() ? fdistance_[s] : Weight::Zero();
    }
  
    const std::vector<Weight> &idistance_;
    const std::vector<Weight> &fdistance_;
    NaturalLess<Weight> less_;
  };
  
  }  // namespace internal
  
  template <class Arc, class ArcFilter>
  struct PruneOptions {
    using StateId = typename Arc::StateId;
    using Weight = typename Arc::Weight;
  
    PruneOptions(const Weight &weight_threshold, StateId state_threshold,
                 ArcFilter filter, std::vector<Weight> *distance = nullptr,
                 float delta = kDelta, bool threshold_initial = false)
        : weight_threshold(std::move(weight_threshold)),
          state_threshold(state_threshold),
          filter(std::move(filter)),
          distance(distance),
          delta(delta),
          threshold_initial(threshold_initial) {}
  
    // Pruning weight threshold.
    Weight weight_threshold;
    // Pruning state threshold.
    StateId state_threshold;
    // Arc filter.
    ArcFilter filter;
    // If non-zero, passes in pre-computed shortest distance to final states.
    const std::vector<Weight> *distance;
    // Determines the degree of convergence required when computing shortest
    // distances.
    float delta;
    // Determines if the shortest path weight is left (true) or right
    // (false) multiplied by the threshold to get the limit for
    // keeping a state or arc (matters if the semiring is not
    // commutative).
    bool threshold_initial;
  };
  
  // Pruning algorithm: this version modifies its input and it takes an options
  // class as an argument. After pruning the FST contains states and arcs that
  // belong to a successful path in the FST whose weight is no more than the
  // weight of the shortest path Times() the provided weight threshold. When the
  // state threshold is not kNoStateId, the output FST is further restricted to
  // have no more than the number of states in opts.state_threshold. Weights must
  // have the path property. The weight of any cycle needs to be bounded; i.e.,
  //
  //   Plus(weight, Weight::One()) == Weight::One()
  template <class Arc, class ArcFilter,
            typename std::enable_if<
                (Arc::Weight::Properties() & kPath) == kPath>::type * = nullptr>
  void Prune(MutableFst<Arc> *fst, const PruneOptions<Arc, ArcFilter> &opts) {
    using StateId = typename Arc::StateId;
    using Weight = typename Arc::Weight;
    using StateHeap = Heap<StateId, internal::PruneCompare<StateId, Weight>>;
    auto ns = fst->NumStates();
    if (ns < 1) return;
    std::vector<Weight> idistance(ns, Weight::Zero());
    std::vector<Weight> tmp;
    if (!opts.distance) {
      tmp.reserve(ns);
      ShortestDistance(*fst, &tmp, true, opts.delta);
    }
    const auto *fdistance = opts.distance ? opts.distance : &tmp;
    if ((opts.state_threshold == 0) || (fdistance->size() <= fst->Start()) ||
        ((*fdistance)[fst->Start()] == Weight::Zero())) {
      fst->DeleteStates();
      return;
    }
    internal::PruneCompare<StateId, Weight> compare(idistance, *fdistance);
    StateHeap heap(compare);
    std::vector<bool> visited(ns, false);
    std::vector<size_t> enqueued(ns, StateHeap::kNoKey);
    std::vector<StateId> dead;
    dead.push_back(fst->AddState());
    NaturalLess<Weight> less;
    auto s = fst->Start();
    const auto limit = opts.threshold_initial ?
        Times(opts.weight_threshold, (*fdistance)[s]) :
        Times((*fdistance)[s], opts.weight_threshold);
    StateId num_visited = 0;
  
    if (!less(limit, (*fdistance)[s])) {
      idistance[s] = Weight::One();
      enqueued[s] = heap.Insert(s);
      ++num_visited;
    }
    while (!heap.Empty()) {
      s = heap.Top();
      heap.Pop();
      enqueued[s] = StateHeap::kNoKey;
      visited[s] = true;
      if (less(limit, Times(idistance[s], fst->Final(s)))) {
        fst->SetFinal(s, Weight::Zero());
      }
      for (MutableArcIterator<MutableFst<Arc>> aiter(fst, s); !aiter.Done();
           aiter.Next()) {
        auto arc = aiter.Value();  // Copy intended.
        if (!opts.filter(arc)) continue;
        const auto weight = Times(Times(idistance[s], arc.weight),
                                  arc.nextstate < fdistance->size() ?
                                  (*fdistance)[arc.nextstate] : Weight::Zero());
        if (less(limit, weight)) {
          arc.nextstate = dead[0];
          aiter.SetValue(arc);
          continue;
        }
        if (less(Times(idistance[s], arc.weight), idistance[arc.nextstate])) {
          idistance[arc.nextstate] = Times(idistance[s], arc.weight);
        }
        if (visited[arc.nextstate]) continue;
        if ((opts.state_threshold != kNoStateId) &&
            (num_visited >= opts.state_threshold)) {
          continue;
        }
        if (enqueued[arc.nextstate] == StateHeap::kNoKey) {
          enqueued[arc.nextstate] = heap.Insert(arc.nextstate);
          ++num_visited;
        } else {
          heap.Update(enqueued[arc.nextstate], arc.nextstate);
        }
      }
    }
    for (StateId i = 0; i < visited.size(); ++i) {
      if (!visited[i]) dead.push_back(i);
    }
    fst->DeleteStates(dead);
  }
  
  template <class Arc, class ArcFilter,
            typename std::enable_if<
                (Arc::Weight::Properties() & kPath) != kPath>::type * = nullptr>
  void Prune(MutableFst<Arc> *fst, const PruneOptions<Arc, ArcFilter> &) {
    FSTERROR() << "Prune: Weight needs to have the path property: "
               << Arc::Weight::Type();
    fst->SetProperties(kError, kError);
  }
  
  // Pruning algorithm: this version modifies its input and takes the
  // pruning threshold as an argument. It deletes states and arcs in the
  // FST that do not belong to a successful path whose weight is more
  // than the weight of the shortest path Times() the provided weight
  // threshold. When the state threshold is not kNoStateId, the output
  // FST is further restricted to have no more than the number of states
  // in opts.state_threshold. Weights must have the path property. The
  // weight of any cycle needs to be bounded; i.e.,
  //
  //   Plus(weight, Weight::One()) == Weight::One()
  template <class Arc>
  void Prune(MutableFst<Arc> *fst, typename Arc::Weight weight_threshold,
             typename Arc::StateId state_threshold = kNoStateId,
             float delta = kDelta) {
    const PruneOptions<Arc, AnyArcFilter<Arc>> opts(
        weight_threshold, state_threshold, AnyArcFilter<Arc>(), nullptr, delta);
    Prune(fst, opts);
  }
  
  // Pruning algorithm: this version writes the pruned input FST to an
  // output MutableFst and it takes an options class as an argument. The
  // output FST contains states and arcs that belong to a successful
  // path in the input FST whose weight is more than the weight of the
  // shortest path Times() the provided weight threshold. When the state
  // threshold is not kNoStateId, the output FST is further restricted
  // to have no more than the number of states in
  // opts.state_threshold. Weights have the path property.  The weight
  // of any cycle needs to be bounded; i.e.,
  //
  //   Plus(weight, Weight::One()) == Weight::One()
  template <class Arc, class ArcFilter,
            typename std::enable_if<IsPath<typename Arc::Weight>::value>::type * =
                nullptr>
  void Prune(const Fst<Arc> &ifst, MutableFst<Arc> *ofst,
             const PruneOptions<Arc, ArcFilter> &opts) {
    using StateId = typename Arc::StateId;
    using Weight = typename Arc::Weight;
    using StateHeap = Heap<StateId, internal::PruneCompare<StateId, Weight>>;
    ofst->DeleteStates();
    ofst->SetInputSymbols(ifst.InputSymbols());
    ofst->SetOutputSymbols(ifst.OutputSymbols());
    if (ifst.Start() == kNoStateId) return;
    NaturalLess<Weight> less;
    if (less(opts.weight_threshold, Weight::One()) ||
        (opts.state_threshold == 0)) {
      return;
    }
    std::vector<Weight> idistance;
    std::vector<Weight> tmp;
    if (!opts.distance) ShortestDistance(ifst, &tmp, true, opts.delta);
    const auto *fdistance = opts.distance ? opts.distance : &tmp;
    if ((fdistance->size() <= ifst.Start()) ||
        ((*fdistance)[ifst.Start()] == Weight::Zero())) {
      return;
    }
    internal::PruneCompare<StateId, Weight> compare(idistance, *fdistance);
    StateHeap heap(compare);
    std::vector<StateId> copy;
    std::vector<size_t> enqueued;
    std::vector<bool> visited;
    auto s = ifst.Start();
    const auto limit = opts.threshold_initial ?
        Times(opts.weight_threshold, (*fdistance)[s]) :
        Times((*fdistance)[s], opts.weight_threshold);
    while (copy.size() <= s) copy.push_back(kNoStateId);
    copy[s] = ofst->AddState();
    ofst->SetStart(copy[s]);
    while (idistance.size() <= s) idistance.push_back(Weight::Zero());
    idistance[s] = Weight::One();
    while (enqueued.size() <= s) {
      enqueued.push_back(StateHeap::kNoKey);
      visited.push_back(false);
    }
    enqueued[s] = heap.Insert(s);
    while (!heap.Empty()) {
      s = heap.Top();
      heap.Pop();
      enqueued[s] = StateHeap::kNoKey;
      visited[s] = true;
      if (!less(limit, Times(idistance[s], ifst.Final(s)))) {
        ofst->SetFinal(copy[s], ifst.Final(s));
      }
      for (ArcIterator<Fst<Arc>> aiter(ifst, s); !aiter.Done(); aiter.Next()) {
        const auto &arc = aiter.Value();
        if (!opts.filter(arc)) continue;
        const auto weight = Times(Times(idistance[s], arc.weight),
                                  arc.nextstate < fdistance->size() ?
                                  (*fdistance)[arc.nextstate] : Weight::Zero());
        if (less(limit, weight)) continue;
        if ((opts.state_threshold != kNoStateId) &&
            (ofst->NumStates() >= opts.state_threshold)) {
          continue;
        }
        while (idistance.size() <= arc.nextstate) {
          idistance.push_back(Weight::Zero());
        }
        if (less(Times(idistance[s], arc.weight), idistance[arc.nextstate])) {
          idistance[arc.nextstate] = Times(idistance[s], arc.weight);
        }
        while (copy.size() <= arc.nextstate) copy.push_back(kNoStateId);
        if (copy[arc.nextstate] == kNoStateId) {
          copy[arc.nextstate] = ofst->AddState();
        }
        ofst->AddArc(copy[s], Arc(arc.ilabel, arc.olabel, arc.weight,
                                  copy[arc.nextstate]));
        while (enqueued.size() <= arc.nextstate) {
          enqueued.push_back(StateHeap::kNoKey);
          visited.push_back(false);
        }
        if (visited[arc.nextstate]) continue;
        if (enqueued[arc.nextstate] == StateHeap::kNoKey) {
          enqueued[arc.nextstate] = heap.Insert(arc.nextstate);
        } else {
          heap.Update(enqueued[arc.nextstate], arc.nextstate);
        }
      }
    }
  }
  
  template <class Arc, class ArcFilter,
            typename std::enable_if<!IsPath<typename Arc::Weight>::value>::type
                * = nullptr>
  void Prune(const Fst<Arc> &, MutableFst<Arc> *ofst,
             const PruneOptions<Arc, ArcFilter> &) {
    FSTERROR() << "Prune: Weight needs to have the path property: "
               << Arc::Weight::Type();
    ofst->SetProperties(kError, kError);
  }
  
  // Pruning algorithm: this version writes the pruned input FST to an
  // output MutableFst and simply takes the pruning threshold as an
  // argument. The output FST contains states and arcs that belong to a
  // successful path in the input FST whose weight is no more than the
  // weight of the shortest path Times() the provided weight
  // threshold. When the state threshold is not kNoStateId, the output
  // FST is further restricted to have no more than the number of states
  // in opts.state_threshold. Weights must have the path property. The
  // weight of any cycle needs to be bounded; i.e.,
  //
  // Plus(weight, Weight::One()) = Weight::One();
  template <class Arc>
  void Prune(const Fst<Arc> &ifst, MutableFst<Arc> *ofst,
             typename Arc::Weight weight_threshold,
             typename Arc::StateId state_threshold = kNoStateId,
             float delta = kDelta) {
    const PruneOptions<Arc, AnyArcFilter<Arc>> opts(
        weight_threshold, state_threshold, AnyArcFilter<Arc>(), nullptr, delta);
    Prune(ifst, ofst, opts);
  }
  
  }  // namespace fst
  
  #endif  // FST_PRUNE_H_