// See www.openfst.org for extensive documentation on this weighted
  // finite-state transducer library.
  //
  // Class to encode and decode an FST.
  
  #ifndef FST_ENCODE_H_
  #define FST_ENCODE_H_
  
  #include <iostream>
  #include <memory>
  #include <string>
  #include <unordered_map>
  #include <utility>
  #include <vector>
  
  #include <fst/log.h>
  #include <fstream>
  
  #include <fst/arc-map.h>
  #include <fst/rmfinalepsilon.h>
  
  
  namespace fst {
  
  enum EncodeType { ENCODE = 1, DECODE = 2 };
  
  static constexpr uint32 kEncodeLabels = 0x0001;
  static constexpr uint32 kEncodeWeights = 0x0002;
  static constexpr uint32 kEncodeFlags = 0x0003;
  
  namespace internal {
  
  static constexpr uint32 kEncodeHasISymbols = 0x0004;
  static constexpr uint32 kEncodeHasOSymbols = 0x0008;
  
  // Identifies stream data as an encode table (and its endianity)
  static const int32 kEncodeMagicNumber = 2129983209;
  
  // The following class encapsulates implementation details for the encoding and
  // decoding of label/weight tuples used for encoding and decoding of FSTs. The
  // EncodeTable is bidirectional. I.e, it stores both the Tuple of encode labels
  // and weights to a unique label, and the reverse.
  template <class Arc>
  class EncodeTable {
   public:
    using Label = typename Arc::Label;
    using Weight = typename Arc::Weight;
  
    // Encoded data consists of arc input/output labels and arc weight.
    struct Tuple {
      Tuple() {}
  
      Tuple(Label ilabel_, Label olabel_, Weight weight_)
          : ilabel(ilabel_), olabel(olabel_), weight(std::move(weight_)) {}
  
      Tuple(const Tuple &tuple)
          : ilabel(tuple.ilabel),
            olabel(tuple.olabel),
            weight(std::move(tuple.weight)) {}
  
      Label ilabel;
      Label olabel;
      Weight weight;
    };
  
    // Comparison object for hashing EncodeTable Tuple(s).
    class TupleEqual {
     public:
      bool operator()(const Tuple *x, const Tuple *y) const {
        return (x->ilabel == y->ilabel && x->olabel == y->olabel &&
                x->weight == y->weight);
      }
    };
  
    // Hash function for EncodeTabe Tuples. Based on the encode flags
    // we either hash the labels, weights or combination of them.
    class TupleKey {
     public:
      TupleKey() : encode_flags_(kEncodeLabels | kEncodeWeights) {}
  
      TupleKey(const TupleKey &key) : encode_flags_(key.encode_flags_) {}
  
      explicit TupleKey(uint32 encode_flags) : encode_flags_(encode_flags) {}
  
      size_t operator()(const Tuple *x) const {
        size_t hash = x->ilabel;
        static constexpr int lshift = 5;
        static constexpr int rshift = CHAR_BIT * sizeof(size_t) - 5;
        if (encode_flags_ & kEncodeLabels) {
          hash = hash << lshift ^ hash >> rshift ^ x->olabel;
        }
        if (encode_flags_ & kEncodeWeights) {
          hash = hash << lshift ^ hash >> rshift ^ x->weight.Hash();
        }
        return hash;
      }
  
     private:
      int32 encode_flags_;
    };
  
    explicit EncodeTable(uint32 encode_flags)
        : flags_(encode_flags), encode_hash_(1024, TupleKey(encode_flags)) {}
  
    using EncodeHash = std::unordered_map<const Tuple *, Label, TupleKey,
                                          TupleEqual>;
  
    // Given an arc, encodes either input/output labels or input/costs or both.
    Label Encode(const Arc &arc) {
      std::unique_ptr<Tuple> tuple(
          new Tuple(arc.ilabel, flags_ & kEncodeLabels ? arc.olabel : 0,
                    flags_ & kEncodeWeights ? arc.weight : Weight::One()));
      auto insert_result = encode_hash_.insert(
          std::make_pair(tuple.get(), encode_tuples_.size() + 1));
      if (insert_result.second) encode_tuples_.push_back(std::move(tuple));
      return insert_result.first->second;
    }
  
    // Given an arc, looks up its encoded label or returns kNoLabel if not found.
    Label GetLabel(const Arc &arc) const {
      const Tuple tuple(arc.ilabel, flags_ & kEncodeLabels ? arc.olabel : 0,
                        flags_ & kEncodeWeights ? arc.weight : Weight::One());
      auto it = encode_hash_.find(&tuple);
      return (it == encode_hash_.end()) ?  kNoLabel : it->second;
    }
  
    // Given an encoded arc label, decodes back to input/output labels and costs.
    const Tuple *Decode(Label key) const {
      if (key < 1 || key > encode_tuples_.size()) {
        LOG(ERROR) << "EncodeTable::Decode: Unknown decode key: " << key;
        return nullptr;
      }
      return encode_tuples_[key - 1].get();
    }
  
    size_t Size() const { return encode_tuples_.size(); }
  
    bool Write(std::ostream &strm, const string &source) const;
  
    static EncodeTable<Arc> *Read(std::istream &strm, const string &source);
  
    uint32 Flags() const { return flags_ & kEncodeFlags; }
  
    const SymbolTable *InputSymbols() const { return isymbols_.get(); }
  
    const SymbolTable *OutputSymbols() const { return osymbols_.get(); }
  
    void SetInputSymbols(const SymbolTable *syms) {
      if (syms) {
        isymbols_.reset(syms->Copy());
        flags_ |= kEncodeHasISymbols;
      } else {
        isymbols_.reset();
        flags_ &= ~kEncodeHasISymbols;
      }
    }
  
    void SetOutputSymbols(const SymbolTable *syms) {
      if (syms) {
        osymbols_.reset(syms->Copy());
        flags_ |= kEncodeHasOSymbols;
      } else {
        osymbols_.reset();
        flags_ &= ~kEncodeHasOSymbols;
      }
    }
  
   private:
    uint32 flags_;
    std::vector<std::unique_ptr<Tuple>> encode_tuples_;
    EncodeHash encode_hash_;
    std::unique_ptr<SymbolTable> isymbols_;  // Pre-encoded input symbol table.
    std::unique_ptr<SymbolTable> osymbols_;  // Pre-encoded output symbol table.
  
    EncodeTable(const EncodeTable &) = delete;
    EncodeTable &operator=(const EncodeTable &) = delete;
  };
  
  template <class Arc>
  bool EncodeTable<Arc>::Write(std::ostream &strm,
                                    const string &source) const {
    WriteType(strm, kEncodeMagicNumber);
    WriteType(strm, flags_);
    const int64 size = encode_tuples_.size();
    WriteType(strm, size);
    for (const auto &tuple : encode_tuples_) {
      WriteType(strm, tuple->ilabel);
      WriteType(strm, tuple->olabel);
      tuple->weight.Write(strm);
    }
    if (flags_ & kEncodeHasISymbols) isymbols_->Write(strm);
    if (flags_ & kEncodeHasOSymbols) osymbols_->Write(strm);
    strm.flush();
    if (!strm) {
      LOG(ERROR) << "EncodeTable::Write: Write failed: " << source;
      return false;
    }
    return true;
  }
  
  template <class Arc>
  EncodeTable<Arc> *EncodeTable<Arc>::Read(std::istream &strm,
                                           const string &source) {
    int32 magic_number = 0;
    ReadType(strm, &magic_number);
    if (magic_number != kEncodeMagicNumber) {
      LOG(ERROR) << "EncodeTable::Read: Bad encode table header: " << source;
      return nullptr;
    }
    uint32 flags;
    ReadType(strm, &flags);
    int64 size;
    ReadType(strm, &size);
    if (!strm) {
      LOG(ERROR) << "EncodeTable::Read: Read failed: " << source;
      return nullptr;
    }
    std::unique_ptr<EncodeTable<Arc>> table(new EncodeTable<Arc>(flags));
    for (int64 i = 0; i < size; ++i) {
      std::unique_ptr<Tuple> tuple(new Tuple());
      ReadType(strm, &tuple->ilabel);
      ReadType(strm, &tuple->olabel);
      tuple->weight.Read(strm);
      if (!strm) {
        LOG(ERROR) << "EncodeTable::Read: Read failed: " << source;
        return nullptr;
      }
      table->encode_tuples_.push_back(std::move(tuple));
      table->encode_hash_[table->encode_tuples_.back().get()] =
          table->encode_tuples_.size();
    }
    if (flags & kEncodeHasISymbols) {
      table->isymbols_.reset(SymbolTable::Read(strm, source));
    }
    if (flags & kEncodeHasOSymbols) {
      table->osymbols_.reset(SymbolTable::Read(strm, source));
    }
    return table.release();
  }
  
  }  // namespace internal
  
  // A mapper to encode/decode weighted transducers. Encoding of an FST is used
  // for performing classical determinization or minimization on a weighted
  // transducer viewing it as an unweighted acceptor over encoded labels.
  //
  // The mapper stores the encoding in a local hash table (EncodeTable). This
  // table is shared (and reference-counted) between the encoder and decoder.
  // A decoder has read-only access to the EncodeTable.
  //
  // The EncodeMapper allows on the fly encoding of the machine. As the
  // EncodeTable is generated the same table may by used to decode the machine
  // on the fly. For example in the following sequence of operations
  //
  //  Encode -> Determinize -> Decode
  //
  // we will use the encoding table generated during the encode step in the
  // decode, even though the encoding is not complete.
  template <class Arc>
  class EncodeMapper {
    using Label = typename Arc::Label;
    using Weight = typename Arc::Weight;
  
   public:
    EncodeMapper(uint32 flags, EncodeType type)
        : flags_(flags),
          type_(type),
          table_(std::make_shared<internal::EncodeTable<Arc>>(flags)),
          error_(false) {}
  
    EncodeMapper(const EncodeMapper &mapper)
        : flags_(mapper.flags_),
          type_(mapper.type_),
          table_(mapper.table_),
          error_(false) {}
  
    // Copy constructor but setting the type, typically to DECODE.
    EncodeMapper(const EncodeMapper &mapper, EncodeType type)
        : flags_(mapper.flags_),
          type_(type),
          table_(mapper.table_),
          error_(mapper.error_) {}
  
    Arc operator()(const Arc &arc);
  
    MapFinalAction FinalAction() const {
      return (type_ == ENCODE && (flags_ & kEncodeWeights))
                 ? MAP_REQUIRE_SUPERFINAL
                 : MAP_NO_SUPERFINAL;
    }
  
    constexpr MapSymbolsAction InputSymbolsAction() const {
      return MAP_CLEAR_SYMBOLS;
    }
  
    constexpr MapSymbolsAction OutputSymbolsAction() const {
      return MAP_CLEAR_SYMBOLS;
    }
  
    uint64 Properties(uint64 inprops) {
      uint64 outprops = inprops;
      if (error_) outprops |= kError;
      uint64 mask = kFstProperties;
      if (flags_ & kEncodeLabels) {
        mask &= kILabelInvariantProperties & kOLabelInvariantProperties;
      }
      if (flags_ & kEncodeWeights) {
        mask &= kILabelInvariantProperties & kWeightInvariantProperties &
                (type_ == ENCODE ? kAddSuperFinalProperties
                                 : kRmSuperFinalProperties);
      }
      return outprops & mask;
    }
  
    uint32 Flags() const { return flags_; }
  
    EncodeType Type() const { return type_; }
  
    bool Write(std::ostream &strm, const string &source) const {
      return table_->Write(strm, source);
    }
  
    bool Write(const string &filename) const {
      std::ofstream strm(filename,
                               std::ios_base::out | std::ios_base::binary);
      if (!strm) {
        LOG(ERROR) << "EncodeMap: Can't open file: " << filename;
        return false;
      }
      return Write(strm, filename);
    }
  
    static EncodeMapper<Arc> *Read(std::istream &strm, const string &source,
                                 EncodeType type = ENCODE) {
      auto *table = internal::EncodeTable<Arc>::Read(strm, source);
      return table ? new EncodeMapper(table->Flags(), type, table) : nullptr;
    }
  
    static EncodeMapper<Arc> *Read(const string &filename,
                                   EncodeType type = ENCODE) {
      std::ifstream strm(filename,
                              std::ios_base::in | std::ios_base::binary);
      if (!strm) {
        LOG(ERROR) << "EncodeMap: Can't open file: " << filename;
        return nullptr;
      }
      return Read(strm, filename, type);
    }
  
    const SymbolTable *InputSymbols() const { return table_->InputSymbols(); }
  
    const SymbolTable *OutputSymbols() const { return table_->OutputSymbols(); }
  
    void SetInputSymbols(const SymbolTable *syms) {
      table_->SetInputSymbols(syms);
    }
  
    void SetOutputSymbols(const SymbolTable *syms) {
      table_->SetOutputSymbols(syms);
    }
  
   private:
    uint32 flags_;
    EncodeType type_;
    std::shared_ptr<internal::EncodeTable<Arc>> table_;
    bool error_;
  
    explicit EncodeMapper(uint32 flags, EncodeType type,
                          internal::EncodeTable<Arc> *table)
        : flags_(flags), type_(type), table_(table), error_(false) {}
  
    EncodeMapper &operator=(const EncodeMapper &) = delete;
  };
  
  template <class Arc>
  Arc EncodeMapper<Arc>::operator()(const Arc &arc) {
    if (type_ == ENCODE) {
      if ((arc.nextstate == kNoStateId && !(flags_ & kEncodeWeights)) ||
          (arc.nextstate == kNoStateId && (flags_ & kEncodeWeights) &&
           arc.weight == Weight::Zero())) {
        return arc;
      } else {
        const auto label = table_->Encode(arc);
        return Arc(label, flags_ & kEncodeLabels ? label : arc.olabel,
                   flags_ & kEncodeWeights ? Weight::One() : arc.weight,
                   arc.nextstate);
      }
    } else {  // type_ == DECODE
      if (arc.nextstate == kNoStateId) {
        return arc;
      } else {
        if (arc.ilabel == 0) return arc;
        if (flags_ & kEncodeLabels && arc.ilabel != arc.olabel) {
          FSTERROR() << "EncodeMapper: Label-encoded arc has different "
                        "input and output labels";
          error_ = true;
        }
        if (flags_ & kEncodeWeights && arc.weight != Weight::One()) {
          FSTERROR() << "EncodeMapper: Weight-encoded arc has non-trivial weight";
          error_ = true;
        }
        const auto tuple = table_->Decode(arc.ilabel);
        if (!tuple) {
          FSTERROR() << "EncodeMapper: Decode failed";
          error_ = true;
          return Arc(kNoLabel, kNoLabel, Weight::NoWeight(), arc.nextstate);
        } else {
          return Arc(tuple->ilabel,
                     flags_ & kEncodeLabels ? tuple->olabel : arc.olabel,
                     flags_ & kEncodeWeights ? tuple->weight : arc.weight,
                     arc.nextstate);
        }
      }
    }
  }
  
  // Complexity: O(E + V).
  template <class Arc>
  inline void Encode(MutableFst<Arc> *fst, EncodeMapper<Arc> *mapper) {
    mapper->SetInputSymbols(fst->InputSymbols());
    mapper->SetOutputSymbols(fst->OutputSymbols());
    ArcMap(fst, mapper);
  }
  
  template <class Arc>
  inline void Decode(MutableFst<Arc> *fst, const EncodeMapper<Arc> &mapper) {
    ArcMap(fst, EncodeMapper<Arc>(mapper, DECODE));
    RmFinalEpsilon(fst);
    fst->SetInputSymbols(mapper.InputSymbols());
    fst->SetOutputSymbols(mapper.OutputSymbols());
  }
  
  // On-the-fly encoding of an input FST.
  //
  // Complexity:
  //
  //   Construction: O(1)
  //   Traversal: O(e + v)
  //
  // where e is the number of arcs visited and v is the number of states visited.
  // Constant time and space to visit an input state or arc is assumed and
  // exclusive of caching.
  template <class Arc>
  class EncodeFst : public ArcMapFst<Arc, Arc, EncodeMapper<Arc>> {
   public:
    using Mapper = EncodeMapper<Arc>;
    using Impl = internal::ArcMapFstImpl<Arc, Arc, Mapper>;
  
    EncodeFst(const Fst<Arc> &fst, Mapper *encoder)
        : ArcMapFst<Arc, Arc, Mapper>(fst, encoder, ArcMapFstOptions()) {
      encoder->SetInputSymbols(fst.InputSymbols());
      encoder->SetOutputSymbols(fst.OutputSymbols());
    }
  
    EncodeFst(const Fst<Arc> &fst, const Mapper &encoder)
        : ArcMapFst<Arc, Arc, Mapper>(fst, encoder, ArcMapFstOptions()) {}
  
    // See Fst<>::Copy() for doc.
    EncodeFst(const EncodeFst<Arc> &fst, bool copy = false)
        : ArcMapFst<Arc, Arc, Mapper>(fst, copy) {}
  
    // Makes a copy of this EncodeFst. See Fst<>::Copy() for further doc.
    EncodeFst<Arc> *Copy(bool safe = false) const override {
      if (safe) {
        FSTERROR() << "EncodeFst::Copy(true): Not allowed";
        GetImpl()->SetProperties(kError, kError);
      }
      return new EncodeFst(*this);
    }
  
   private:
    using ImplToFst<Impl>::GetImpl;
    using ImplToFst<Impl>::GetMutableImpl;
  };
  
  // On-the-fly decoding of an input FST.
  //
  // Complexity:
  //
  //   Construction: O(1).
  //   Traversal: O(e + v)
  //
  // Constant time and space to visit an input state or arc is assumed and
  // exclusive of caching.
  template <class Arc>
  class DecodeFst : public ArcMapFst<Arc, Arc, EncodeMapper<Arc>> {
   public:
    using Mapper = EncodeMapper<Arc>;
    using Impl = internal::ArcMapFstImpl<Arc, Arc, Mapper>;
    using ImplToFst<Impl>::GetImpl;
  
    DecodeFst(const Fst<Arc> &fst, const Mapper &encoder)
        : ArcMapFst<Arc, Arc, Mapper>(fst, Mapper(encoder, DECODE),
                                      ArcMapFstOptions()) {
      GetMutableImpl()->SetInputSymbols(encoder.InputSymbols());
      GetMutableImpl()->SetOutputSymbols(encoder.OutputSymbols());
    }
  
    // See Fst<>::Copy() for doc.
    DecodeFst(const DecodeFst<Arc> &fst, bool safe = false)
        : ArcMapFst<Arc, Arc, Mapper>(fst, safe) {}
  
    // Makes a copy of this DecodeFst. See Fst<>::Copy() for further doc.
    DecodeFst<Arc> *Copy(bool safe = false) const override {
      return new DecodeFst(*this, safe);
    }
  
   private:
    using ImplToFst<Impl>::GetMutableImpl;
  };
  
  // Specialization for EncodeFst.
  template <class Arc>
  class StateIterator<EncodeFst<Arc>>
      : public StateIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>> {
   public:
    explicit StateIterator(const EncodeFst<Arc> &fst)
        : StateIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>>(fst) {}
  };
  
  // Specialization for EncodeFst.
  template <class Arc>
  class ArcIterator<EncodeFst<Arc>>
      : public ArcIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>> {
   public:
    ArcIterator(const EncodeFst<Arc> &fst, typename Arc::StateId s)
        : ArcIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>>(fst, s) {}
  };
  
  // Specialization for DecodeFst.
  template <class Arc>
  class StateIterator<DecodeFst<Arc>>
      : public StateIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>> {
   public:
    explicit StateIterator(const DecodeFst<Arc> &fst)
        : StateIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>>(fst) {}
  };
  
  // Specialization for DecodeFst.
  template <class Arc>
  class ArcIterator<DecodeFst<Arc>>
      : public ArcIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>> {
   public:
    ArcIterator(const DecodeFst<Arc> &fst, typename Arc::StateId s)
        : ArcIterator<ArcMapFst<Arc, Arc, EncodeMapper<Arc>>>(fst, s) {}
  };
  
  // Useful aliases when using StdArc.
  
  using StdEncodeFst = EncodeFst<StdArc>;
  
  using StdDecodeFst = DecodeFst<StdArc>;
  
  }  // namespace fst
  
  #endif  // FST_ENCODE_H_