// lat/sausages.h

// Copyright 2012  Johns Hopkins University (Author: Daniel Povey)
//           2015  Guoguo Chen
//           2019  Dogan Can

// 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.


#ifndef KALDI_LAT_SAUSAGES_H_
#define KALDI_LAT_SAUSAGES_H_

#include <vector>
#include <map>

#include "base/kaldi-common.h"
#include "util/common-utils.h"
#include "fstext/fstext-lib.h"
#include "lat/kaldi-lattice.h"

namespace kaldi {

/// The implementation of the Minimum Bayes Risk decoding method described in
///  "Minimum Bayes Risk decoding and system combination based on a recursion for
///  edit distance", Haihua Xu, Daniel Povey, Lidia Mangu and Jie Zhu, Computer
///  Speech and Language, 2011
/// This is a slightly more principled way to do Minimum Bayes Risk (MBR) decoding
/// than the standard "Confusion Network" method.  Note: MBR decoding aims to
/// minimize the expected word error rate, assuming the lattice encodes the
/// true uncertainty about what was spoken; standard Viterbi decoding gives the
/// most likely utterance, which corresponds to minimizing the expected sentence
/// error rate.
///
/// In addition to giving the MBR output, we also provide a way to get a
/// "Confusion Network" or informally "sausage"-like structure.  This is a
/// linear sequence of bins, and in each bin, there is a distribution over
/// words (or epsilon, meaning no word).  This is useful for estimating
/// confidence.  Note: due to the way these sausages are made, typically there
/// will be, between each bin representing a high-confidence word, a bin
/// in which epsilon (no word) is the most likely word.  Inside these bins
/// is where we put possible insertions.

struct MinimumBayesRiskOptions {
  /// Boolean configuration parameter: if true, we actually update the hypothesis
  /// to do MBR decoding (if false, our output is the MAP decoded output, but we
  /// output the stats too (i.e. the confidences)).
  bool decode_mbr;
  /// Boolean configuration parameter: if true, the 1-best path will 'keep' the <eps> bins,
  bool print_silence;

  MinimumBayesRiskOptions() : decode_mbr(true), print_silence(false)
  { }
  void Register(OptionsItf *opts) {
    opts->Register("decode-mbr", &decode_mbr, "If true, do Minimum Bayes Risk "
                   "decoding (else, Maximum a Posteriori)");
    opts->Register("print-silence", &print_silence, "Keep the inter-word '<eps>' "
                   "bins in the 1-best output (ctm, <eps> can be a 'silence' or a 'deleted' word)");
  }
};

/// This class does the word-level Minimum Bayes Risk computation, and gives you
/// either the 1-best MBR output together with the expected Bayes Risk,
/// or a sausage-like structure.
class MinimumBayesRisk {
 public:
  /// Initialize with compact lattice-- any acoustic scaling etc., is assumed
  /// to have been done already.
  /// This does the whole computation.  You get the output with
  /// GetOneBest(), GetBayesRisk(), and GetSausageStats().
  MinimumBayesRisk(const CompactLattice &clat,
                   MinimumBayesRiskOptions opts = MinimumBayesRiskOptions());

  // Uses the provided <words> as <R_> instead of using the lattice best path.
  // Note that the default value of opts.decode_mbr is true. If you provide 1-best
  // hypothesis from MAP decoding, the output ctm from MBR decoding may be
  // mismatched with the provided <words> (<words> would be used as the starting
  // point of optimization).
  MinimumBayesRisk(const CompactLattice &clat,
                   const std::vector<int32> &words,
                   MinimumBayesRiskOptions opts = MinimumBayesRiskOptions());
  // Uses the provided <words> as <R_> and <times> of bins instead of using the lattice best path.
  // Note that the default value of opts.decode_mbr is true. If you provide 1-best
  // hypothesis from MAP decoding, the output ctm from MBR decoding may be
  // mismatched with the provided <words> (<words> would be used as the starting
  // point of optimization).
  MinimumBayesRisk(const CompactLattice &clat,
                   const std::vector<int32> &words,
                   const std::vector<std::pair<BaseFloat,BaseFloat> > &times,
                   MinimumBayesRiskOptions opts = MinimumBayesRiskOptions());

  const std::vector<int32> &GetOneBest() const { // gets one-best (with no epsilons)
    return R_;
  }

  const std::vector<std::vector<std::pair<BaseFloat, BaseFloat> > > GetTimes() const {
    return times_; // returns average (start,end) times for each word in each
    // bin. These are raw averages without any processing, i.e. time intervals
    // from different bins can overlap.
  }

  const std::vector<std::pair<BaseFloat, BaseFloat> > GetSausageTimes() const {
    return sausage_times_; // returns average (start,end) times for each bin.
    // This is typically the weighted average of the times in GetTimes() but can
    // be slightly different if the times for the bins overlap, in which case
    // the times returned by this method do not overlap unlike the times
    // returned by GetTimes().
  }

  const std::vector<std::pair<BaseFloat, BaseFloat> > &GetOneBestTimes() const {
    return one_best_times_; // returns average (start,end) times for each word
    // corresponding to an entry in the one-best output.  This is typically the
    // appropriate subset of the times in GetTimes() but can be slightly
    // different if the times for the one-best words overlap, in which case
    // the times returned by this method do not overlap unlike the times
    // returned by GetTimes().
  }

  /// Outputs the confidences for the one-best transcript.
  const std::vector<BaseFloat> &GetOneBestConfidences() const {
    return one_best_confidences_;
  }

  /// Returns the expected WER over this sentence (assuming model correctness).
  BaseFloat GetBayesRisk() const { return L_; }

  const std::vector<std::vector<std::pair<int32, BaseFloat> > > &GetSausageStats() const {
    return gamma_;
  }

 private:
  void PrepareLatticeAndInitStats(CompactLattice *clat);

  /// Minimum-Bayes-Risk Decode. Top-level algorithm.  Figure 6 of the paper.
  void MbrDecode();

  /// Without the 'penalize' argument this gives us the basic edit-distance
  /// function l(a,b), as in the paper.
  /// With the 'penalize' argument it can be interpreted as the edit distance
  /// plus the 'delta' from the paper, except that we make a kind of conceptual
  /// bug-fix and only apply the delta if the edit-distance was not already
  /// zero.  This bug-fix was necessary in order to force all the stats to show
  /// up, that should show up, and applying the bug-fix makes the sausage stats
  /// significantly less sparse.
  inline double l(int32 a, int32 b, bool penalize = false) {
    if (a == b) return 0.0;
    else return (penalize ? 1.0 + delta() : 1.0);
  }

  /// returns r_q, in one-based indexing, as in the paper.
  inline int32 r(int32 q) { return R_[q-1]; }


  /// Figure 4 of the paper; called from AccStats (Fig. 5)
  double EditDistance(int32 N, int32 Q,
                      Vector<double> &alpha,
                      Matrix<double> &alpha_dash,
                      Vector<double> &alpha_dash_arc);

  /// Figure 5 of the paper.  Outputs to gamma_ and L_.
  void AccStats();

  /// Removes epsilons (symbol 0) from a vector
  static void RemoveEps(std::vector<int32> *vec);

  // Ensures that between each word in "vec" and at the beginning and end, is
  // epsilon (0).  (But if no words in vec, just one epsilon)
  static void NormalizeEps(std::vector<int32> *vec);

  // delta() is a constant used in the algorithm, which penalizes
  // the use of certain epsilon transitions in the edit-distance which would cause
  // words not to show up in the accumulated edit-distance statistics.
  // There has been a conceptual bug-fix versus the way it was presented in
  // the paper: we now add delta only if the edit-distance was not already
  // zero.
  static inline BaseFloat delta() { return 1.0e-05; }


  /// Function used to increment map.
  static inline void AddToMap(int32 i, double d, std::map<int32, double> *gamma) {
    if (d == 0) return;
    std::pair<const int32, double> pr(i, d);
    std::pair<std::map<int32, double>::iterator, bool> ret = gamma->insert(pr);
    if (!ret.second) // not inserted, so add to contents.
      ret.first->second += d;
  }

  struct Arc {
    int32 word;
    int32 start_node;
    int32 end_node;
    BaseFloat loglike;
  };

  MinimumBayesRiskOptions opts_;


  /// Arcs in the topologically sorted acceptor form of the word-level lattice,
  /// with one final-state.  Contains (word-symbol, log-likelihood on arc ==
  /// negated cost).  Indexed from zero.
  std::vector<Arc> arcs_;

  /// For each node in the lattice, a list of arcs entering that node. Indexed
  /// from 1 (first node == 1).
  std::vector<std::vector<int32> > pre_;

  std::vector<int32> state_times_; // time of each state in the word lattice,
  // indexed from 1 (same index as into pre_)

  std::vector<int32> R_; // current 1-best word sequence, normalized to have
  // epsilons between each word and at the beginning and end.  R in paper...
  // caution: indexed from zero, not from 1 as in paper.

  double L_; // current averaged edit-distance between lattice and R_.
  // \hat{L} in paper.

  std::vector<std::vector<std::pair<int32, BaseFloat> > > gamma_;
  // The stats we accumulate; these are pairs of (posterior, word-id), and note
  // that word-id may be epsilon.  Caution: indexed from zero, not from 1 as in
  // paper.  We sort in reverse order on the second member (posterior), so more
  // likely word is first.

  std::vector<std::vector<std::pair<BaseFloat, BaseFloat> > > times_;
  // The average start and end times for words in each confusion-network bin.
  // This is like an average over arcs, of the tau_b and tau_e quantities in
  // Appendix C of the paper.  Indexed from zero, like gamma_ and R_.

  std::vector<std::pair<BaseFloat, BaseFloat> > sausage_times_;
  // The average start and end times for each confusion-network bin.  This
  // is like an average over words, of the tau_b and tau_e quantities in
  // Appendix C of the paper.  Indexed from zero, like gamma_ and R_.

  std::vector<std::pair<BaseFloat, BaseFloat> > one_best_times_;
  // The average start and end times for words in the one best output.  This
  // is like an average over the arcs, of the tau_b and tau_e quantities in
  // Appendix C of the paper. Indexed from zero, like gamma_ and R_.

  std::vector<BaseFloat> one_best_confidences_;
  // vector of confidences for the 1-best output (which could be
  // the MAP output if opts_.decode_mbr == false, or the MBR output otherwise).
  // Indexed by the same index as one_best_times_.

  struct GammaCompare{
    // should be like operator <.  But we want reverse order
    // on the 2nd element (posterior), so it'll be like operator
    // > that looks first at the posterior.
    bool operator () (const std::pair<int32, BaseFloat> &a,
                      const std::pair<int32, BaseFloat> &b) const {
      if (a.second > b.second) return true;
      else if (a.second < b.second) return false;
      else return a.first > b.first;
    }
  };
};

}  // namespace kaldi

#endif  // KALDI_LAT_SAUSAGES_H_