// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.
//
// Classes to factor weights in an FST.

#ifndef FST_FACTOR_WEIGHT_H_
#define FST_FACTOR_WEIGHT_H_

#include <algorithm>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

#include <fst/log.h>

#include <fst/cache.h>
#include <fst/test-properties.h>


namespace fst {

constexpr uint32 kFactorFinalWeights = 0x00000001;
constexpr uint32 kFactorArcWeights = 0x00000002;

template <class Arc>
struct FactorWeightOptions : CacheOptions {
  using Label = typename Arc::Label;

  float delta;
  uint32 mode;         // Factor arc weights and/or final weights.
  Label final_ilabel;  // Input label of arc when factoring final weights.
  Label final_olabel;  // Output label of arc when factoring final weights.
  bool increment_final_ilabel;  // When factoring final w' results in > 1 arcs
  bool increment_final_olabel;  // at state, increment labels to make distinct?

  explicit FactorWeightOptions(const CacheOptions &opts, float delta = kDelta,
                               uint32 mode = kFactorArcWeights |
                                             kFactorFinalWeights,
                               Label final_ilabel = 0, Label final_olabel = 0,
                               bool increment_final_ilabel = false,
                               bool increment_final_olabel = false)
      : CacheOptions(opts),
        delta(delta),
        mode(mode),
        final_ilabel(final_ilabel),
        final_olabel(final_olabel),
        increment_final_ilabel(increment_final_ilabel),
        increment_final_olabel(increment_final_olabel) {}

  explicit FactorWeightOptions(float delta = kDelta,
                               uint32 mode = kFactorArcWeights |
                                             kFactorFinalWeights,
                               Label final_ilabel = 0, Label final_olabel = 0,
                               bool increment_final_ilabel = false,
                               bool increment_final_olabel = false)
      : delta(delta),
        mode(mode),
        final_ilabel(final_ilabel),
        final_olabel(final_olabel),
        increment_final_ilabel(increment_final_ilabel),
        increment_final_olabel(increment_final_olabel) {}
};

// A factor iterator takes as argument a weight w and returns a sequence of
// pairs of weights (xi, yi) such that the sum of the products xi times yi is
// equal to w. If w is fully factored, the iterator should return nothing.
//
// template <class W>
// class FactorIterator {
//  public:
//   explicit FactorIterator(W w);
//
//   bool Done() const;
//
//   void Next();
//
//   std::pair<W, W> Value() const;
//
//   void Reset();
// }

// Factors trivially.
template <class W>
class IdentityFactor {
 public:
  explicit IdentityFactor(const W &weight) {}

  bool Done() const { return true; }

  void Next() {}

  std::pair<W, W> Value() const { return std::make_pair(W::One(), W::One()); }

  void Reset() {}
};

// Factors a StringWeight w as 'ab' where 'a' is a label.
template <typename Label, StringType S = STRING_LEFT>
class StringFactor {
 public:
  explicit StringFactor(const StringWeight<Label, S> &weight)
      : weight_(weight), done_(weight.Size() <= 1) {}

  bool Done() const { return done_; }

  void Next() { done_ = true; }

  std::pair<StringWeight<Label, S>, StringWeight<Label, S>> Value() const {
    using Weight = StringWeight<Label, S>;
    typename Weight::Iterator siter(weight_);
    Weight w1(siter.Value());
    Weight w2;
    for (siter.Next(); !siter.Done(); siter.Next()) w2.PushBack(siter.Value());
    return std::make_pair(w1, w2);
  }

  void Reset() { done_ = weight_.Size() <= 1; }

 private:
  const StringWeight<Label, S> weight_;
  bool done_;
};

// Factor a GallicWeight using StringFactor.
template <class Label, class W, GallicType G = GALLIC_LEFT>
class GallicFactor {
 public:
  using GW = GallicWeight<Label, W, G>;

  explicit GallicFactor(const GW &weight)
      : weight_(weight), done_(weight.Value1().Size() <= 1) {}

  bool Done() const { return done_; }

  void Next() { done_ = true; }

  std::pair<GW, GW> Value() const {
    StringFactor<Label, GallicStringType(G)> siter(weight_.Value1());
    GW w1(siter.Value().first, weight_.Value2());
    GW w2(siter.Value().second, W::One());
    return std::make_pair(w1, w2);
  }

  void Reset() { done_ = weight_.Value1().Size() <= 1; }

 private:
  const GW weight_;
  bool done_;
};

// Specialization for the (general) GALLIC type GallicWeight.
template <class Label, class W>
class GallicFactor<Label, W, GALLIC> {
 public:
  using GW = GallicWeight<Label, W, GALLIC>;
  using GRW = GallicWeight<Label, W, GALLIC_RESTRICT>;

  explicit GallicFactor(const GW &weight)
      : iter_(weight),
        done_(weight.Size() == 0 ||
              (weight.Size() == 1 && weight.Back().Value1().Size() <= 1)) {}

  bool Done() const { return done_ || iter_.Done(); }

  void Next() { iter_.Next(); }

  void Reset() { iter_.Reset(); }

  std::pair<GW, GW> Value() const {
    const auto weight = iter_.Value();
    StringFactor<Label, GallicStringType(GALLIC_RESTRICT)> siter(
        weight.Value1());
    GRW w1(siter.Value().first, weight.Value2());
    GRW w2(siter.Value().second, W::One());
    return std::make_pair(GW(w1), GW(w2));
  }

 private:
  UnionWeightIterator<GRW, GallicUnionWeightOptions<Label, W>> iter_;
  bool done_;
};

namespace internal {

// Implementation class for FactorWeight
template <class Arc, class FactorIterator>
class FactorWeightFstImpl : public CacheImpl<Arc> {
 public:
  using Label = typename Arc::Label;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  using FstImpl<Arc>::SetType;
  using FstImpl<Arc>::SetProperties;
  using FstImpl<Arc>::SetInputSymbols;
  using FstImpl<Arc>::SetOutputSymbols;

  using CacheBaseImpl<CacheState<Arc>>::PushArc;
  using CacheBaseImpl<CacheState<Arc>>::HasStart;
  using CacheBaseImpl<CacheState<Arc>>::HasFinal;
  using CacheBaseImpl<CacheState<Arc>>::HasArcs;
  using CacheBaseImpl<CacheState<Arc>>::SetArcs;
  using CacheBaseImpl<CacheState<Arc>>::SetFinal;
  using CacheBaseImpl<CacheState<Arc>>::SetStart;

  struct Element {
    Element() {}

    Element(StateId s, Weight weight_) : state(s), weight(std::move(weight_)) {}

    StateId state;  // Input state ID.
    Weight weight;  // Residual weight.
  };

  FactorWeightFstImpl(const Fst<Arc> &fst, const FactorWeightOptions<Arc> &opts)
      : CacheImpl<Arc>(opts),
        fst_(fst.Copy()),
        delta_(opts.delta),
        mode_(opts.mode),
        final_ilabel_(opts.final_ilabel),
        final_olabel_(opts.final_olabel),
        increment_final_ilabel_(opts.increment_final_ilabel),
        increment_final_olabel_(opts.increment_final_olabel) {
    SetType("factor_weight");
    const auto props = fst.Properties(kFstProperties, false);
    SetProperties(FactorWeightProperties(props), kCopyProperties);
    SetInputSymbols(fst.InputSymbols());
    SetOutputSymbols(fst.OutputSymbols());
    if (mode_ == 0) {
      LOG(WARNING) << "FactorWeightFst: Factor mode is set to 0; "
                   << "factoring neither arc weights nor final weights";
    }
  }

  FactorWeightFstImpl(const FactorWeightFstImpl<Arc, FactorIterator> &impl)
      : CacheImpl<Arc>(impl),
        fst_(impl.fst_->Copy(true)),
        delta_(impl.delta_),
        mode_(impl.mode_),
        final_ilabel_(impl.final_ilabel_),
        final_olabel_(impl.final_olabel_),
        increment_final_ilabel_(impl.increment_final_ilabel_),
        increment_final_olabel_(impl.increment_final_olabel_) {
    SetType("factor_weight");
    SetProperties(impl.Properties(), kCopyProperties);
    SetInputSymbols(impl.InputSymbols());
    SetOutputSymbols(impl.OutputSymbols());
  }

  StateId Start() {
    if (!HasStart()) {
      const auto s = fst_->Start();
      if (s == kNoStateId) return kNoStateId;
      SetStart(FindState(Element(fst_->Start(), Weight::One())));
    }
    return CacheImpl<Arc>::Start();
  }

  Weight Final(StateId s) {
    if (!HasFinal(s)) {
      const auto &element = elements_[s];
      // TODO(sorenj): fix so cast is unnecessary
      const auto weight =
          element.state == kNoStateId
              ? element.weight
              : (Weight)Times(element.weight, fst_->Final(element.state));
      FactorIterator siter(weight);
      if (!(mode_ & kFactorFinalWeights) || siter.Done()) {
        SetFinal(s, weight);
      } else {
        SetFinal(s, Weight::Zero());
      }
    }
    return CacheImpl<Arc>::Final(s);
  }

  size_t NumArcs(StateId s) {
    if (!HasArcs(s)) Expand(s);
    return CacheImpl<Arc>::NumArcs(s);
  }

  size_t NumInputEpsilons(StateId s) {
    if (!HasArcs(s)) Expand(s);
    return CacheImpl<Arc>::NumInputEpsilons(s);
  }

  size_t NumOutputEpsilons(StateId s) {
    if (!HasArcs(s)) Expand(s);
    return CacheImpl<Arc>::NumOutputEpsilons(s);
  }

  uint64 Properties() const override { return Properties(kFstProperties); }

  // Sets error if found, and returns other FST impl properties.
  uint64 Properties(uint64 mask) const override {
    if ((mask & kError) && fst_->Properties(kError, false)) {
      SetProperties(kError, kError);
    }
    return FstImpl<Arc>::Properties(mask);
  }

  void InitArcIterator(StateId s, ArcIteratorData<Arc> *data) {
    if (!HasArcs(s)) Expand(s);
    CacheImpl<Arc>::InitArcIterator(s, data);
  }

  // Finds state corresponding to an element, creating new state if element not
  // found.
  StateId FindState(const Element &element) {
    if (!(mode_ & kFactorArcWeights) && element.weight == Weight::One() &&
        element.state != kNoStateId) {
      while (unfactored_.size() <= element.state)
        unfactored_.push_back(kNoStateId);
      if (unfactored_[element.state] == kNoStateId) {
        unfactored_[element.state] = elements_.size();
        elements_.push_back(element);
      }
      return unfactored_[element.state];
    } else {
      const auto insert_result =
          element_map_.insert(std::make_pair(element, elements_.size()));
      if (insert_result.second) {
        elements_.push_back(element);
      }
      return insert_result.first->second;
    }
  }

  // Computes the outgoing transitions from a state, creating new destination
  // states as needed.
  void Expand(StateId s) {
    const auto element = elements_[s];
    if (element.state != kNoStateId) {
      for (ArcIterator<Fst<Arc>> ait(*fst_, element.state); !ait.Done();
           ait.Next()) {
        const auto &arc = ait.Value();
        const auto weight = Times(element.weight, arc.weight);
        FactorIterator fiter(weight);
        if (!(mode_ & kFactorArcWeights) || fiter.Done()) {
          const auto dest = FindState(Element(arc.nextstate, Weight::One()));
          PushArc(s, Arc(arc.ilabel, arc.olabel, weight, dest));
        } else {
          for (; !fiter.Done(); fiter.Next()) {
            const auto &pair = fiter.Value();
            const auto dest =
                FindState(Element(arc.nextstate, pair.second.Quantize(delta_)));
            PushArc(s, Arc(arc.ilabel, arc.olabel, pair.first, dest));
          }
        }
      }
    }
    if ((mode_ & kFactorFinalWeights) &&
        ((element.state == kNoStateId) ||
         (fst_->Final(element.state) != Weight::Zero()))) {
      const auto weight =
          element.state == kNoStateId
              ? element.weight
              : Times(element.weight, fst_->Final(element.state));
      auto ilabel = final_ilabel_;
      auto olabel = final_olabel_;
      for (FactorIterator fiter(weight); !fiter.Done(); fiter.Next()) {
        const auto &pair = fiter.Value();
        const auto dest =
            FindState(Element(kNoStateId, pair.second.Quantize(delta_)));
        PushArc(s, Arc(ilabel, olabel, pair.first, dest));
        if (increment_final_ilabel_) ++ilabel;
        if (increment_final_olabel_) ++olabel;
      }
    }
    SetArcs(s);
  }

 private:
  // Equality function for Elements, assume weights have been quantized.
  class ElementEqual {
   public:
    bool operator()(const Element &x, const Element &y) const {
      return x.state == y.state && x.weight == y.weight;
    }
  };

  // Hash function for Elements to Fst states.
  class ElementKey {
   public:
    size_t operator()(const Element &x) const {
      static constexpr auto prime = 7853;
      return static_cast<size_t>(x.state * prime + x.weight.Hash());
    }
  };

  using ElementMap =
      std::unordered_map<Element, StateId, ElementKey, ElementEqual>;

  std::unique_ptr<const Fst<Arc>> fst_;
  float delta_;
  uint32 mode_;         // Factoring arc and/or final weights.
  Label final_ilabel_;  // ilabel of arc created when factoring final weights.
  Label final_olabel_;  // olabel of arc created when factoring final weights.
  bool increment_final_ilabel_;    // When factoring final weights results in
  bool increment_final_olabel_;    // mutiple arcs, increment labels?
  std::vector<Element> elements_;  // mapping from FST state to Element.
  ElementMap element_map_;         // mapping from Element to FST state.
  // Mapping between old/new StateId for states that do not need to be factored
  // when mode_ is 0 or kFactorFinalWeights.
  std::vector<StateId> unfactored_;
};

}  // namespace internal

// FactorWeightFst takes as template parameter a FactorIterator as defined
// above. The result of weight factoring is a transducer equivalent to the
// input whose path weights have been factored according to the FactorIterator.
// States and transitions will be added as necessary. The algorithm is a
// generalization to arbitrary weights of the second step of the input
// epsilon-normalization algorithm.
//
// This class attaches interface to implementation and handles reference
// counting, delegating most methods to ImplToFst.
template <class A, class FactorIterator>
class FactorWeightFst
    : public ImplToFst<internal::FactorWeightFstImpl<A, FactorIterator>> {
 public:
  using Arc = A;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  using Store = DefaultCacheStore<Arc>;
  using State = typename Store::State;
  using Impl = internal::FactorWeightFstImpl<Arc, FactorIterator>;

  friend class ArcIterator<FactorWeightFst<Arc, FactorIterator>>;
  friend class StateIterator<FactorWeightFst<Arc, FactorIterator>>;

  explicit FactorWeightFst(const Fst<Arc> &fst)
      : ImplToFst<Impl>(
            std::make_shared<Impl>(fst, FactorWeightOptions<Arc>())) {}

  FactorWeightFst(const Fst<Arc> &fst, const FactorWeightOptions<Arc> &opts)
      : ImplToFst<Impl>(std::make_shared<Impl>(fst, opts)) {}

  // See Fst<>::Copy() for doc.
  FactorWeightFst(const FactorWeightFst<Arc, FactorIterator> &fst, bool copy)
      : ImplToFst<Impl>(fst, copy) {}

  // Get a copy of this FactorWeightFst. See Fst<>::Copy() for further doc.
  FactorWeightFst<Arc, FactorIterator> *Copy(bool copy = false) const override {
    return new FactorWeightFst<Arc, FactorIterator>(*this, copy);
  }

  inline void InitStateIterator(StateIteratorData<Arc> *data) const override;

  void InitArcIterator(StateId s, ArcIteratorData<Arc> *data) const override {
    GetMutableImpl()->InitArcIterator(s, data);
  }

 private:
  using ImplToFst<Impl>::GetImpl;
  using ImplToFst<Impl>::GetMutableImpl;

  FactorWeightFst &operator=(const FactorWeightFst &) = delete;
};

// Specialization for FactorWeightFst.
template <class Arc, class FactorIterator>
class StateIterator<FactorWeightFst<Arc, FactorIterator>>
    : public CacheStateIterator<FactorWeightFst<Arc, FactorIterator>> {
 public:
  explicit StateIterator(const FactorWeightFst<Arc, FactorIterator> &fst)
      : CacheStateIterator<FactorWeightFst<Arc, FactorIterator>>(
            fst, fst.GetMutableImpl()) {}
};

// Specialization for FactorWeightFst.
template <class Arc, class FactorIterator>
class ArcIterator<FactorWeightFst<Arc, FactorIterator>>
    : public CacheArcIterator<FactorWeightFst<Arc, FactorIterator>> {
 public:
  using StateId = typename Arc::StateId;

  ArcIterator(const FactorWeightFst<Arc, FactorIterator> &fst, StateId s)
      : CacheArcIterator<FactorWeightFst<Arc, FactorIterator>>(
            fst.GetMutableImpl(), s) {
    if (!fst.GetImpl()->HasArcs(s)) fst.GetMutableImpl()->Expand(s);
  }
};

template <class Arc, class FactorIterator>
inline void FactorWeightFst<Arc, FactorIterator>::InitStateIterator(
    StateIteratorData<Arc> *data) const {
  data->base = new StateIterator<FactorWeightFst<Arc, FactorIterator>>(*this);
}

}  // namespace fst

#endif  // FST_FACTOR_WEIGHT_H_