// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.
//
// Functions and classes to disambiguate an FST.

#ifndef FST_DISAMBIGUATE_H_
#define FST_DISAMBIGUATE_H_

#include <list>
#include <map>
#include <set>
#include <utility>
#include <vector>

#include <fst/arcsort.h>
#include <fst/compose.h>
#include <fst/connect.h>
#include <fst/determinize.h>
#include <fst/dfs-visit.h>
#include <fst/project.h>
#include <fst/prune.h>
#include <fst/state-map.h>
#include <fst/state-table.h>
#include <fst/union-find.h>
#include <fst/verify.h>


namespace fst {

template <class Arc>
struct DisambiguateOptions : public DeterminizeOptions<Arc> {
  using Label = typename Arc::Label;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  explicit DisambiguateOptions(float delta = kDelta,
                               Weight weight = Weight::Zero(),
                               StateId n = kNoStateId, Label label = 0)
      : DeterminizeOptions<Arc>(delta, std::move(weight), n, label,
                                DETERMINIZE_FUNCTIONAL) {}
};

namespace internal {

// A determinization filter based on a subset element relation. The relation is
// assumed to be reflexive and symmetric.
template <class Arc, class Relation>
class RelationDeterminizeFilter {
 public:
  using Label = typename Arc::Label;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  using FilterState = IntegerFilterState<StateId>;
  using StateTuple = DeterminizeStateTuple<Arc, FilterState>;
  using Subset = typename StateTuple::Subset;
  using Element = typename StateTuple::Element;
  using LabelMap = std::multimap<Label, DeterminizeArc<StateTuple>>;

  // This is needed (e.g.) to go into the gallic domain for transducers; there
  // is no need to rebind the relation since its use here only depends on the
  // state IDs.
  template <class A>
  struct rebind {
    using Other = RelationDeterminizeFilter<A, Relation>;
  };

  explicit RelationDeterminizeFilter(const Fst<Arc> &fst)
      : fst_(fst.Copy()), r_(new Relation()), s_(kNoStateId), head_(nullptr) {}

  // Ownership of the relation is given to this class.
  RelationDeterminizeFilter(const Fst<Arc> &fst, Relation *r)
      : fst_(fst.Copy()), r_(r), s_(kNoStateId), head_(0) {}

  // Ownership of the relation is given to this class.
  RelationDeterminizeFilter(const Fst<Arc> &fst, Relation *r,
                            std::vector<StateId> *head)
      : fst_(fst.Copy()), r_(r), s_(kNoStateId), head_(head) {}

  // This is needed, e.g., to go into the gallic domain for transducers.
  // Ownership of the templated filter argument is given to this class.
  template <class Filter>
  RelationDeterminizeFilter(const Fst<Arc> &fst, Filter *filter)
      : fst_(fst.Copy()),
        r_(new Relation(filter->GetRelation())),
        s_(kNoStateId),
        head_(filter->GetHeadStates()) {
    delete filter;
  }

  // Copy constructor; the FST can be passed if it has been deep-copied.
  RelationDeterminizeFilter(const RelationDeterminizeFilter &filter,
                            const Fst<Arc> *fst = nullptr)
      : fst_(fst ? fst->Copy() : filter.fst_->Copy()),
        r_(new Relation(*filter.r_)),
        s_(kNoStateId),
        head_() {}

  FilterState Start() const { return FilterState(fst_->Start()); }

  void SetState(StateId s, const StateTuple &tuple) {
    if (s_ != s) {
      s_ = s;
      tuple_ = &tuple;
      const auto head = tuple.filter_state.GetState();
      is_final_ = fst_->Final(head) != Weight::Zero();
      if (head_) {
        if (head_->size() <= s) head_->resize(s + 1, kNoStateId);
        (*head_)[s] = head;
      }
    }
  }

  // Filters transition, possibly modifying label map. Returns true if arc is
  // added to label map.
  bool FilterArc(const Arc &arc, const Element &src_element,
                 const Element &dest_element, LabelMap *label_map) const;

  // Filters super-final transition, returning new final weight.
  Weight FilterFinal(const Weight final_weight, const Element &element) const {
    return is_final_ ? final_weight : Weight::Zero();
  }

  static uint64 Properties(uint64 props) {
    return props & ~(kIDeterministic | kODeterministic);
  }

  const Relation &GetRelation() { return *r_; }

  std::vector<StateId> *GetHeadStates() { return head_; }

 private:
  // Pairs arc labels with state tuples with possible heads and empty subsets.
  void InitLabelMap(LabelMap *label_map) const;

  std::unique_ptr<Fst<Arc>> fst_;  // Input FST.
  std::unique_ptr<Relation> r_;    // Relation compatible with inv. trans. fnc.
  StateId s_;                      // Current state.
  const StateTuple *tuple_;        // Current tuple.
  bool is_final_;                  // Is the current head state final?
  std::vector<StateId> *head_;     // Head state for a given state,
                                   // owned by the Disambiguator.
};

template <class Arc, class Relation>
bool RelationDeterminizeFilter<Arc, Relation>::FilterArc(
    const Arc &arc, const Element &src_element, const Element &dest_element,
    LabelMap *label_map) const {
  bool added = false;
  if (label_map->empty()) InitLabelMap(label_map);
  // Adds element to state tuple if element state is related to tuple head.
  for (auto liter = label_map->lower_bound(arc.ilabel);
       liter != label_map->end() && liter->first == arc.ilabel; ++liter) {
    auto *dest_tuple = liter->second.dest_tuple;
    const auto dest_head = dest_tuple->filter_state.GetState();
    if ((*r_)(dest_element.state_id, dest_head)) {
      dest_tuple->subset.push_front(dest_element);
      added = true;
    }
  }
  return added;
}

template <class Arc, class Relation>
void RelationDeterminizeFilter<Arc, Relation>::InitLabelMap(
    LabelMap *label_map) const {
  const auto src_head = tuple_->filter_state.GetState();
  Label label = kNoLabel;
  StateId nextstate = kNoStateId;
  for (ArcIterator<Fst<Arc>> aiter(*fst_, src_head); !aiter.Done();
       aiter.Next()) {
    const auto &arc = aiter.Value();
    // Continues if multiarc.
    if (arc.ilabel == label && arc.nextstate == nextstate) continue;
    DeterminizeArc<StateTuple> det_arc(arc);
    det_arc.dest_tuple->filter_state = FilterState(arc.nextstate);
    label_map->insert(std::make_pair(arc.ilabel, det_arc));
    label = arc.ilabel;
    nextstate = arc.nextstate;
  }
}

// Helper class to disambiguate an FST via Disambiguate().
template <class Arc>
class Disambiguator {
 public:
  using Label = typename Arc::Label;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  // IDs arcs with state ID and arc position. Arc position -1 indicates final
  // (super-final transition).
  using ArcId = std::pair<StateId, ssize_t>;

  Disambiguator() : error_(false) {}

  void Disambiguate(
      const Fst<Arc> &ifst, MutableFst<Arc> *ofst,
      const DisambiguateOptions<Arc> &opts = DisambiguateOptions<Arc>()) {
    VectorFst<Arc> sfst(ifst);
    Connect(&sfst);
    ArcSort(&sfst, ArcCompare());
    PreDisambiguate(sfst, ofst, opts);
    ArcSort(ofst, ArcCompare());
    FindAmbiguities(*ofst);
    RemoveSplits(ofst);
    MarkAmbiguities();
    RemoveAmbiguities(ofst);
    if (error_) ofst->SetProperties(kError, kError);
  }

 private:
  // Comparison functor for comparing input labels and next states of arcs. This
  // sort order facilitates the predisambiguation.
  class ArcCompare {
   public:
    bool operator()(const Arc &arc1, const Arc &arc2) const {
      return arc1.ilabel < arc2.ilabel ||
             (arc1.ilabel == arc2.ilabel && arc1.nextstate < arc2.nextstate);
    }

    uint64 Properties(uint64 props) const {
      return (props & kArcSortProperties) | kILabelSorted |
             (props & kAcceptor ? kOLabelSorted : 0);
    }
  };

  // Comparison functor for comparing transitions represented by their arc ID.
  // This sort order facilitates ambiguity detection.
  class ArcIdCompare {
   public:
    explicit ArcIdCompare(const std::vector<StateId> &head) : head_(head) {}

    bool operator()(const ArcId &a1, const ArcId &a2) const {
      // Sort first by source head state...
      const auto src1 = a1.first;
      const auto src2 = a2.first;
      const auto head1 = head_[src1];
      const auto head2 = head_[src2];
      if (head1 < head2) return true;
      if (head2 < head1) return false;
      // ...then by source state...
      if (src1 < src2) return true;
      if (src2 < src1) return false;
      // ...then by position.
      return a1.second < a2.second;
    }

   private:
    const std::vector<StateId> &head_;
  };

  // A relation that determines if two states share a common future.
  class CommonFuture {
   public:
    using StateTable = GenericComposeStateTable<Arc, TrivialFilterState>;
    using StateTuple = typename StateTable::StateTuple;

    // Needed for compilation with DeterminizeRelationFilter.
    CommonFuture() {
      FSTERROR() << "Disambiguate::CommonFuture: FST not provided";
    }

    explicit CommonFuture(const Fst<Arc> &ifst) {
      using M = Matcher<Fst<Arc>>;
      ComposeFstOptions<Arc, M, NullComposeFilter<M>> opts;
      // Ensures composition is between acceptors.
      const bool trans = ifst.Properties(kNotAcceptor, true);
      const auto *fsa =
          trans ? new ProjectFst<Arc>(ifst, PROJECT_INPUT) : &ifst;
      opts.state_table = new StateTable(*fsa, *fsa);
      const ComposeFst<Arc> cfst(*fsa, *fsa, opts);
      std::vector<bool> coaccess;
      uint64 props = 0;
      SccVisitor<Arc> scc_visitor(nullptr, nullptr, &coaccess, &props);
      DfsVisit(cfst, &scc_visitor);
      for (StateId s = 0; s < coaccess.size(); ++s) {
        if (coaccess[s]) {
          related_.insert(opts.state_table->Tuple(s).StatePair());
        }
      }
      if (trans) delete fsa;
    }

    bool operator()(const StateId s1, StateId s2) const {
      return related_.count(std::make_pair(s1, s2)) > 0;
    }

   private:
    // States s1 and s2 resp. are in this relation iff they there is a
    // path from s1 to a final state that has the same label as some
    // path from s2 to a final state.
    std::set<std::pair<StateId, StateId>> related_;
  };

  using ArcIdMap = std::multimap<ArcId, ArcId, ArcIdCompare>;

  // Inserts candidate into the arc ID map.
  inline void InsertCandidate(StateId s1, StateId s2, const ArcId &a1,
                              const ArcId &a2) {
    candidates_->insert(head_[s1] > head_[s2] ? std::make_pair(a1, a2)
                                              : std::make_pair(a2, a1));
  }

  // Returns the arc corresponding to ArcId a.
  static Arc GetArc(const Fst<Arc> &fst, ArcId aid) {
    if (aid.second == -1) {  // Returns super-final transition.
      return Arc(kNoLabel, kNoLabel, fst.Final(aid.first), kNoStateId);
    } else {
      ArcIterator<Fst<Arc>> aiter(fst, aid.first);
      aiter.Seek(aid.second);
      return aiter.Value();
    }
  }

  // Outputs an equivalent FST whose states are subsets of states that have a
  // future path in common.
  void PreDisambiguate(const ExpandedFst<Arc> &ifst, MutableFst<Arc> *ofst,
                       const DisambiguateOptions<Arc> &opts);

  // Finds transitions that are ambiguous candidates in the result of
  // PreDisambiguate.
  void FindAmbiguities(const ExpandedFst<Arc> &fst);

  // Finds transition pairs that are ambiguous candidates from two specified
  // source states.
  void FindAmbiguousPairs(const ExpandedFst<Arc> &fst, StateId s1, StateId s2);

  // Marks ambiguous transitions to be removed.
  void MarkAmbiguities();

  // Deletes spurious ambiguous transitions (due to quantization).
  void RemoveSplits(MutableFst<Arc> *ofst);

  // Deletes actual ambiguous transitions.
  void RemoveAmbiguities(MutableFst<Arc> *ofst);

  // States s1 and s2 are in this relation iff there is a path from the initial
  // state to s1 that has the same label as some path from the initial state to
  // s2. We store only state pairs s1, s2 such that s1 <= s2.
  std::set<std::pair<StateId, StateId>> coreachable_;

  // Queue of disambiguation-related states to be processed. We store only
  // state pairs s1, s2 such that s1 <= s2.
  std::list<std::pair<StateId, StateId>> queue_;

  // Head state in the pre-disambiguation for a given state.
  std::vector<StateId> head_;

  // Maps from a candidate ambiguous arc A to each ambiguous candidate arc B
  // with the same label and destination state as A, whose source state s' is
  // coreachable with the source state s of A, and for which head(s') < head(s).
  std::unique_ptr<ArcIdMap> candidates_;

  // Set of ambiguous transitions to be removed.
  std::set<ArcId> ambiguous_;

  // States to merge due to quantization issues.
  std::unique_ptr<UnionFind<StateId>> merge_;
  // Marks error condition.
  bool error_;

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

template <class Arc>
void Disambiguator<Arc>::PreDisambiguate(const ExpandedFst<Arc> &ifst,
                                         MutableFst<Arc> *ofst,
                                         const DisambiguateOptions<Arc> &opts) {
  using CommonDivisor = DefaultCommonDivisor<Weight>;
  using Filter = RelationDeterminizeFilter<Arc, CommonFuture>;
  // Subset elements with states s1 and s2 (resp.) are in this relation iff they
  // there is a path from s1 to a final state that has the same label as some
  // path from s2 to a final state.
  auto *common_future = new CommonFuture(ifst);
  DeterminizeFstOptions<Arc, CommonDivisor, Filter> nopts;
  nopts.delta = opts.delta;
  nopts.subsequential_label = opts.subsequential_label;
  nopts.filter = new Filter(ifst, common_future, &head_);
  // The filter takes ownership of 'common_future', and determinization takes
  // ownership of the filter itself.
  nopts.gc_limit = 0;  // Cache only the last state for fastest copy.
  if (opts.weight_threshold != Weight::Zero() ||
      opts.state_threshold != kNoStateId) {
    /* TODO(riley): fails regression test; understand why
    if (ifst.Properties(kAcceptor, true)) {
      std::vector<Weight> idistance, odistance;
      ShortestDistance(ifst, &idistance, true);
      DeterminizeFst<Arc> dfst(ifst, &idistance, &odistance, nopts);
      PruneOptions< Arc, AnyArcFilter<Arc>> popts(opts.weight_threshold,
                                                   opts.state_threshold,
                                                   AnyArcFilter<Arc>(),
                                                   &odistance);
      Prune(dfst, ofst, popts);
      } else */ {
      *ofst = DeterminizeFst<Arc>(ifst, nopts);
      Prune(ofst, opts.weight_threshold, opts.state_threshold);
    }
  } else {
    *ofst = DeterminizeFst<Arc>(ifst, nopts);
  }
  head_.resize(ofst->NumStates(), kNoStateId);
}

template <class Arc>
void Disambiguator<Arc>::FindAmbiguities(const ExpandedFst<Arc> &fst) {
  if (fst.Start() == kNoStateId) return;
  candidates_.reset(new ArcIdMap(ArcIdCompare(head_)));
  const auto start_pr = std::make_pair(fst.Start(), fst.Start());
  coreachable_.insert(start_pr);
  queue_.push_back(start_pr);
  while (!queue_.empty()) {
    const auto &pr = queue_.front();
    const auto s1 = pr.first;
    const auto s2 = pr.second;
    queue_.pop_front();
    FindAmbiguousPairs(fst, s1, s2);
  }
}

template <class Arc>
void Disambiguator<Arc>::FindAmbiguousPairs(const ExpandedFst<Arc> &fst,
                                            StateId s1, StateId s2) {
  if (fst.NumArcs(s2) > fst.NumArcs(s1)) FindAmbiguousPairs(fst, s2, s1);
  SortedMatcher<Fst<Arc>> matcher(fst, MATCH_INPUT);
  matcher.SetState(s2);
  for (ArcIterator<Fst<Arc>> aiter(fst, s1); !aiter.Done(); aiter.Next()) {
    const auto &arc1 = aiter.Value();
    const ArcId a1(s1, aiter.Position());
    if (matcher.Find(arc1.ilabel)) {
      for (; !matcher.Done(); matcher.Next()) {
        const auto &arc2 = matcher.Value();
        // Continues on implicit epsilon match.
        if (arc2.ilabel == kNoLabel) continue;
        const ArcId a2(s2, matcher.Position());
        // Actual transition is ambiguous.
        if (s1 != s2 && arc1.nextstate == arc2.nextstate) {
          InsertCandidate(s1, s2, a1, a2);
        }
        const auto spr = arc1.nextstate <= arc2.nextstate
                             ? std::make_pair(arc1.nextstate, arc2.nextstate)
                             : std::make_pair(arc2.nextstate, arc1.nextstate);
        // Not already marked as coreachable?
        if (coreachable_.insert(spr).second) {
          // Only possible if state split by quantization issues.
          if (spr.first != spr.second &&
              head_[spr.first] == head_[spr.second]) {
            if (!merge_) {
              merge_.reset(new UnionFind<StateId>(fst.NumStates(), kNoStateId));
              merge_->MakeAllSet(fst.NumStates());
            }
            merge_->Union(spr.first, spr.second);
          } else {
            queue_.push_back(spr);
          }
        }
      }
    }
  }
  // Super-final transition is ambiguous.
  if (s1 != s2 && fst.Final(s1) != Weight::Zero() &&
      fst.Final(s2) != Weight::Zero()) {
    const ArcId a1(s1, -1);
    const ArcId a2(s2, -1);
    InsertCandidate(s1, s2, a1, a2);
  }
}

template <class Arc>
void Disambiguator<Arc>::MarkAmbiguities() {
  if (!candidates_) return;
  for (auto it = candidates_->begin(); it != candidates_->end(); ++it) {
    const auto a = it->first;
    const auto b = it->second;
    // If b is not to be removed, then a is.
    if (ambiguous_.count(b) == 0) ambiguous_.insert(a);
  }
  coreachable_.clear();
  candidates_.reset();
}

template <class Arc>
void Disambiguator<Arc>::RemoveSplits(MutableFst<Arc> *ofst) {
  if (!merge_) return;
  // Merges split states to remove spurious ambiguities.
  for (StateIterator<MutableFst<Arc>> siter(*ofst); !siter.Done();
       siter.Next()) {
    for (MutableArcIterator<MutableFst<Arc>> aiter(ofst, siter.Value());
         !aiter.Done(); aiter.Next()) {
      auto arc = aiter.Value();
      const auto nextstate = merge_->FindSet(arc.nextstate);
      if (nextstate != arc.nextstate) {
        arc.nextstate = nextstate;
        aiter.SetValue(arc);
      }
    }
  }
  // Repeats search for actual ambiguities on modified FST.
  coreachable_.clear();
  merge_.reset();
  candidates_.reset();
  FindAmbiguities(*ofst);
  if (merge_) {  // Shouldn't get here; sanity test.
    FSTERROR() << "Disambiguate: Unable to remove spurious ambiguities";
    error_ = true;
    return;
  }
}

template <class Arc>
void Disambiguator<Arc>::RemoveAmbiguities(MutableFst<Arc> *ofst) {
  if (ambiguous_.empty()) return;
  // Adds dead state to redirect ambiguous transitions to be removed.
  const auto dead = ofst->AddState();
  for (auto it = ambiguous_.begin(); it != ambiguous_.end(); ++it) {
    const auto pos = it->second;
    if (pos >= 0) {  // Actual transition.
      MutableArcIterator<MutableFst<Arc>> aiter(ofst, it->first);
      aiter.Seek(pos);
      auto arc = aiter.Value();
      arc.nextstate = dead;
      aiter.SetValue(arc);
    } else {  // Super-final transition.
      ofst->SetFinal(it->first, Weight::Zero());
    }
  }
  Connect(ofst);
  ambiguous_.clear();
}

}  // namespace internal

// Disambiguates a weighted FST. This version writes the disambiguated FST to an
// output MutableFst. The result will be an equivalent FST that has the
// property that there are not two distinct paths from the initial state to a
// final state with the same input labeling.
//
// The weights must be (weakly) left divisible (valid for Tropical and
// LogWeight).
//
// Complexity:
//
//   Disambiguable: exponential (polynomial in the size of the output).
//   Non-disambiguable: does not terminate.
//
// The disambiguable transducers include all automata and functional transducers
// that are unweighted or that are acyclic or that are unambiguous.
//
// For more information, see:
//
// Mohri, M. and Riley, M. 2015. On the disambiguation of weighted automata.
// In CIAA, pages 263-278.
template <class Arc>
void Disambiguate(
    const Fst<Arc> &ifst, MutableFst<Arc> *ofst,
    const DisambiguateOptions<Arc> &opts = DisambiguateOptions<Arc>()) {
  internal::Disambiguator<Arc> disambiguator;
  disambiguator.Disambiguate(ifst, ofst, opts);
}

}  // namespace fst

#endif  // FST_DISAMBIGUATE_H_