// See www.openfst.org for extensive documentation on this weighted
// finite-state transducer library.

#ifndef FST_SCRIPT_FST_CLASS_H_
#define FST_SCRIPT_FST_CLASS_H_

#include <algorithm>
#include <limits>
#include <string>
#include <type_traits>

#include <fst/expanded-fst.h>
#include <fst/fst.h>
#include <fst/mutable-fst.h>
#include <fst/vector-fst.h>
#include <fst/script/arc-class.h>
#include <fst/script/weight-class.h>

// Classes to support "boxing" all existing types of FST arcs in a single
// FstClass which hides the arc types. This allows clients to load
// and work with FSTs without knowing the arc type. These classes are only
// recommended for use in high-level scripting applications. Most users should
// use the lower-level templated versions corresponding to these classes.

namespace fst {
namespace script {

// Abstract base class defining the set of functionalities implemented in all
// impls and passed through by all bases. Below FstClassBase the class
// hierarchy bifurcates; FstClassImplBase serves as the base class for all
// implementations (of which FstClassImpl is currently the only one) and
// FstClass serves as the base class for all interfaces.

class FstClassBase {
 public:
  virtual const string &ArcType() const = 0;
  virtual WeightClass Final(int64) const = 0;
  virtual const string &FstType() const = 0;
  virtual const SymbolTable *InputSymbols() const = 0;
  virtual size_t NumArcs(int64) const = 0;
  virtual size_t NumInputEpsilons(int64) const = 0;
  virtual size_t NumOutputEpsilons(int64) const = 0;
  virtual const SymbolTable *OutputSymbols() const = 0;
  virtual uint64 Properties(uint64, bool) const = 0;
  virtual int64 Start() const = 0;
  virtual const string &WeightType() const = 0;
  virtual bool ValidStateId(int64) const = 0;
  virtual bool Write(const string &) const = 0;
  virtual bool Write(std::ostream &, const string &) const = 0;
  virtual ~FstClassBase() {}
};

// Adds all the MutableFst methods.
class FstClassImplBase : public FstClassBase {
 public:
  virtual bool AddArc(int64, const ArcClass &) = 0;
  virtual int64 AddState() = 0;
  virtual FstClassImplBase *Copy() = 0;
  virtual bool DeleteArcs(int64, size_t) = 0;
  virtual bool DeleteArcs(int64) = 0;
  virtual bool DeleteStates(const std::vector<int64> &) = 0;
  virtual void DeleteStates() = 0;
  virtual SymbolTable *MutableInputSymbols() = 0;
  virtual SymbolTable *MutableOutputSymbols() = 0;
  virtual int64 NumStates() const = 0;
  virtual bool ReserveArcs(int64, size_t) = 0;
  virtual void ReserveStates(int64) = 0;
  virtual void SetInputSymbols(SymbolTable *) = 0;
  virtual bool SetFinal(int64, const WeightClass &) = 0;
  virtual void SetOutputSymbols(SymbolTable *) = 0;
  virtual void SetProperties(uint64, uint64) = 0;
  virtual bool SetStart(int64) = 0;
  ~FstClassImplBase() override {}
};

// Containiner class wrapping an Fst<Arc>, hiding its arc type. Whether this
// Fst<Arc> pointer refers to a special kind of FST (e.g. a MutableFst) is
// known by the type of interface class that owns the pointer to this
// container.

template <class Arc>
class FstClassImpl : public FstClassImplBase {
 public:
  explicit FstClassImpl(Fst<Arc> *impl, bool should_own = false)
      : impl_(should_own ? impl : impl->Copy()) {}

  explicit FstClassImpl(const Fst<Arc> &impl) : impl_(impl.Copy()) {}

  // Warning: calling this method casts the FST to a mutable FST.
  bool AddArc(int64 s, const ArcClass &ac) final {
    if (!ValidStateId(s)) return false;
    // Note that we do not check that the destination state is valid, so users
    // can add arcs before they add the corresponding states. Verify can be
    // used to determine whether any arc has a nonexisting destination.
    Arc arc(ac.ilabel, ac.olabel, *ac.weight.GetWeight<typename Arc::Weight>(),
            ac.nextstate);
    static_cast<MutableFst<Arc> *>(impl_.get())->AddArc(s, arc);
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  int64 AddState() final {
    return static_cast<MutableFst<Arc> *>(impl_.get())->AddState();
  }

  const string &ArcType() const final { return Arc::Type(); }

  FstClassImpl *Copy() final { return new FstClassImpl<Arc>(impl_.get()); }

  // Warning: calling this method casts the FST to a mutable FST.
  bool DeleteArcs(int64 s, size_t n) final {
    if (!ValidStateId(s)) return false;
    static_cast<MutableFst<Arc> *>(impl_.get())->DeleteArcs(s, n);
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  bool DeleteArcs(int64 s) final {
    if (!ValidStateId(s)) return false;
    static_cast<MutableFst<Arc> *>(impl_.get())->DeleteArcs(s);
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  bool DeleteStates(const std::vector<int64> &dstates) final {
    for (const auto &state : dstates)
      if (!ValidStateId(state)) return false;
    // Warning: calling this method with any integers beyond the precision of
    // the underlying FST will result in truncation.
    std::vector<typename Arc::StateId> typed_dstates(dstates.size());
    std::copy(dstates.begin(), dstates.end(), typed_dstates.begin());
    static_cast<MutableFst<Arc> *>(impl_.get())->DeleteStates(typed_dstates);
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  void DeleteStates() final {
    static_cast<MutableFst<Arc> *>(impl_.get())->DeleteStates();
  }

  WeightClass Final(int64 s) const final {
    if (!ValidStateId(s)) return WeightClass::NoWeight(WeightType());
    WeightClass w(impl_->Final(s));
    return w;
  }

  const string &FstType() const final { return impl_->Type(); }

  const SymbolTable *InputSymbols() const final {
    return impl_->InputSymbols();
  }

  // Warning: calling this method casts the FST to a mutable FST.
  SymbolTable *MutableInputSymbols() final {
    return static_cast<MutableFst<Arc> *>(impl_.get())->MutableInputSymbols();
  }

  // Warning: calling this method casts the FST to a mutable FST.
  SymbolTable *MutableOutputSymbols() final {
    return static_cast<MutableFst<Arc> *>(impl_.get())->MutableOutputSymbols();
  }

  // Signals failure by returning size_t max.
  size_t NumArcs(int64 s) const final {
    return ValidStateId(s) ? impl_->NumArcs(s)
                           : std::numeric_limits<size_t>::max();
  }

  // Signals failure by returning size_t max.
  size_t NumInputEpsilons(int64 s) const final {
    return ValidStateId(s) ? impl_->NumInputEpsilons(s)
                           : std::numeric_limits<size_t>::max();
  }

  // Signals failure by returning size_t max.
  size_t NumOutputEpsilons(int64 s) const final {
    return ValidStateId(s) ? impl_->NumOutputEpsilons(s)
                           : std::numeric_limits<size_t>::max();
  }

  // Warning: calling this method casts the FST to a mutable FST.
  int64 NumStates() const final {
    return static_cast<MutableFst<Arc> *>(impl_.get())->NumStates();
  }

  uint64 Properties(uint64 mask, bool test) const final {
    return impl_->Properties(mask, test);
  }

  // Warning: calling this method casts the FST to a mutable FST.
  bool ReserveArcs(int64 s, size_t n) final {
    if (!ValidStateId(s)) return false;
    static_cast<MutableFst<Arc> *>(impl_.get())->ReserveArcs(s, n);
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  void ReserveStates(int64 s) final {
    static_cast<MutableFst<Arc> *>(impl_.get())->ReserveStates(s);
  }

  const SymbolTable *OutputSymbols() const final {
    return impl_->OutputSymbols();
  }

  // Warning: calling this method casts the FST to a mutable FST.
  void SetInputSymbols(SymbolTable *isyms) final {
    static_cast<MutableFst<Arc> *>(impl_.get())->SetInputSymbols(isyms);
  }

  // Warning: calling this method casts the FST to a mutable FST.
  bool SetFinal(int64 s, const WeightClass &weight) final {
    if (!ValidStateId(s)) return false;
    static_cast<MutableFst<Arc> *>(impl_.get())
        ->SetFinal(s, *weight.GetWeight<typename Arc::Weight>());
    return true;
  }

  // Warning: calling this method casts the FST to a mutable FST.
  void SetOutputSymbols(SymbolTable *osyms) final {
    static_cast<MutableFst<Arc> *>(impl_.get())->SetOutputSymbols(osyms);
  }

  // Warning: calling this method casts the FST to a mutable FST.
  void SetProperties(uint64 props, uint64 mask) final {
    static_cast<MutableFst<Arc> *>(impl_.get())->SetProperties(props, mask);
  }

  // Warning: calling this method casts the FST to a mutable FST.
  bool SetStart(int64 s) final {
    if (!ValidStateId(s)) return false;
    static_cast<MutableFst<Arc> *>(impl_.get())->SetStart(s);
    return true;
  }

  int64 Start() const final { return impl_->Start(); }

  bool ValidStateId(int64 s) const final {
    // This cowardly refuses to count states if the FST is not yet expanded.
    if (!Properties(kExpanded, true)) {
      FSTERROR() << "Cannot get number of states for unexpanded FST";
      return false;
    }
    // If the FST is already expanded, CountStates calls NumStates.
    if (s < 0 || s >= CountStates(*impl_)) {
      FSTERROR() << "State ID " << s << " not valid";
      return false;
    }
    return true;
  }

  const string &WeightType() const final { return Arc::Weight::Type(); }

  bool Write(const string &fname) const final { return impl_->Write(fname); }

  bool Write(std::ostream &ostr, const string &fname) const final {
    const FstWriteOptions opts(fname);
    return impl_->Write(ostr, opts);
  }

  ~FstClassImpl() override {}

  Fst<Arc> *GetImpl() const { return impl_.get(); }

 private:
  std::unique_ptr<Fst<Arc>> impl_;
};

// BASE CLASS DEFINITIONS

class MutableFstClass;

class FstClass : public FstClassBase {
 public:
  FstClass() : impl_(nullptr) {}

  template <class Arc>
  explicit FstClass(const Fst<Arc> &fst) : impl_(new FstClassImpl<Arc>(fst)) {}

  FstClass(const FstClass &other)
      : impl_(other.impl_ == nullptr ? nullptr : other.impl_->Copy()) {}

  FstClass &operator=(const FstClass &other) {
    impl_.reset(other.impl_ == nullptr ? nullptr : other.impl_->Copy());
    return *this;
  }

  WeightClass Final(int64 s) const final { return impl_->Final(s); }

  const string &ArcType() const final { return impl_->ArcType(); }

  const string &FstType() const final { return impl_->FstType(); }

  const SymbolTable *InputSymbols() const final {
    return impl_->InputSymbols();
  }

  size_t NumArcs(int64 s) const final { return impl_->NumArcs(s); }

  size_t NumInputEpsilons(int64 s) const final {
    return impl_->NumInputEpsilons(s);
  }

  size_t NumOutputEpsilons(int64 s) const final {
    return impl_->NumOutputEpsilons(s);
  }

  const SymbolTable *OutputSymbols() const final {
    return impl_->OutputSymbols();
  }

  uint64 Properties(uint64 mask, bool test) const final {
    // Special handling for FSTs with a null impl.
    if (!impl_) return kError & mask;
    return impl_->Properties(mask, test);
  }

  static FstClass *Read(const string &fname);

  static FstClass *Read(std::istream &istrm, const string &source);

  int64 Start() const final { return impl_->Start(); }

  bool ValidStateId(int64 s) const final { return impl_->ValidStateId(s); }

  const string &WeightType() const final { return impl_->WeightType(); }

  // Helper that logs an ERROR if the weight type of an FST and a WeightClass
  // don't match.

  bool WeightTypesMatch(const WeightClass &weight, const string &op_name) const;

  bool Write(const string &fname) const final { return impl_->Write(fname); }

  bool Write(std::ostream &ostr, const string &fname) const final {
    return impl_->Write(ostr, fname);
  }

  ~FstClass() override {}

  // These methods are required by IO registration.

  template <class Arc>
  static FstClassImplBase *Convert(const FstClass &other) {
    FSTERROR() << "Doesn't make sense to convert any class to type FstClass";
    return nullptr;
  }

  template <class Arc>
  static FstClassImplBase *Create() {
    FSTERROR() << "Doesn't make sense to create an FstClass with a "
               << "particular arc type";
    return nullptr;
  }

  template <class Arc>
  const Fst<Arc> *GetFst() const {
    if (Arc::Type() != ArcType()) {
      return nullptr;
    } else {
      FstClassImpl<Arc> *typed_impl =
          static_cast<FstClassImpl<Arc> *>(impl_.get());
      return typed_impl->GetImpl();
    }
  }

  template <class Arc>
  static FstClass *Read(std::istream &stream, const FstReadOptions &opts) {
    if (!opts.header) {
      LOG(ERROR) << "FstClass::Read: Options header not specified";
      return nullptr;
    }
    const FstHeader &hdr = *opts.header;
    if (hdr.Properties() & kMutable) {
      return ReadTypedFst<MutableFstClass, MutableFst<Arc>>(stream, opts);
    } else {
      return ReadTypedFst<FstClass, Fst<Arc>>(stream, opts);
    }
  }

 protected:
  explicit FstClass(FstClassImplBase *impl) : impl_(impl) {}

  const FstClassImplBase *GetImpl() const { return impl_.get(); }

  FstClassImplBase *GetImpl() { return impl_.get(); }

  // Generic template method for reading an arc-templated FST of type
  // UnderlyingT, and returning it wrapped as FstClassT, with appropriat
  // error checking. Called from arc-templated Read() static methods.
  template <class FstClassT, class UnderlyingT>
  static FstClassT *ReadTypedFst(std::istream &stream,
                                 const FstReadOptions &opts) {
    std::unique_ptr<UnderlyingT> u(UnderlyingT::Read(stream, opts));
    return u ? new FstClassT(*u) : nullptr;
  }

 private:
  std::unique_ptr<FstClassImplBase> impl_;
};

// Specific types of FstClass with special properties

class MutableFstClass : public FstClass {
 public:
  bool AddArc(int64 s, const ArcClass &ac) {
    if (!WeightTypesMatch(ac.weight, "AddArc")) return false;
    return GetImpl()->AddArc(s, ac);
  }

  int64 AddState() { return GetImpl()->AddState(); }

  bool DeleteArcs(int64 s, size_t n) { return GetImpl()->DeleteArcs(s, n); }

  bool DeleteArcs(int64 s) { return GetImpl()->DeleteArcs(s); }

  bool DeleteStates(const std::vector<int64> &dstates) {
    return GetImpl()->DeleteStates(dstates);
  }

  void DeleteStates() { GetImpl()->DeleteStates(); }

  SymbolTable *MutableInputSymbols() {
    return GetImpl()->MutableInputSymbols();
  }

  SymbolTable *MutableOutputSymbols() {
    return GetImpl()->MutableOutputSymbols();
  }

  int64 NumStates() const { return GetImpl()->NumStates(); }

  bool ReserveArcs(int64 s, size_t n) { return GetImpl()->ReserveArcs(s, n); }

  void ReserveStates(int64 s) { GetImpl()->ReserveStates(s); }

  static MutableFstClass *Read(const string &fname, bool convert = false);

  void SetInputSymbols(SymbolTable *isyms) {
    GetImpl()->SetInputSymbols(isyms);
  }

  bool SetFinal(int64 s, const WeightClass &weight) {
    if (!WeightTypesMatch(weight, "SetFinal")) return false;
    return GetImpl()->SetFinal(s, weight);
  }

  void SetOutputSymbols(SymbolTable *osyms) {
    GetImpl()->SetOutputSymbols(osyms);
  }

  void SetProperties(uint64 props, uint64 mask) {
    GetImpl()->SetProperties(props, mask);
  }

  bool SetStart(int64 s) { return GetImpl()->SetStart(s); }

  template <class Arc>
  explicit MutableFstClass(const MutableFst<Arc> &fst) : FstClass(fst) {}

  // These methods are required by IO registration.

  template <class Arc>
  static FstClassImplBase *Convert(const FstClass &other) {
    FSTERROR() << "Doesn't make sense to convert any class to type "
               << "MutableFstClass";
    return nullptr;
  }

  template <class Arc>
  static FstClassImplBase *Create() {
    FSTERROR() << "Doesn't make sense to create a MutableFstClass with a "
               << "particular arc type";
    return nullptr;
  }

  template <class Arc>
  MutableFst<Arc> *GetMutableFst() {
    Fst<Arc> *fst = const_cast<Fst<Arc> *>(this->GetFst<Arc>());
    MutableFst<Arc> *mfst = static_cast<MutableFst<Arc> *>(fst);
    return mfst;
  }

  template <class Arc>
  static MutableFstClass *Read(std::istream &stream,
                               const FstReadOptions &opts) {
    std::unique_ptr<MutableFst<Arc>> mfst(MutableFst<Arc>::Read(stream, opts));
    return mfst ? new MutableFstClass(*mfst) : nullptr;
  }

 protected:
  explicit MutableFstClass(FstClassImplBase *impl) : FstClass(impl) {}
};

class VectorFstClass : public MutableFstClass {
 public:
  explicit VectorFstClass(FstClassImplBase *impl) : MutableFstClass(impl) {}

  explicit VectorFstClass(const FstClass &other);

  explicit VectorFstClass(const string &arc_type);

  static VectorFstClass *Read(const string &fname);

  template <class Arc>
  static VectorFstClass *Read(std::istream &stream,
                              const FstReadOptions &opts) {
    std::unique_ptr<VectorFst<Arc>> mfst(VectorFst<Arc>::Read(stream, opts));
    return mfst ? new VectorFstClass(*mfst) : nullptr;
  }

  template <class Arc>
  explicit VectorFstClass(const VectorFst<Arc> &fst) : MutableFstClass(fst) {}

  template <class Arc>
  static FstClassImplBase *Convert(const FstClass &other) {
    return new FstClassImpl<Arc>(new VectorFst<Arc>(*other.GetFst<Arc>()),
                                 true);
  }

  template <class Arc>
  static FstClassImplBase *Create() {
    return new FstClassImpl<Arc>(new VectorFst<Arc>(), true);
  }
};

}  // namespace script
}  // namespace fst

#endif  // FST_SCRIPT_FST_CLASS_H_