// lat/sausages.cc

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

#include "lat/sausages.h"
#include "lat/lattice-functions.h"

namespace kaldi {

// this is Figure 6 in the paper.
void MinimumBayesRisk::MbrDecode() {

  for (size_t counter = 0; ; counter++) {
    NormalizeEps(&R_);
    AccStats(); // writes to gamma_
    double delta_Q = 0.0; // change in objective function.

    one_best_times_.clear();
    one_best_confidences_.clear();

    // Caution: q in the line below is (q-1) in the algorithm
    // in the paper; both R_ and gamma_ are indexed by q-1.
    for (size_t q = 0; q < R_.size(); q++) {
      if (opts_.decode_mbr) { // This loop updates R_ [indexed same as gamma_].
        // gamma_[i] is sorted in reverse order so most likely one is first.
        const std::vector<std::pair<int32, BaseFloat> > &this_gamma = gamma_[q];
        double old_gamma = 0, new_gamma = this_gamma[0].second;
        int32 rq = R_[q], rhat = this_gamma[0].first; // rq: old word, rhat: new.
        for (size_t j = 0; j < this_gamma.size(); j++)
          if (this_gamma[j].first == rq) old_gamma = this_gamma[j].second;
        delta_Q += (old_gamma - new_gamma); // will be 0 or negative; a bound on
        // change in error.
        if (rq != rhat)
          KALDI_VLOG(2) << "Changing word " << rq << " to " << rhat;
        R_[q] = rhat;
      }
      // build the outputs (time, confidences),
      if (R_[q] != 0 || opts_.print_silence) {
        // see which 'item' from the sausage-bin should we select,
        // (not necessarily the 1st one when MBR decoding disabled)
        int32 s = 0;
        for (int32 j=0; j<gamma_[q].size(); j++) {
          if (gamma_[q][j].first == R_[q]) {
            s = j;
            break;
          }
        }
        one_best_times_.push_back(times_[q][s]);
        // post-process the times,
        size_t i = one_best_times_.size();
        if (i > 1 && one_best_times_[i-2].second > one_best_times_[i-1].first) {
          // It's quite possible for this to happen, but it seems like it would
          // have a bad effect on the downstream processing, so we fix it here.
          // We resolve overlaps by redistributing the available time interval.
          BaseFloat prev_right = i > 2 ? one_best_times_[i-3].second : 0.0;
          BaseFloat left = std::max(prev_right,
                                    std::min(one_best_times_[i-2].first,
                                             one_best_times_[i-1].first));
          BaseFloat right = std::max(one_best_times_[i-2].second,
                                     one_best_times_[i-1].second);
          BaseFloat first_dur =
              one_best_times_[i-2].second - one_best_times_[i-2].first;
          BaseFloat second_dur =
              one_best_times_[i-1].second - one_best_times_[i-1].first;
          BaseFloat mid = first_dur > 0 ? left + (right - left) * first_dur /
                                     (first_dur + second_dur) : left;
          one_best_times_[i-2].first = left;
          one_best_times_[i-2].second = one_best_times_[i-1].first = mid;
          one_best_times_[i-1].second = right;
        }
        BaseFloat confidence = 0.0;
        for (int32 j = 0; j < gamma_[q].size(); j++) {
          if (gamma_[q][j].first == R_[q]) {
            confidence = gamma_[q][j].second;
            break;
          }
        }
        one_best_confidences_.push_back(confidence);
      }
    }
    KALDI_VLOG(2) << "Iter = " << counter << ", delta-Q = " << delta_Q;
    if (delta_Q == 0) break;
    if (counter > 100) {
      KALDI_WARN << "Iterating too many times in MbrDecode; stopping.";
      break;
    }
  }
  if (!opts_.print_silence) RemoveEps(&R_);
}

struct Int32IsZero {
  bool operator() (int32 i) { return (i == 0); }
};
// static
void MinimumBayesRisk::RemoveEps(std::vector<int32> *vec) {
  Int32IsZero pred;
  vec->erase(std::remove_if (vec->begin(), vec->end(), pred),
             vec->end());
}

// static
void MinimumBayesRisk::NormalizeEps(std::vector<int32> *vec) {
  RemoveEps(vec);
  vec->resize(1 + vec->size() * 2);
  int32 s = vec->size();
  for (int32 i = s/2 - 1; i >= 0; i--) {
    (*vec)[i*2 + 1] = (*vec)[i];
    (*vec)[i*2 + 2] = 0;
  }
  (*vec)[0] = 0;
}

double MinimumBayesRisk::EditDistance(int32 N, int32 Q,
                                      Vector<double> &alpha,
                                      Matrix<double> &alpha_dash,
                                      Vector<double> &alpha_dash_arc) {
  alpha(1) = 0.0; // = log(1).  Line 5.
  alpha_dash(1, 0) = 0.0; // Line 5.
  for (int32 q = 1; q <= Q; q++)
    alpha_dash(1, q) = alpha_dash(1, q-1) + l(0, r(q)); // Line 7.
  for (int32 n = 2; n <= N; n++) {
    double alpha_n = kLogZeroDouble;
    for (size_t i = 0; i < pre_[n].size(); i++) {
      const Arc &arc = arcs_[pre_[n][i]];
      alpha_n = LogAdd(alpha_n, alpha(arc.start_node) + arc.loglike);
    }
    alpha(n) = alpha_n; // Line 10.
    // Line 11 omitted: matrix was initialized to zero.
    for (size_t i = 0; i < pre_[n].size(); i++) {
      const Arc &arc = arcs_[pre_[n][i]];
      int32 s_a = arc.start_node, w_a = arc.word;
      BaseFloat p_a = arc.loglike;
      for (int32 q = 0; q <= Q; q++) {
        if (q == 0) {
          alpha_dash_arc(q) = // line 15.
              alpha_dash(s_a, q) + l(w_a, 0, true);
        } else {  // a1,a2,a3 are the 3 parts of min expression of line 17.
          int32 r_q = r(q);
          double a1 = alpha_dash(s_a, q-1) + l(w_a, r_q),
              a2 = alpha_dash(s_a, q) + l(w_a, 0, true),
              a3 = alpha_dash_arc(q-1) + l(0, r_q);
          alpha_dash_arc(q) = std::min(a1, std::min(a2, a3));
        }
        // line 19:
        alpha_dash(n, q) += Exp(alpha(s_a) + p_a - alpha(n)) * alpha_dash_arc(q);
      }
    }
  }
  return alpha_dash(N, Q); // line 23.
}

// Figure 5 in the paper.
void MinimumBayesRisk::AccStats() {
  using std::map;

  int32 N = static_cast<int32>(pre_.size()) - 1,
      Q = static_cast<int32>(R_.size());

  Vector<double> alpha(N+1); // index (1...N)
  Matrix<double> alpha_dash(N+1, Q+1); // index (1...N, 0...Q)
  Vector<double> alpha_dash_arc(Q+1); // index 0...Q
  Matrix<double> beta_dash(N+1, Q+1); // index (1...N, 0...Q)
  Vector<double> beta_dash_arc(Q+1); // index 0...Q
  std::vector<char> b_arc(Q+1); // integer in {1,2,3}; index 1...Q
  std::vector<map<int32, double> > gamma(Q+1); // temp. form of gamma.
  // index 1...Q [word] -> occ.

  // The tau maps below are the sums over arcs with the same word label
  // of the tau_b and tau_e timing quantities mentioned in Appendix C of
  // the paper... we are using these to get averaged times for both the
  // the sausage bins and the 1-best output.
  std::vector<map<int32, double> > tau_b(Q+1), tau_e(Q+1);

  double Ltmp = EditDistance(N, Q, alpha, alpha_dash, alpha_dash_arc);
  if (L_ != 0 && Ltmp > L_) { // L_ != 0 is to rule out 1st iter.
    KALDI_WARN << "Edit distance increased: " << Ltmp << " > "
               << L_;
  }
  L_ = Ltmp;
  KALDI_VLOG(2) << "L = " << L_;
  // omit line 10: zero when initialized.
  beta_dash(N, Q) = 1.0; // Line 11.
  for (int32 n = N; n >= 2; n--) {
    for (size_t i = 0; i < pre_[n].size(); i++) {
      const Arc &arc = arcs_[pre_[n][i]];
      int32 s_a = arc.start_node, w_a = arc.word;
      BaseFloat p_a = arc.loglike;
      alpha_dash_arc(0) = alpha_dash(s_a, 0) + l(w_a, 0, true); // line 14.
      for (int32 q = 1; q <= Q; q++) { // this loop == lines 15-18.
        int32 r_q = r(q);
        double a1 = alpha_dash(s_a, q-1) + l(w_a, r_q),
            a2 = alpha_dash(s_a, q) + l(w_a, 0, true),
            a3 = alpha_dash_arc(q-1) + l(0, r_q);
        if (a1 <= a2) {
          if (a1 <= a3) { b_arc[q] = 1; alpha_dash_arc(q) = a1; }
          else { b_arc[q] = 3; alpha_dash_arc(q) = a3; }
        } else {
          if (a2 <= a3) { b_arc[q] = 2; alpha_dash_arc(q) = a2; }
          else { b_arc[q] = 3; alpha_dash_arc(q) = a3; }
        }
      }
      beta_dash_arc.SetZero(); // line 19.
      for (int32 q = Q; q >= 1; q--) {
        // line 21:
        beta_dash_arc(q) += Exp(alpha(s_a) + p_a - alpha(n)) * beta_dash(n, q);
        switch (static_cast<int>(b_arc[q])) { // lines 22 and 23:
          case 1:
            beta_dash(s_a, q-1) += beta_dash_arc(q);
            // next: gamma(q, w(a)) += beta_dash_arc(q)
            AddToMap(w_a, beta_dash_arc(q), &(gamma[q]));
            // next: accumulating times, see decl for tau_b,tau_e
            AddToMap(w_a, state_times_[s_a] * beta_dash_arc(q), &(tau_b[q]));
            AddToMap(w_a, state_times_[n] * beta_dash_arc(q), &(tau_e[q]));
            break;
          case 2:
            beta_dash(s_a, q) += beta_dash_arc(q);
            break;
          case 3:
            beta_dash_arc(q-1) += beta_dash_arc(q);
            // next: gamma(q, epsilon) += beta_dash_arc(q)
            AddToMap(0, beta_dash_arc(q), &(gamma[q]));
            // next: accumulating times, see decl for tau_b,tau_e
            // WARNING: there was an error in Appendix C.  If we followed
            // the instructions there the next line would say state_times_[sa], but
            // it would be wrong.  I will try to publish an erratum.
            AddToMap(0, state_times_[n] * beta_dash_arc(q), &(tau_b[q]));
            AddToMap(0, state_times_[n] * beta_dash_arc(q), &(tau_e[q]));
            break;
          default:
            KALDI_ERR << "Invalid b_arc value"; // error in code.
        }
      }
      beta_dash_arc(0) += Exp(alpha(s_a) + p_a - alpha(n)) * beta_dash(n, 0);
      beta_dash(s_a, 0) += beta_dash_arc(0); // line 26.
    }
  }
  beta_dash_arc.SetZero(); // line 29.
  for (int32 q = Q; q >= 1; q--) {
    beta_dash_arc(q) += beta_dash(1, q);
    beta_dash_arc(q-1) += beta_dash_arc(q);
    AddToMap(0, beta_dash_arc(q), &(gamma[q]));
    // the statements below are actually redundant because
    // state_times_[1] is zero.
    AddToMap(0, state_times_[1] * beta_dash_arc(q), &(tau_b[q]));
    AddToMap(0, state_times_[1] * beta_dash_arc(q), &(tau_e[q]));
  }
  for (int32 q = 1; q <= Q; q++) { // a check (line 35)
    double sum = 0.0;
    for (map<int32, double>::iterator iter = gamma[q].begin();
         iter != gamma[q].end(); ++iter) sum += iter->second;
    if (fabs(sum - 1.0) > 0.1)
      KALDI_WARN << "sum of gamma[" << q << ",s] is " << sum;
  }
  // The next part is where we take gamma, and convert
  // to the class member gamma_, which is using a different
  // data structure and indexed from zero, not one.
  gamma_.clear();
  gamma_.resize(Q);
  for (int32 q = 1; q <= Q; q++) {
    for (map<int32, double>::iterator iter = gamma[q].begin();
         iter != gamma[q].end(); ++iter)
      gamma_[q-1].push_back(
          std::make_pair(iter->first, static_cast<BaseFloat>(iter->second)));
    // sort gamma_[q-1] from largest to smallest posterior.
    GammaCompare comp;
    std::sort(gamma_[q-1].begin(), gamma_[q-1].end(), comp);
  }
  // We do the same conversion for the state times tau_b and tau_e:
  // they get turned into the times_ data member, which has zero-based
  // indexing.
  times_.clear();
  times_.resize(Q);
  sausage_times_.clear();
  sausage_times_.resize(Q);
  for (int32 q = 1; q <= Q; q++) {
    double t_b = 0.0, t_e = 0.0;
    for (std::vector<std::pair<int32, BaseFloat>>::iterator iter = gamma_[q-1].begin();
         iter != gamma_[q-1].end(); ++iter) {
      double w_b = tau_b[q][iter->first], w_e = tau_e[q][iter->first];
      if (w_b > w_e)
        KALDI_WARN << "Times out of order";  // this is quite bad.
      times_[q-1].push_back(
          std::make_pair(static_cast<BaseFloat>(w_b / iter->second),
                         static_cast<BaseFloat>(w_e / iter->second)));
      t_b += w_b;
      t_e += w_e;
    }
    sausage_times_[q-1].first = t_b;
    sausage_times_[q-1].second = t_e;
    if (sausage_times_[q-1].first > sausage_times_[q-1].second)
      KALDI_WARN << "Times out of order";  // this is quite bad.
    if (q > 1 && sausage_times_[q-2].second > sausage_times_[q-1].first) {
      // We previously had a warning here, but now we'll just set both
      // those values to their average.  It's quite possible for this
      // condition to happen, but it seems like it would have a bad effect
      // on the downstream processing, so we fix it.
      sausage_times_[q-2].second = sausage_times_[q-1].first =
          0.5 * (sausage_times_[q-2].second + sausage_times_[q-1].first);
    }
  }
}

void MinimumBayesRisk::PrepareLatticeAndInitStats(CompactLattice *clat) {
  KALDI_ASSERT(clat != NULL);

  CreateSuperFinal(clat); // Add super-final state to clat... this is
  // one of the requirements of the MBR algorithm, as mentioned in the
  // paper (i.e. just one final state).

  // Topologically sort the lattice, if not already sorted.
  kaldi::uint64 props = clat->Properties(fst::kFstProperties, false);
  if (!(props & fst::kTopSorted)) {
    if (fst::TopSort(clat) == false)
      KALDI_ERR << "Cycles detected in lattice.";
  }
  CompactLatticeStateTimes(*clat, &state_times_); // work out times of
  // the states in clat
  state_times_.push_back(0); // we'll convert to 1-based numbering.
  for (size_t i = state_times_.size()-1; i > 0; i--)
    state_times_[i] = state_times_[i-1];

  // Now we convert the information in "clat" into a special internal
  // format (pre_, post_ and arcs_) which allows us to access the
  // arcs preceding any given state.
  // Note: in our internal format the states will be numbered from 1,
  // which involves adding 1 to the OpenFst states.
  int32 N = clat->NumStates();
  pre_.resize(N+1);

  // Careful: "Arc" is a class-member struct, not an OpenFst type of arc as one
  // would normally assume.
  for (int32 n = 1; n <= N; n++) {
    for (fst::ArcIterator<CompactLattice> aiter(*clat, n-1);
         !aiter.Done();
         aiter.Next()) {
      const CompactLatticeArc &carc = aiter.Value();
      Arc arc; // in our local format.
      arc.word = carc.ilabel; // == carc.olabel
      arc.start_node = n;
      arc.end_node = carc.nextstate + 1; // convert to 1-based.
      arc.loglike = - (carc.weight.Weight().Value1() +
                       carc.weight.Weight().Value2());
      // loglike: sum graph/LM and acoustic cost, and negate to
      // convert to loglikes.  We assume acoustic scaling is already done.

      pre_[arc.end_node].push_back(arcs_.size()); // record index of this arc.
      arcs_.push_back(arc);
    }
  }
}

MinimumBayesRisk::MinimumBayesRisk(const CompactLattice &clat_in,
                                   MinimumBayesRiskOptions opts) : opts_(opts) {
  CompactLattice clat(clat_in); // copy.

  PrepareLatticeAndInitStats(&clat);

  // We don't need to look at clat.Start() or clat.Final(state):
  // we know clat.Start() == 0 since it's topologically sorted,
  // and clat.Final(state) is Zero() except for One() at the last-
  // numbered state, thanks to CreateSuperFinal and the topological
  // sorting.

  { // Now set R_ to one best in the FST.
    RemoveAlignmentsFromCompactLattice(&clat); // will be more efficient
    // in best-path if we do this.
    Lattice lat;
    ConvertLattice(clat, &lat); // convert from CompactLattice to Lattice.
    fst::VectorFst<fst::StdArc> fst;
    ConvertLattice(lat, &fst); // convert from lattice to normal FST.
    fst::VectorFst<fst::StdArc> fst_shortest_path;
    fst::ShortestPath(fst, &fst_shortest_path); // take shortest path of FST.
    std::vector<int32> alignment, words;
    fst::TropicalWeight weight;
    GetLinearSymbolSequence(fst_shortest_path, &alignment, &words, &weight);
    KALDI_ASSERT(alignment.empty()); // we removed the alignment.
    R_ = words;
    L_ = 0.0; // Set current edit-distance to 0 [just so we know
    // when we're on the 1st iter.]
  }

  MbrDecode();

}

MinimumBayesRisk::MinimumBayesRisk(const CompactLattice &clat_in,
                                   const std::vector<int32> &words,
                                   MinimumBayesRiskOptions opts) : opts_(opts) {
  CompactLattice clat(clat_in); // copy.

  PrepareLatticeAndInitStats(&clat);

  R_ = words;
  L_ = 0.0;

  MbrDecode();
}

MinimumBayesRisk::MinimumBayesRisk(const CompactLattice &clat_in,
                                   const std::vector<int32> &words,
                                   const std::vector<std::pair<BaseFloat,BaseFloat> > &times,
                                   MinimumBayesRiskOptions opts) : opts_(opts) {
  CompactLattice clat(clat_in); // copy.

  PrepareLatticeAndInitStats(&clat);

  R_ = words;
  sausage_times_ = times;
  L_ = 0.0;

  MbrDecode();
}


}  // namespace kaldi