// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.
//
// Compose an MPDT and an FST.

#ifndef FST_EXTENSIONS_MPDT_COMPOSE_H_
#define FST_EXTENSIONS_MPDT_COMPOSE_H_

#include <list>

#include <fst/extensions/mpdt/mpdt.h>
#include <fst/extensions/pdt/compose.h>
#include <fst/compose.h>

namespace fst {

template <class Filter>
class MPdtParenFilter {
 public:
  using FST1 = typename Filter::FST1;
  using FST2 = typename Filter::FST2;
  using Arc = typename Filter::Arc;
  using Label = typename Arc::Label;
  using StateId = typename Arc::StateId;
  using Weight = typename Arc::Weight;

  using Matcher1 = typename Filter::Matcher1;
  using Matcher2 = typename Filter::Matcher2;

  using StackId = StateId;
  using ParenStack = internal::MPdtStack<StackId, Label>;
  using FilterState1 = typename Filter::FilterState;
  using FilterState2 = IntegerFilterState<StackId>;
  using FilterState = PairFilterState<FilterState1, FilterState2>;

  MPdtParenFilter(const FST1 &fst1, const FST2 &fst2,
                  Matcher1 *matcher1 = nullptr, Matcher2 *matcher2 = nullptr,
                  const std::vector<std::pair<Label, Label>> *parens = nullptr,
                  const std::vector<Label> *assignments = nullptr,
                  bool expand = false, bool keep_parens = true)
      : filter_(fst1, fst2, matcher1, matcher2),
        parens_(parens ? *parens : std::vector<std::pair<Label, Label>>()),
        assignments_(assignments ? *assignments : std::vector<Label>()),
        expand_(expand),
        keep_parens_(keep_parens),
        fs_(FilterState::NoState()),
        stack_(parens_, assignments_),
        paren_id_(-1) {
    if (parens) {
      for (const auto &pair : *parens) {
        parens_.push_back(pair);
        GetMatcher1()->AddOpenParen(pair.first);
        GetMatcher2()->AddOpenParen(pair.first);
        if (!expand_) {
          GetMatcher1()->AddCloseParen(pair.second);
          GetMatcher2()->AddCloseParen(pair.second);
        }
      }
    }
  }

  MPdtParenFilter(const MPdtParenFilter &filter, bool safe = false)
      : filter_(filter.filter_, safe),
        parens_(filter.parens_),
        expand_(filter.expand_),
        keep_parens_(filter.keep_parens_),
        fs_(FilterState::NoState()),
        stack_(filter.parens_, filter.assignments_),
        paren_id_(-1) {}

  FilterState Start() const {
    return FilterState(filter_.Start(), FilterState2(0));
  }

  void SetState(StateId s1, StateId s2, const FilterState &fs) {
    fs_ = fs;
    filter_.SetState(s1, s2, fs_.GetState1());
    if (!expand_) return;
    const auto paren_id = stack_.Top(fs.GetState2().GetState());
    if (paren_id != paren_id_) {
      if (paren_id_ != -1) {
        GetMatcher1()->RemoveCloseParen(parens_[paren_id_].second);
        GetMatcher2()->RemoveCloseParen(parens_[paren_id_].second);
      }
      paren_id_ = paren_id;
      if (paren_id_ != -1) {
        GetMatcher1()->AddCloseParen(parens_[paren_id_].second);
        GetMatcher2()->AddCloseParen(parens_[paren_id_].second);
      }
    }
  }

  FilterState FilterArc(Arc *arc1, Arc *arc2) const {
    const auto fs1 = filter_.FilterArc(arc1, arc2);
    const auto &fs2 = fs_.GetState2();
    if (fs1 == FilterState1::NoState()) return FilterState::NoState();
    if (arc1->olabel == kNoLabel && arc2->ilabel) {  // arc2 parentheses.
      if (keep_parens_) {
        arc1->ilabel = arc2->ilabel;
      } else if (arc2->ilabel) {
        arc2->olabel = arc1->ilabel;
      }
      return FilterParen(arc2->ilabel, fs1, fs2);
    } else if (arc2->ilabel == kNoLabel && arc1->olabel) {  // arc1 parentheses
      if (keep_parens_) {
        arc2->olabel = arc1->olabel;
      } else {
        arc1->ilabel = arc2->olabel;
      }
      return FilterParen(arc1->olabel, fs1, fs2);
    } else {
      return FilterState(fs1, fs2);
    }
  }

  void FilterFinal(Weight *w1, Weight *w2) const {
    if (fs_.GetState2().GetState() != 0) *w1 = Weight::Zero();
    filter_.FilterFinal(w1, w2);
  }

  // Returns respective matchers; ownership stays with filter.

  Matcher1 *GetMatcher1() { return filter_.GetMatcher1(); }

  Matcher2 *GetMatcher2() { return filter_.GetMatcher2(); }

  uint64 Properties(uint64 iprops) const {
    const auto oprops = filter_.Properties(iprops);
    return oprops & kILabelInvariantProperties & kOLabelInvariantProperties;
  }

 private:
  const FilterState FilterParen(Label label, const FilterState1 &fs1,
                                const FilterState2 &fs2) const {
    if (!expand_) return FilterState(fs1, fs2);
    const auto stack_id = stack_.Find(fs2.GetState(), label);
    if (stack_id < 0) {
      return FilterState::NoState();
    } else {
      return FilterState(fs1, FilterState2(stack_id));
    }
  }

  Filter filter_;
  std::vector<std::pair<Label, Label>> parens_;
  std::vector<Label> assignments_;
  bool expand_;       // Expands to FST?
  bool keep_parens_;  // Retains parentheses in output?
  FilterState fs_;    // Current filter state.
  mutable ParenStack stack_;
  ssize_t paren_id_;
};

// Class to setup composition options for MPDT composition. Default is to take
// the MPDT as the first composition argument.
template <class Arc, bool left_pdt = true>
class MPdtComposeFstOptions
    : public ComposeFstOptions<Arc, ParenMatcher<Fst<Arc>>,
                               MPdtParenFilter<AltSequenceComposeFilter<
                                   ParenMatcher<Fst<Arc>> >> > {
 public:
  using Label = typename Arc::Label;
  using MPdtMatcher = ParenMatcher<Fst<Arc>>;
  using MPdtFilter = MPdtParenFilter<AltSequenceComposeFilter<MPdtMatcher>>;

  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::matcher1;
  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::matcher2;
  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::filter;

  MPdtComposeFstOptions(const Fst<Arc> &ifst1,
                        const std::vector<std::pair<Label, Label>> &parens,
                        const std::vector<typename Arc::Label> &assignments,
                        const Fst<Arc> &ifst2, bool expand = false,
                        bool keep_parens = true) {
    matcher1 = new MPdtMatcher(ifst1, MATCH_OUTPUT, kParenList);
    matcher2 = new MPdtMatcher(ifst2, MATCH_INPUT, kParenLoop);
    filter = new MPdtFilter(ifst1, ifst2, matcher1, matcher2, &parens,
                            &assignments, expand, keep_parens);
  }
};

// Class to setup composition options for PDT with FST composition.
// Specialization is for the FST as the first composition argument.
template <class Arc>
class MPdtComposeFstOptions<Arc, false>
    : public ComposeFstOptions<
          Arc, ParenMatcher<Fst<Arc>>,
          MPdtParenFilter<SequenceComposeFilter<ParenMatcher<Fst<Arc>> >> > {
 public:
  using Label = typename Arc::Label;
  using MPdtMatcher = ParenMatcher<Fst<Arc>>;
  using MPdtFilter = MPdtParenFilter<SequenceComposeFilter<MPdtMatcher>>;

  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::matcher1;
  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::matcher2;
  using ComposeFstOptions<Arc, MPdtMatcher, MPdtFilter>::filter;

  MPdtComposeFstOptions(const Fst<Arc> &ifst1, const Fst<Arc> &ifst2,
                        const std::vector<std::pair<Label, Label>> &parens,
                        const std::vector<typename Arc::Label> &assignments,
                        bool expand = false, bool keep_parens = true) {
    matcher1 = new MPdtMatcher(ifst1, MATCH_OUTPUT, kParenLoop);
    matcher2 = new MPdtMatcher(ifst2, MATCH_INPUT, kParenList);
    filter = new MPdtFilter(ifst1, ifst2, matcher1, matcher2, &parens,
                            &assignments, expand, keep_parens);
  }
};

struct MPdtComposeOptions {
  bool connect;                  // Connect output?
  PdtComposeFilter filter_type;  // Which pre-defined filter to use.

  explicit MPdtComposeOptions(bool connect = true,
                              PdtComposeFilter filter_type = PAREN_FILTER)
      : connect(connect), filter_type(filter_type) {}
};

// Composes multi-pushdown transducer (MPDT) encoded as an FST (1st arg) and an
// FST (2nd arg) with the result also an MPDT encoded as an FST (3rd arg). In
// theMPDTs, some transitions are labeled with open or close parentheses (and
// associated with a stack). To be interpreted as an MPDT, the parents on each
// stack must balance on a path (see MPdtExpand()). The open-close parenthesis
// label pairs are passed using the parens arguments, and the stack assignments
// are passed using the assignments argument.
template <class Arc>
void Compose(
    const Fst<Arc> &ifst1,
    const std::vector<std::pair<typename Arc::Label, typename Arc::Label>>
        &parens,
    const std::vector<typename Arc::Label> &assignments, const Fst<Arc> &ifst2,
    MutableFst<Arc> *ofst,
    const MPdtComposeOptions &opts = MPdtComposeOptions()) {
  bool expand = opts.filter_type != PAREN_FILTER;
  bool keep_parens = opts.filter_type != EXPAND_FILTER;
  MPdtComposeFstOptions<Arc, true> copts(ifst1, parens, assignments, ifst2,
                                         expand, keep_parens);
  copts.gc_limit = 0;
  *ofst = ComposeFst<Arc>(ifst1, ifst2, copts);
  if (opts.connect) Connect(ofst);
}

// Composes an FST (1st arg) and a multi-pushdown transducer (MPDT) encoded as
// an FST (2nd arg) with the result also an MPDT encoded as an FST (3rd arg).
// In the MPDTs, some transitions are labeled with open or close parentheses
// (and associated with a stack). To be interpreted as an MPDT, the parents on
// each stack must balance on a path (see MPdtExpand()). The open-close
// parenthesis label pairs are passed using the parens arguments, and the stack
// assignments are passed using the assignments argument.
template <class Arc>
void Compose(
    const Fst<Arc> &ifst1, const Fst<Arc> &ifst2,
    const std::vector<std::pair<typename Arc::Label, typename Arc::Label>>
        &parens,
    const std::vector<typename Arc::Label> &assignments, MutableFst<Arc> *ofst,
    const MPdtComposeOptions &opts = MPdtComposeOptions()) {
  bool expand = opts.filter_type != PAREN_FILTER;
  bool keep_parens = opts.filter_type != EXPAND_FILTER;
  MPdtComposeFstOptions<Arc, false> copts(ifst1, ifst2, parens, assignments,
                                          expand, keep_parens);
  copts.gc_limit = 0;
  *ofst = ComposeFst<Arc>(ifst1, ifst2, copts);
  if (opts.connect) Connect(ofst);
}

}  // namespace fst

#endif  // FST_EXTENSIONS_MPDT_COMPOSE_H_