// lat/phone-align-lattice.cc
  
  // Copyright 2012-2013  Microsoft Corporation
  //                      Johns Hopkins University (Author: Daniel Povey)
  
  // 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/phone-align-lattice.h"
  #include "hmm/transition-model.h"
  #include "util/stl-utils.h"
  
  namespace kaldi {
  
  class LatticePhoneAligner {
   public:
    typedef CompactLatticeArc::StateId StateId;
    typedef CompactLatticeArc::Label Label;
  
    class ComputationState { /// The state of the computation in which,
      /// along a single path in the lattice, we work out the phone
      /// boundaries and output phone-aligned arcs. [These may or may not have
      /// words on them; the word symbols are not aligned with anything.
     public:
  
      /// Advance the computation state by adding the symbols and weights
      /// from this arc.  Gets rid of the weight and puts it in "weight" which
      /// will be put on the output arc; this keeps the state-space small.
      void Advance(const CompactLatticeArc &arc, const PhoneAlignLatticeOptions &opts,
                   LatticeWeight *weight) {
        const std::vector<int32> &string = arc.weight.String();
        transition_ids_.insert(transition_ids_.end(),
                               string.begin(), string.end());
        if (arc.ilabel != 0 && !opts.replace_output_symbols) // note: arc.ilabel==arc.olabel (acceptor)
          word_labels_.push_back(arc.ilabel);
        *weight = Times(weight_, arc.weight.Weight());
        weight_ = LatticeWeight::One();
      }
  
      /// If it can output a whole phone, it will do so, will put it in arc_out,
      /// and return true; else it will return false.  If it detects an error
      /// condition and *error = false, it will set *error to true and print
      /// a warning.  In this case it will still output phone arcs, they will
      /// just be inaccurate.  Of course once *error is set, something has gone
      /// wrong so don't trust the output too fully.
      /// Note: the "next_state" of the arc will not be set, you have to do that
      /// yourself.
      bool OutputPhoneArc(const TransitionModel &tmodel,
                          const PhoneAlignLatticeOptions &opts,
                          CompactLatticeArc *arc_out,
                          bool *error);
  
      /// This will succeed (and output the arc) if we have >1 word in words_;
      /// the arc won't have any transition-ids on it.  This is intended to fix
      /// a particular pathology where too many words were pending and we had
      /// blowup.
      bool OutputWordArc(const TransitionModel &tmodel,
                         const PhoneAlignLatticeOptions &opts,
                         CompactLatticeArc *arc_out,
                         bool *error);
  
      bool IsEmpty() { return (transition_ids_.empty() && word_labels_.empty()); }
  
      /// FinalWeight() will return "weight" if both transition_ids
      /// and word_labels are empty, otherwise it will return
      /// Weight::Zero().
      LatticeWeight FinalWeight() { return (IsEmpty() ? weight_ : LatticeWeight::Zero()); }
  
      /// This function may be called when you reach the end of
      /// the lattice and this structure hasn't voluntarily
      /// output words using "OutputArc".  If IsEmpty() == false,
      /// then you can call this function and it will output
      /// an arc.  The only
      /// non-error state in which this happens, is when a word
      /// (or silence) has ended, but we don't know that it's
      /// ended because we haven't seen the first transition-id
      /// from the next word.  Otherwise (error state), the output
      /// will consist of partial words, and this will only
      /// happen for lattices that were somehow broken, i.e.
      /// had not reached the final state.
      void OutputArcForce(const TransitionModel &tmodel,
                          const PhoneAlignLatticeOptions &opts,
                          CompactLatticeArc *arc_out,
                          bool *error);
  
      size_t Hash() const {
        VectorHasher<int32> vh;
        return vh(transition_ids_) + 90647 * vh(word_labels_);
        // 90647 is an arbitrary largish prime number.
        // We don't bother including the weight in the hash--
        // we don't really expect duplicates with the same vectors
        // but different weights, and anyway, this is only an
        // efficiency issue.
      }
  
      // Just need an arbitrary complete order.
      bool operator == (const ComputationState &other) const {
        return (transition_ids_ == other.transition_ids_
                && word_labels_ == other.word_labels_
                && weight_ == other.weight_);
      }
  
      ComputationState(): weight_(LatticeWeight::One()) { } // initial state.
      ComputationState(const ComputationState &other):
          transition_ids_(other.transition_ids_), word_labels_(other.word_labels_),
          weight_(other.weight_) { }
     private:
      std::vector<int32> transition_ids_;
      std::vector<int32> word_labels_;
      LatticeWeight weight_; // contains two floats.
    };
  
  
    struct Tuple {
      Tuple(StateId input_state, ComputationState comp_state):
          input_state(input_state), comp_state(comp_state) {}
      StateId input_state;
      ComputationState comp_state;
    };
  
    struct TupleHash {
      size_t operator() (const Tuple &state) const {
        return state.input_state + 102763 * state.comp_state.Hash();
        // 102763 is just an arbitrary prime number
      }
    };
    struct TupleEqual {
      bool operator () (const Tuple &state1, const Tuple &state2) const {
        // treat this like operator ==
        return (state1.input_state == state2.input_state
                && state1.comp_state == state2.comp_state);
      }
    };
  
    typedef unordered_map<Tuple, StateId, TupleHash, TupleEqual> MapType;
  
    StateId GetStateForTuple(const Tuple &tuple, bool add_to_queue) {
      MapType::iterator iter = map_.find(tuple);
      if (iter == map_.end()) { // not in map.
        StateId output_state = lat_out_->AddState();
        map_[tuple] = output_state;
        if (add_to_queue)
          queue_.push_back(std::make_pair(tuple, output_state));
        return output_state;
      } else {
        return iter->second;
      }
    }
  
    void ProcessFinal(Tuple tuple, StateId output_state) {
      // ProcessFinal is only called if the input_state has
      // final-prob of One().  [else it should be zero.  This
      // is because we called CreateSuperFinal().]
  
      if (tuple.comp_state.IsEmpty()) { // computation state doesn't have
        // anything pending.
        std::vector<int32> empty_vec;
        CompactLatticeWeight cw(tuple.comp_state.FinalWeight(), empty_vec);
        lat_out_->SetFinal(output_state, Plus(lat_out_->Final(output_state), cw));
      } else {
        // computation state has something pending, i.e. input or
        // output symbols that need to be flushed out.  Note: OutputArc() would
        // have returned false or we wouldn't have been called, so we have to
        // force it out.
        CompactLatticeArc lat_arc;
        tuple.comp_state.OutputArcForce(tmodel_, opts_, &lat_arc, &error_);
        lat_arc.nextstate = GetStateForTuple(tuple, true); // true == add to queue.
        // The final-prob stuff will get called again from ProcessQueueElement().
        // Note: because we did CreateSuperFinal(), this final-state on the input
        // lattice will have no output arcs (and unit final-prob), so there will be
        // no complications with processing the arcs from this state (there won't
        // be any).
        KALDI_ASSERT(output_state != lat_arc.nextstate);
        lat_out_->AddArc(output_state, lat_arc);
      }
    }
  
  
    void ProcessQueueElement() {
      KALDI_ASSERT(!queue_.empty());
      Tuple tuple = queue_.back().first;
      StateId output_state = queue_.back().second;
      queue_.pop_back();
  
      // First thing is-- we see whether the computation-state has something
      // pending that it wants to output.  In this case we don't do
      // anything further.  This is a chosen behavior similar to the
      // epsilon-sequencing rules encoded by the filters in
      // composition.
      CompactLatticeArc lat_arc;
      Tuple tuple2(tuple); // temp
      if (tuple.comp_state.OutputPhoneArc(tmodel_, opts_, &lat_arc, &error_) ||
          tuple.comp_state.OutputWordArc(tmodel_, opts_, &lat_arc, &error_)) {
        // note: this function changes the tuple (when it returns true).
        lat_arc.nextstate = GetStateForTuple(tuple, true); // true == add to queue,
        // if not already present.
        KALDI_ASSERT(output_state != lat_arc.nextstate);
        lat_out_->AddArc(output_state, lat_arc);
      } else {
        // when there's nothing to output, we'll process arcs from the input-state.
        // note: it would in a sense be valid to do both (i.e. process the stuff
        // above, and also these), but this is a bit like the epsilon-sequencing
        // stuff in composition: we avoid duplicate arcs by doing it this way.
  
        if (lat_.Final(tuple.input_state) != CompactLatticeWeight::Zero()) {
          KALDI_ASSERT(lat_.Final(tuple.input_state) == CompactLatticeWeight::One());
          // ... since we did CreateSuperFinal.
          ProcessFinal(tuple, output_state);
        }
        // Now process the arcs.  Note: final-state shouldn't have any arcs.
        for(fst::ArcIterator<CompactLattice> aiter(lat_, tuple.input_state);
            !aiter.Done(); aiter.Next()) {
          const CompactLatticeArc &arc = aiter.Value();
          Tuple next_tuple(tuple);
          LatticeWeight weight;
          next_tuple.comp_state.Advance(arc, opts_, &weight);
          next_tuple.input_state = arc.nextstate;
          StateId next_output_state = GetStateForTuple(next_tuple, true); // true == add to queue,
          // if not already present.
          // We add an epsilon arc here (as the input and output happens
          // separately)... the epsilons will get removed later.
          KALDI_ASSERT(next_output_state != output_state);
          lat_out_->AddArc(output_state,
                           CompactLatticeArc(0, 0,
                              CompactLatticeWeight(weight, std::vector<int32>()),
                              next_output_state));
        }
      }
    }
  
    LatticePhoneAligner(const CompactLattice &lat,
                        const TransitionModel &tmodel,
                        const PhoneAlignLatticeOptions &opts,
                       CompactLattice *lat_out):
        lat_(lat), tmodel_(tmodel), opts_(opts), lat_out_(lat_out),
        error_(false) {
      fst::CreateSuperFinal(&lat_); // Creates a super-final state, so the
      // only final-probs are One().
    }
  
    // Removes epsilons; also removes unreachable states...
    // not sure if these would exist if original was connected.
    // This also replaces the temporary symbols for the silence
    // and partial-words, with epsilons, if we wanted epsilons.
    void RemoveEpsilonsFromLattice() {
      RmEpsilon(lat_out_, true); // true = connect.
    }
  
    bool AlignLattice() {
      lat_out_->DeleteStates();
      if (lat_.Start() == fst::kNoStateId) {
        KALDI_WARN << "Trying to word-align empty lattice.";
        return false;
      }
      ComputationState initial_comp_state;
      Tuple initial_tuple(lat_.Start(), initial_comp_state);
      StateId start_state = GetStateForTuple(initial_tuple, true); // True = add this to queue.
      lat_out_->SetStart(start_state);
  
      while (!queue_.empty())
        ProcessQueueElement();
  
      if (opts_.remove_epsilon)
        RemoveEpsilonsFromLattice();
  
      return !error_;
    }
  
    CompactLattice lat_;
    const TransitionModel &tmodel_;
    const PhoneAlignLatticeOptions &opts_;
    CompactLattice *lat_out_;
  
    std::vector<std::pair<Tuple, StateId> > queue_;
    MapType map_; // map from tuples to StateId.
    bool error_;
  };
  
  bool LatticePhoneAligner::ComputationState::OutputPhoneArc(
      const TransitionModel &tmodel,
      const PhoneAlignLatticeOptions &opts,
      CompactLatticeArc *arc_out,
      bool *error) {
    if (transition_ids_.empty()) return false;
    int32 phone = tmodel.TransitionIdToPhone(transition_ids_[0]);
    // we assume the start of transition_ids_ is the start of the phone;
    // this is a precondition.
    size_t len = transition_ids_.size(), i;
    // Keep going till we reach a "final" transition-id; note, if
    // reorder==true, we have to go a bit further after this.
    for (i = 0; i < len; i++) {
      int32 tid = transition_ids_[i];
      int32 this_phone = tmodel.TransitionIdToPhone(tid);
      if (this_phone != phone && ! *error) { // error condition: should have
                                             // reached final transition-id first.
        *error = true;
        KALDI_WARN << phone << " -> " << this_phone;
        KALDI_WARN << "Phone changed before final transition-id found "
            "[broken lattice or mismatched model or wrong --reorder option?]";
      }
      if (tmodel.IsFinal(tid))
        break;
    }
    if (i == len) return false; // fell off loop.
    i++; // go past the one for which IsFinal returned true.
    if (opts.reorder) // we have to consume the following self-loop transition-ids.
      while (i < len && tmodel.IsSelfLoop(transition_ids_[i])) i++;
    if (i == len) return false; // we don't know if it ends here... so can't output arc.
  
    // interpret i as the number of transition-ids to consume.
    std::vector<int32> tids_out(transition_ids_.begin(),
                                transition_ids_.begin()+i);
  
    Label output_label = 0;
    if (!word_labels_.empty()) {
      output_label = word_labels_[0];
      word_labels_.erase(word_labels_.begin(), word_labels_.begin()+1);
    }
    if (opts.replace_output_symbols)
      output_label = phone;
    *arc_out = CompactLatticeArc(output_label, output_label,
                                 CompactLatticeWeight(weight_, tids_out),
                                 fst::kNoStateId);
    transition_ids_.erase(transition_ids_.begin(), transition_ids_.begin()+i);
    weight_ = LatticeWeight::One(); // we just output the weight.
    return true;
  }
  
  bool LatticePhoneAligner::ComputationState::OutputWordArc(
      const TransitionModel &tmodel,
      const PhoneAlignLatticeOptions &opts,
      CompactLatticeArc *arc_out,
      bool *error) {
    // output a word but no phones.
    if (word_labels_.size() < 2) return false;
  
    int32 output_label = word_labels_[0];
    word_labels_.erase(word_labels_.begin(), word_labels_.begin()+1);
  
    *arc_out = CompactLatticeArc(output_label, output_label,
                                 CompactLatticeWeight(weight_, std::vector<int32>()),
                                 fst::kNoStateId);
    weight_ = LatticeWeight::One(); // we just output the weight, so set it to one.
    return true;
  }
  
  
  void LatticePhoneAligner::ComputationState::OutputArcForce(
      const TransitionModel &tmodel,
      const PhoneAlignLatticeOptions &opts,
      CompactLatticeArc *arc_out,
      bool *error) {
    KALDI_ASSERT(!IsEmpty());
  
    int32 phone = -1; // This value -1 will never be used,
    // although it might not be obvious from superficially checking
    // the code.  IsEmpty() would be true if we had transition_ids_.empty()
    // and opts.replace_output_symbols, so we would already die by assertion;
    // in fact, this function would neve be called.
  
    if (!transition_ids_.empty()) { // Do some checking here.
      int32 tid = transition_ids_[0];
      phone = tmodel.TransitionIdToPhone(tid);
      int32 num_final = 0;
      for (int32 i = 0; i < transition_ids_.size(); i++) { // A check.
        int32 this_tid = transition_ids_[i];
        int32 this_phone = tmodel.TransitionIdToPhone(this_tid);
        bool is_final = tmodel.IsFinal(this_tid); // should be exactly one.
        if (is_final) num_final++;
        if (this_phone != phone && ! *error) {
          KALDI_WARN << "Mismatch in phone: error in lattice or mismatched transition model?";
          *error = true;
        }
      }
      if (num_final != 1 && ! *error) {
        KALDI_WARN << "Problem phone-aligning lattice: saw " << num_final
                   << " final-states in last phone in lattice (forced out?) "
                   << "Producing partial lattice.";
        *error = true;
      }
    }
  
    Label output_label = 0;
    if (!word_labels_.empty()) {
      output_label = word_labels_[0];
      word_labels_.erase(word_labels_.begin(), word_labels_.begin()+1);
    }
    if (opts.replace_output_symbols)
      output_label = phone;
    *arc_out = CompactLatticeArc(output_label, output_label,
                                 CompactLatticeWeight(weight_, transition_ids_),
                                 fst::kNoStateId);
    transition_ids_.clear();
    weight_ = LatticeWeight::One(); // we just output the weight.
  }
  
  bool PhoneAlignLattice(const CompactLattice &lat,
                         const TransitionModel &tmodel,
                         const PhoneAlignLatticeOptions &opts,
                         CompactLattice *lat_out) {
    LatticePhoneAligner aligner(lat, tmodel, opts, lat_out);
    return aligner.AlignLattice();
  }
  
  
  }  // namespace kaldi