// fstext/trivial-factor-weight.h

// Copyright 2009-2011  Microsoft Corporation

// See ../../COPYING for clarification regarding multiple authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//  http://www.apache.org/licenses/LICENSE-2.0
//
// THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED
// WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE,
// MERCHANTABLITY OR NON-INFRINGEMENT.
// See the Apache 2 License for the specific language governing permissions and
// limitations under the License.
//
//
// This is a modified file from the OpenFST Library v1.2.7 available at
// http://www.openfst.org and released under the Apache License Version 2.0.
//
//
// See ../../COPYING for clarification regarding multiple authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Copyright 2005-2010 Google, Inc.
// Author: allauzen@google.com (Cyril Allauzen)


#ifndef KALDI_FSTEXT_TRIVIAL_FACTOR_WEIGHT_H_
#define KALDI_FSTEXT_TRIVIAL_FACTOR_WEIGHT_H_


// TrivialFactorWeight.h This is an extension to factor-weight.h in the OpenFst
// code.  It is a version of FactorWeight that creates separate states (with
// input epsilons) rather than pushing the factors forward.  This is for
// converting from Gallic FSTs, where you want the result to be a bit more
// trivial with input epsilons inserted where there are multiple output symbols.
// This has the advantage that it always works, for any input (also I just
// prefer this approach).

#include <unordered_map>
using std::unordered_map;

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

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

namespace fst {


template <class Arc>
struct TrivialFactorWeightOptions : CacheOptions {
  typedef typename Arc::Label Label;
  float delta;
  Label extra_ilabel;  // input label of extra arcs
  Label extra_olabel;  // output label of extra arcs

  TrivialFactorWeightOptions(const CacheOptions &opts, float d,
                      Label il = 0, Label ol = 0)
      : CacheOptions(opts), delta(d), extra_ilabel(il), extra_olabel(ol) {}

  explicit TrivialFactorWeightOptions(
      float d, Label il = 0, Label ol = 0)
      : delta(d), extra_ilabel(il), extra_olabel(ol) {}

  TrivialFactorWeightOptions(): delta(kDelta), extra_ilabel(0), extra_olabel(0) {}

};

namespace internal {

// Implementation class for TrivialFactorWeight
template <class A, class F>
class TrivialFactorWeightFstImpl
    : public CacheImpl<A> {
 public:
  using CacheImpl<A>::PushArc;
  using FstImpl<A>::SetType;
  using FstImpl<A>::SetProperties;
  using FstImpl<A>::Properties;
  using FstImpl<A>::SetInputSymbols;
  using FstImpl<A>::SetOutputSymbols;

  using CacheBaseImpl< CacheState<A> >::HasStart;
  using CacheBaseImpl< CacheState<A> >::HasFinal;
  using CacheBaseImpl< CacheState<A> >::HasArcs;

  typedef A Arc;
  typedef typename A::Label Label;
  typedef typename A::Weight Weight;
  typedef typename A::StateId StateId;
  typedef F FactorIterator;

  typedef DefaultCacheStore<A> Store;
  typedef typename Store::State State;

  struct Element {
    Element() {}

    Element(StateId s, Weight w) : state(s), weight(w) {}

    StateId state;     // Input state Id
    Weight weight;     // Residual weight
  };

  TrivialFactorWeightFstImpl(const Fst<A> &fst, const TrivialFactorWeightOptions<A> &opts)
      : CacheImpl<A>(opts),
        fst_(fst.Copy()),
        delta_(opts.delta),
        extra_ilabel_(opts.extra_ilabel),
        extra_olabel_(opts.extra_olabel) {
    SetType("factor-weight");
    uint64 props = fst.Properties(kFstProperties, false);
    SetProperties(FactorWeightProperties(props), kCopyProperties);

    SetInputSymbols(fst.InputSymbols());
    SetOutputSymbols(fst.OutputSymbols());
  }

  TrivialFactorWeightFstImpl(const TrivialFactorWeightFstImpl<A, F> &impl)
      : CacheImpl<A>(impl),
        fst_(impl.fst_->Copy(true)),
        delta_(impl.delta_),
        extra_ilabel_(impl.extra_ilabel_),
        extra_olabel_(impl.extra_olabel_) {
    SetType("factor-weight");
    SetProperties(impl.Properties(), kCopyProperties);
    SetInputSymbols(impl.InputSymbols());
    SetOutputSymbols(impl.OutputSymbols());
  }

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

  Weight Final(StateId s) {
    if (!HasFinal(s)) {
      const Element &e = elements_[s];
      Weight w;
      if (e.state == kNoStateId) {  // extra state inserted to represent final weights.
        FactorIterator fit(e.weight);
        if (fit.Done()) {  // cannot be factored.
          w = e.weight;  // so it's final
        } else {
          w = Weight::Zero();  // need another transition.
        }
      } else {
        if (e.weight != Weight::One()) {  // Not a real state.
          w = Weight::Zero();
        } else {  // corresponds to a "real" state.
          w = fst_->Final(e.state);
          FactorIterator fit(w);
          if (!fit.Done()) // we would have intermediate states representing this final state.
            w = Weight::Zero();
        }
      }
      this->SetFinal(s, w);
      return w;
    } else {
      return CacheImpl<A>::Final(s);
    }
  }

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

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

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

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


  // Find state corresponding to an element. Create new state
  // if element not found.
  StateId FindState(const Element &e) {
    typename ElementMap::iterator eit = element_map_.find(e);
    if (eit != element_map_.end()) {
      return (*eit).second;
    } else {
      StateId s = elements_.size();
      elements_.push_back(e);
      element_map_.insert(pair<const Element, StateId>(e, s));
      return s;
    }
  }

  // Computes the outgoing transitions from a state, creating new destination
  // states as needed.
  void Expand(StateId s) {
    CHECK(static_cast<size_t>(s) < elements_.size());
    Element e = elements_[s];
    if (e.weight != Weight::One()) {
      FactorIterator fit(e.weight);
      if (fit.Done()) {  // Cannot be factored-> create a link to dest state directly
        if (e.state != kNoStateId) {
          StateId dest = FindState(Element(e.state, Weight::One()));
          PushArc(s, Arc(extra_ilabel_, extra_olabel_, e.weight, dest));
        } // else we're done.  This is a final state.
      } else {  // Can be factored.
        const pair<Weight, Weight> &p = fit.Value();
        StateId dest = FindState(Element(e.state, p.second.Quantize(delta_)));
        PushArc(s, Arc(extra_ilabel_, extra_olabel_, p.first, dest));
      }
    } else {  // Unit weight.  This corresponds to a "real" state.
      CHECK(e.state != kNoStateId);
      for (ArcIterator< Fst<A> > ait(*fst_, e.state);
           !ait.Done();
           ait.Next()) {
        const A &arc = ait.Value();
        FactorIterator fit(arc.weight);
        if (fit.Done()) {  // cannot be factored->just link directly to dest.
          StateId dest = FindState(Element(arc.nextstate, Weight::One()));
          PushArc(s, Arc(arc.ilabel, arc.olabel, arc.weight, dest));
        } else {
          const pair<Weight, Weight> &p = fit.Value();
          StateId dest = FindState(Element(arc.nextstate, p.second.Quantize(delta_)));
          PushArc(s, Arc(arc.ilabel, arc.olabel, p.first, dest));
        }
      }
      // See if we have to add arcs for final-states [only if final-weight is factorable].
      Weight final_w = fst_->Final(e.state);
      if (final_w != Weight::Zero()) {
        FactorIterator fit(final_w);
        if (!fit.Done()) {
          const pair<Weight, Weight> &p = fit.Value();
          StateId dest = FindState(Element(kNoStateId, p.second.Quantize(delta_)));
          PushArc(s, Arc(extra_ilabel_, extra_olabel_, p.first, dest));
        }
      }
    }
    this->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 {
      return static_cast<size_t>(x.state * kPrime + x.weight.Hash());
    }
   private:
    static const int kPrime = 7853;
  };

  typedef unordered_map<Element, StateId, ElementKey, ElementEqual> ElementMap;

  std::unique_ptr<const Fst<A>> fst_;
  float delta_;
  uint32 mode_;               // factoring arc and/or final weights
  Label extra_ilabel_;        // ilabel of arc created when factoring final w's
  Label extra_olabel_;        // olabel of arc created when factoring final w's
  vector<Element> elements_;  // mapping Fst state to Elements
  ElementMap element_map_;    // mapping Elements to Fst state

};

}  // namespace internal

/// TrivialFactorWeightFst 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.
/// This algorithm differs from the one implemented in FactorWeightFst
/// in that it does not attempt to push the extra weight forward to the
/// next state: it uses a sequence of "extra" intermediate state, and
/// outputs the remaining weight right away.  This ensures that it will
/// always succeed, even for Gallic representations of FSTs that have cycles
/// with more output than input symbols.

/// Note that the code below was modified from factor-weight.h by just
/// search-and-replacing "FactorWeight" by "TrivialFactorWeight".


template <class A, class F>
class TrivialFactorWeightFst :
    public ImplToFst<internal::TrivialFactorWeightFstImpl<A, F>> {
 public:
  friend class ArcIterator< TrivialFactorWeightFst<A, F> >;
  friend class StateIterator< TrivialFactorWeightFst<A, F> >;

  typedef A Arc;
  typedef typename A::Weight Weight;
  typedef typename A::StateId StateId;
  typedef DefaultCacheStore<Arc> Store;
  typedef typename Store::State State;
  typedef internal::TrivialFactorWeightFstImpl<A, F> Impl;

  explicit TrivialFactorWeightFst(const Fst<A> &fst)
      : ImplToFst<Impl>(std::make_shared<Impl>(fst, TrivialFactorWeightOptions<A>())) {}

  TrivialFactorWeightFst(const Fst<A> &fst,  const TrivialFactorWeightOptions<A> &opts)
      : ImplToFst<Impl>(std::make_shared<Impl>(fst, opts)) {}

  // See Fst<>::Copy() for doc.
  TrivialFactorWeightFst(const TrivialFactorWeightFst<A, F> &fst, bool copy)
      : ImplToFst<Impl>(fst, copy) {}

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

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

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

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

  TrivialFactorWeightFst &operator=(const TrivialFactorWeightFst &fst) = delete;
};


// Specialization for TrivialFactorWeightFst.
template<class A, class F>
class StateIterator< TrivialFactorWeightFst<A, F> >
    : public CacheStateIterator< TrivialFactorWeightFst<A, F> > {
 public:
  explicit StateIterator(const TrivialFactorWeightFst<A, F> &fst)
      : CacheStateIterator< TrivialFactorWeightFst<A, F> >(fst, fst.GetMutableImpl()) {}
};


// Specialization for TrivialFactorWeightFst.
template <class A, class F>
class ArcIterator< TrivialFactorWeightFst<A, F> >
    : public CacheArcIterator< TrivialFactorWeightFst<A, F> > {
 public:
  typedef typename A::StateId StateId;

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

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




}  // namespace fst

#endif