// lat/minimize-lattice.cc
  
  // Copyright 2009-2011  Saarland University (Author: Arnab Ghoshal)
  //           2012-2013  Johns Hopkins University (Author: Daniel Povey);  Chao Weng;
  //                      Bagher BabaAli
  //                2014  Guoguo Chen
  
  // 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.
  
  
  #include "lat/minimize-lattice.h"
  #include "hmm/transition-model.h"
  #include "util/stl-utils.h"
  
  namespace fst {
  
  /*
    Process the states in reverse topological order.
    For each state, compute a hash-value that will be the same for states
    that can be combined.  Then for each pair of states with the
    same hash value, check that the "to-states" map to the
    same equivalence class and that the weights are sufficiently similar.
  */
    
  template<class Weight, class IntType> class CompactLatticeMinimizer {
   public:
    typedef CompactLatticeWeightTpl<Weight, IntType> CompactWeight;
    typedef ArcTpl<CompactWeight> CompactArc;
    typedef typename CompactArc::StateId StateId;
    typedef typename CompactArc::Label Label;
    typedef size_t HashType;
    
    CompactLatticeMinimizer(MutableFst<CompactArc> *clat,
                            float delta = fst::kDelta):
        clat_(clat), delta_(delta) { }
  
    bool Minimize() {
      if (clat_->Properties(kTopSorted, true) == 0) {
        if (!TopSort(clat_)) {
          KALDI_WARN << "Topological sorting of state-level lattice failed "
              "(probably your lexicon has empty words or your LM has epsilon cycles; this "
              " is a bad idea.)";
          return false;
        }
      }
      ComputeStateHashValues();
      ComputeStateMap();
      ModifyModel();
      return true;
    }
    
    static HashType ConvertStringToHashValue(const std::vector<IntType> &vec) {
      const HashType prime = 53281;
      kaldi::VectorHasher<IntType> h;
      HashType ans = static_cast<HashType>(h(vec));
      if (ans == 0)  ans = prime;
      // We don't allow a zero answer, as this can cause too many values to be the
      // same.
      return ans;
    }
    
    static void InitHashValue(const CompactWeight &final_weight, HashType *h) {
      const HashType prime1 = 33317, prime2 = 607; // it's pretty random.
      if (final_weight == CompactWeight::Zero()) *h = prime1;
      else *h = prime2 * ConvertStringToHashValue(final_weight.String());
    }
  
    // It's important that this function and UpdateHashValueForFinalProb be
    // insensitive to the order in which it's called, as the order of the arcs
    // won't necessarily be the same for different equivalent states.
    static void UpdateHashValueForTransition(const CompactWeight &weight,
                                             Label label,
                                             HashType &next_state_hash,
                                             HashType *h) {
      const HashType prime1 = 1447, prime2 = 51907;
      if (label == 0) label = prime2; // Zeros will cause problems.
      *h += prime1 * label *
          (1 + ConvertStringToHashValue(weight.String()) * next_state_hash);
      // Above, the "1 +" is to ensure that if somehow we get zeros due to
      // weird word sequences, they don't propagate.
    }
  
    void ComputeStateHashValues() {
      // Note: clat_ is topologically sorted, and StateId is
      // signed.  Each state's hash value is only a function of toplogically-later
      // states' hash values.
      state_hashes_.resize(clat_->NumStates());
      for (StateId s = clat_->NumStates() - 1; s >= 0; s--) {
        HashType this_hash;
        InitHashValue(clat_->Final(s), &this_hash);
        for (ArcIterator<MutableFst<CompactArc> > aiter(*clat_, s);
             !aiter.Done(); aiter.Next()) {
          const CompactArc &arc = aiter.Value();
          HashType next_hash;
          if (arc.nextstate > s) {
            next_hash = state_hashes_[arc.nextstate];
          } else {
            KALDI_ASSERT(s == arc.nextstate &&
                         "Lattice not topologically sorted [code error]");
            next_hash = 1;
            KALDI_WARN << "Minimizing lattice with self-loops "
                "(lattices should not have self-loops)";
          }
          UpdateHashValueForTransition(arc.weight, arc.ilabel,
                                       next_hash, &this_hash);
        }
        state_hashes_[s] = this_hash;
      }
    }
  
  
  
    struct EquivalenceSorter {
      // This struct has an operator () which you can interpret as a less-than (<)
      // operator for arcs.  We sort on ilabel; since the lattice is supposed to
      // be deterministic, this should completely determine the ordering (there
      // should not be more than one arc with the same ilabel, out of the same
      // state).  For identical ilabels we next sort on the nextstate, simply to
      // better handle non-deterministic input (we do our best on this, without
      // guaranteeing full minimization).  We could sort on the strings next, but
      // this would be an unnecessary hassle as we only really need good
      // performance on deterministic input.
      bool operator () (const CompactArc &a, const CompactArc &b) const {
        if (a.ilabel < b.ilabel) return true;
        else if (a.ilabel > b.ilabel) return false;
        else if (a.nextstate < b.nextstate) return true;
        else return false;
      }
    };
  
    
    // This function works out whether s and t are equivalent, assuming
    // we have already partitioned all topologically-later states into
    // equivalence classes (i.e. set up state_map_).
    bool Equivalent(StateId s, StateId t) const {
      if (!ApproxEqual(clat_->Final(s), clat_->Final(t), delta_))
        return false;
      if (clat_->NumArcs(s) != clat_->NumArcs(t))
        return false;
      std::vector<CompactArc> s_arcs;
      std::vector<CompactArc> t_arcs;
      for (int32 iter = 0; iter <= 1; iter++) {
        StateId state = (iter == 0 ? s : t);
        std::vector<CompactArc> &arcs = (iter == 0 ? s_arcs : t_arcs);
        arcs.reserve(clat_->NumArcs(s));
        for (ArcIterator<MutableFst<CompactArc> > aiter(*clat_, state);
             !aiter.Done(); aiter.Next()) {
          CompactArc arc = aiter.Value();
          if (arc.nextstate == state) {
            // This is a special case for states that have self-loops.  If two
            // states have an identical self-loop arc, they may be equivalent.
            arc.nextstate = kNoStateId;
          } else {
            KALDI_ASSERT(arc.nextstate > state);
            //while (state_map_[arc.nextstate] != arc.nextstate)
            arc.nextstate = state_map_[arc.nextstate];
            arcs.push_back(arc);
          }
        }
        EquivalenceSorter s;
        std::sort(arcs.begin(), arcs.end(), s);
      }
      KALDI_ASSERT(s_arcs.size() == t_arcs.size());
      for (size_t i = 0; i < s_arcs.size(); i++) {
        if (s_arcs[i].nextstate != t_arcs[i].nextstate) return false;
        KALDI_ASSERT(s_arcs[i].ilabel == s_arcs[i].olabel); // CompactLattices are
                                                            // supposed to be
                                                            // acceptors.
        if (s_arcs[i].ilabel != t_arcs[i].ilabel) return false;
        // We've already mapped to equivalence classes.
        if (s_arcs[i].nextstate != t_arcs[i].nextstate) return false;
        if (!ApproxEqual(s_arcs[i].weight, t_arcs[i].weight)) return false;
      }
      return true;
    }
    
    void ComputeStateMap() {
      // We have to compute the state mapping in reverse topological order also,
      // since the equivalence test relies on later states being already sorted
      // out into equivalence classes (by state_map_).
      StateId num_states = clat_->NumStates();
      unordered_map<HashType, std::vector<StateId> > hash_groups_;
      
      for (StateId s = 0; s < num_states; s++)
        hash_groups_[state_hashes_[s]].push_back(s);
  
      state_map_.resize(num_states);
      for (StateId s = 0; s < num_states; s++)
        state_map_[s] = s; // Default mapping.
      
  
      { // This block is just diagnostic.
        typedef typename unordered_map<HashType,
                std::vector<StateId> >::const_iterator HashIter;
        size_t max_size = 0;
        for (HashIter iter = hash_groups_.begin(); iter != hash_groups_.end();
             ++iter)
          max_size = std::max(max_size, iter->second.size());
        if (max_size > 1000) {
          KALDI_WARN << "Largest equivalence group (using hash) is " << max_size
                     << ", minimization might be slow.";
        }
      }
  
      for (StateId s = num_states - 1; s >= 0; s--) {
        HashType hash = state_hashes_[s];
        const std::vector<StateId> &equivalence_class = hash_groups_[hash];
        KALDI_ASSERT(!equivalence_class.empty());
        for (size_t i = 0; i < equivalence_class.size(); i++) {
          StateId t = equivalence_class[i];
          // Below, there is no point doing the test if state_map_[t] != t, because
          // in that case we will, before after this, be comparing with another state
          // that is equivalent to t.
          if (t > s && state_map_[t] == t && Equivalent(s, t)) {
            state_map_[s] = t;
            break;
          }
        }
      }
    }
  
    void ModifyModel() {    
      // Modifies the model according to state_map_;
  
      StateId num_removed = 0;
      StateId num_states = clat_->NumStates();
      for (StateId s = 0; s < num_states; s++)
        if (state_map_[s] != s)
          num_removed++;
      KALDI_VLOG(3) << "Removing " << num_removed << " of "
                    << num_states << " states.";
      if (num_removed == 0) return; // Nothing to do.
      
      clat_->SetStart(state_map_[clat_->Start()]);
  
      for (StateId s = 0; s < num_states; s++) {
        if (state_map_[s] != s)
          continue; // There is no point modifying states we're removing.
        for (MutableArcIterator<MutableFst<CompactArc> > aiter(clat_, s);
             !aiter.Done(); aiter.Next()) {
          CompactArc arc = aiter.Value();
          StateId mapped_nextstate = state_map_[arc.nextstate];
          if (mapped_nextstate != arc.nextstate) {
            arc.nextstate = mapped_nextstate;
            aiter.SetValue(arc);
          }
        }
      }
      fst::Connect(clat_);
    }
   private:
    MutableFst<ArcTpl<CompactLatticeWeightTpl<Weight, IntType> > > *clat_;
    float delta_;
    std::vector<HashType> state_hashes_;
    std::vector<StateId> state_map_; // maps each state to itself or to some
                                     // equivalent state.  Within each equivalence
                                     // class, we pick one arbitrarily.
  };
  
  template<class Weight, class IntType>
  bool MinimizeCompactLattice(
      MutableFst<ArcTpl<CompactLatticeWeightTpl<Weight, IntType> > > *clat,
      float delta) {
    CompactLatticeMinimizer<Weight, IntType> minimizer(clat, delta);
    return minimizer.Minimize();
  }
  
  // Instantiate for CompactLattice type.
  template
  bool MinimizeCompactLattice<kaldi::LatticeWeight, kaldi::int32>(
      MutableFst<kaldi::CompactLatticeArc> *clat, float delta);
    
  
  }  // namespace fst