// gmm/ebw-diag-gmm.cc
  
  // Copyright 2009-2011  Arnab Ghoshal, Petr Motlicek
  
  // 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 <algorithm>  // for std::max
  #include <string>
  #include <vector>
  
  #include "gmm/diag-gmm.h"
  #include "gmm/ebw-diag-gmm.h"
  
  namespace kaldi {
  
  // This function is used inside the EBW update routines.
  // returns true if all variances were positive.
  static bool EBWUpdateGaussian(
      BaseFloat D,
      GmmFlagsType flags,
      const VectorBase<double> &orig_mean,
      const VectorBase<double> &orig_var,
      const VectorBase<double> &x_stats,
      const VectorBase<double> &x2_stats,
      double occ,
      VectorBase<double> *mean,
      VectorBase<double> *var,
      double *auxf_impr) {
    if (! (flags&(kGmmMeans|kGmmVariances))) { // nothing to do.
      if (auxf_impr) *auxf_impr = 0.0;
      mean->CopyFromVec(orig_mean);
      var->CopyFromVec(orig_var);
      return true; 
    }   
    KALDI_ASSERT(!( (flags&kGmmVariances) && !(flags&kGmmMeans))
                 && "We didn't make the update cover this case sensibly (update vars not means)");
    
    mean->SetZero();
    var->SetZero();
    mean->AddVec(D, orig_mean);
    var->AddVec2(D, orig_mean);
    var->AddVec(D, orig_var);
    mean->AddVec(1.0, x_stats);
    var->AddVec(1.0, x2_stats);
    BaseFloat scale = 1.0 / (occ + D);
    mean->Scale(scale);
    var->Scale(scale);
    var->AddVec2(-1.0, *mean);
    
    if (!(flags&kGmmVariances)) var->CopyFromVec(orig_var);
    if (!(flags&kGmmMeans)) mean->CopyFromVec(orig_mean);
  
    // Return false if any NaN's.
    for (int32 i = 0; i < mean->Dim(); i++) {
      double m =  ((*mean)(i)), v = ((*var)(i));
      if (m!=m || v!=v || m-m != 0 || v-v != 0) {
        return false;
      }
    }
    
    if (var->Min() > 0.0) {
      if (auxf_impr != NULL) {
        // work out auxf improvement.  
        BaseFloat old_auxf = 0.0, new_auxf = 0.0;
        int32 dim = orig_mean.Dim();
        for (int32 i = 0; i < dim; i++) {
          BaseFloat mean_diff = (*mean)(i) - orig_mean(i);
          old_auxf += (occ+D) * -0.5 * (Log(orig_var(i)) +
                                        ((*var)(i) + mean_diff*mean_diff)
                                        / orig_var(i));
          new_auxf += (occ+D) * -0.5 * (Log((*var)(i)) + 1.0);
          
        }
        *auxf_impr = new_auxf - old_auxf;
      }
      return true;
    } else return false;
  }
  
  // Update Gaussian parameters only (no weights)
  void UpdateEbwDiagGmm(const AccumDiagGmm &num_stats, // with I-smoothing, if used.
                        const AccumDiagGmm &den_stats,
                        GmmFlagsType flags,
                        const EbwOptions &opts,
                        DiagGmm *gmm,
                        BaseFloat *auxf_change_out,
                        BaseFloat *count_out,
                        int32 *num_floored_out) {
    GmmFlagsType acc_flags = num_stats.Flags();
    if (flags & ~acc_flags)
      KALDI_ERR << "Incompatible flags: you requested to update flags \""
                << GmmFlagsToString(flags) << "\" but accumulators have only \""
                << GmmFlagsToString(acc_flags) << '"';
    
    // It could be that the num stats actually contain the difference between
    // num and den (for mean and var stats), and den stats only have the weights.
    bool den_has_stats;
    if (den_stats.Flags() != acc_flags) {
      den_has_stats = false;
      if (den_stats.Flags() != kGmmWeights) 
        KALDI_ERR << "Incompatible flags: num stats have flags \""
                  << GmmFlagsToString(acc_flags) << "\" vs. den stats \""
                  << GmmFlagsToString(den_stats.Flags()) << '"';
    } else {
      den_has_stats = true;
    }
    int32 num_comp = num_stats.NumGauss();
    int32 dim = num_stats.Dim();
    KALDI_ASSERT(num_stats.NumGauss() == den_stats.NumGauss());
    KALDI_ASSERT(num_stats.Dim() == gmm->Dim());
    KALDI_ASSERT(gmm->NumGauss() == num_comp);
    
    if ( !(flags & (kGmmMeans | kGmmVariances)) ) {
      return; // Nothing to update.
    }
    
    // copy DiagGMM model and transform this to the normal case
    DiagGmmNormal diaggmmnormal;
    gmm->ComputeGconsts();
    diaggmmnormal.CopyFromDiagGmm(*gmm);
  
    // go over all components
    Vector<double> mean(dim), var(dim), mean_stats(dim), var_stats(dim);
  
    for (int32 g = 0; g < num_comp; g++) {
      BaseFloat num_count = num_stats.occupancy()(g),
          den_count = den_stats.occupancy()(g);
      if (num_count == 0.0 && den_count == 0.0) {
        KALDI_VLOG(2) << "Not updating Gaussian " << g << " since counts are zero";
        continue;
      }
      mean_stats.CopyFromVec(num_stats.mean_accumulator().Row(g));
      if (den_has_stats)
        mean_stats.AddVec(-1.0, den_stats.mean_accumulator().Row(g));
      if (flags & kGmmVariances) {
        var_stats.CopyFromVec(num_stats.variance_accumulator().Row(g));
        if (den_has_stats)
          var_stats.AddVec(-1.0, den_stats.variance_accumulator().Row(g));
      }
      double D = (opts.tau + opts.E * den_count) / 2;
      if (D+num_count-den_count <= 0.0) {
        // ensure +ve-- can be problem if num count == 0 and E=2.
        D = -1.0001*(num_count-den_count) + 1.0e-10;
        KALDI_ASSERT(D+num_count-den_count > 0.0);
      }
      // We initialize to half the value of D that would be dictated by E (and
      // tau); this is part of the strategy used to ensure that the value of D we
      // use is at least twice the value that would ensure positive variances.
  
      int32 iter, max_iter = 100;
      for (iter = 0; iter < max_iter; iter++) { // will normally break from the loop
        // the first time.
        if (EBWUpdateGaussian(D, flags,
                              diaggmmnormal.means_.Row(g),
                              diaggmmnormal.vars_.Row(g),
                              mean_stats, var_stats, num_count-den_count,
                              &mean, &var, NULL)) {
          // Succeeded in getting all +ve vars at this value of D.
          // So double D and commit changes.
          D *= 2.0;
          double auxf_impr = 0.0;
          bool ans = EBWUpdateGaussian(D, flags,
                                       diaggmmnormal.means_.Row(g),
                                       diaggmmnormal.vars_.Row(g),
                                       mean_stats, var_stats, num_count-den_count,
                                       &mean, &var, &auxf_impr);
          if (!ans) {
            KALDI_WARN << "Something went wrong in the EBW update. Check that your"
                "previous update phase looks reasonable, probably your model is "
                "already ruined.  Reverting to the old values";
          } else {
            if (auxf_change_out) *auxf_change_out += auxf_impr;
            if (count_out) *count_out += den_count; // The idea is that for MMI, this will
            // reflect the actual #frames trained on (the numerator one would be I-smoothed).
            // In general (e.g. for MPE), we won't know the #frames.
            diaggmmnormal.means_.CopyRowFromVec(mean, g);
            diaggmmnormal.vars_.CopyRowFromVec(var, g);
          }
          break;
        } else {
          // small step
          D *= 1.1; 
        }
      }
      if (iter > 0 && num_floored_out != NULL) (*num_floored_out)++;
      if (iter == max_iter) KALDI_WARN << "Dropped off end of loop, recomputing D. (unexpected.)";
    }
    // copy to natural representation according to flags.
    diaggmmnormal.CopyToDiagGmm(gmm, flags);
    gmm->ComputeGconsts();
  }
  
  
  void UpdateEbwWeightsDiagGmm(const AccumDiagGmm &num_stats, // should have no I-smoothing
                               const AccumDiagGmm &den_stats,
                               const EbwWeightOptions &opts,
                               DiagGmm *gmm,
                               BaseFloat *auxf_change_out,
                               BaseFloat *count_out) {
  
    DiagGmmNormal diaggmmnormal;
    gmm->ComputeGconsts();
    diaggmmnormal.CopyFromDiagGmm(*gmm);
  
    Vector<double> weights(diaggmmnormal.weights_),
        num_occs(num_stats.occupancy()),
        den_occs(den_stats.occupancy());
    if (opts.tau == 0.0 &&
        num_occs.Sum() + den_occs.Sum() < opts.min_num_count_weight_update) {
      KALDI_LOG << "Not updating weights for this state because total count is "
                << num_occs.Sum() + den_occs.Sum() << " < "
                << opts.min_num_count_weight_update;
      if (count_out)
        *count_out += num_occs.Sum();
      return;
    }
    num_occs.AddVec(opts.tau, weights);
    KALDI_ASSERT(weights.Dim() == num_occs.Dim() && num_occs.Dim() == den_occs.Dim());
    if (weights.Dim() == 1) return; // Nothing to do: only one mixture.
    double weight_auxf_at_start = 0.0, weight_auxf_at_end = 0.0;
    
    int32 num_comp = weights.Dim();
    for (int32 g = 0; g < num_comp; g++) {   // c.f. eq. 4.32 in Dan Povey's thesis.
      weight_auxf_at_start +=
          num_occs(g) * log (weights(g))
          - den_occs(g) * weights(g) / diaggmmnormal.weights_(g);
    }
    for (int32 iter = 0; iter < 50; iter++) {
      Vector<double> k_jm(num_comp); // c.f. eq. 4.35
      double max_m = 0.0;
      for (int32 g = 0; g < num_comp; g++)
        max_m = std::max(max_m, den_occs(g)/diaggmmnormal.weights_(g));
      for (int32 g = 0; g < num_comp; g++)
        k_jm(g) = max_m - den_occs(g)/diaggmmnormal.weights_(g);
      for (int32 g = 0; g < num_comp; g++) // c.f. eq. 4.34
        weights(g) = num_occs(g) + k_jm(g)*weights(g);
      weights.Scale(1.0 / weights.Sum()); // c.f. eq. 4.34 (denominator)
    }
    for (int32 g = 0; g < num_comp; g++) {   // weight flooring.
      if (weights(g) < opts.min_gaussian_weight)
        weights(g) = opts.min_gaussian_weight;
    }
    weights.Scale(1.0 / weights.Sum()); // renormalize after flooring..
    // floor won't be exact now but doesn't really matter.
  
    for (int32 g = 0; g < num_comp; g++) {   // c.f. eq. 4.32 in Dan Povey's thesis.
      weight_auxf_at_end +=
          num_occs(g) * log (weights(g))
          - den_occs(g) * weights(g) / diaggmmnormal.weights_(g);
    }
  
    if (auxf_change_out)
      *auxf_change_out += weight_auxf_at_end - weight_auxf_at_start;
    if (count_out)
      *count_out += num_occs.Sum(); // only really valid for MMI [not MPE, or MMI
    // with canceled stats]
    
    diaggmmnormal.weights_.CopyFromVec(weights);
  
    // copy to natural representation
    diaggmmnormal.CopyToDiagGmm(gmm, kGmmAll);
    gmm->ComputeGconsts();
  }
  
  void UpdateEbwAmDiagGmm(const AccumAmDiagGmm &num_stats, // with I-smoothing, if used.
                          const AccumAmDiagGmm &den_stats,
                          GmmFlagsType flags,
                          const EbwOptions &opts,
                          AmDiagGmm *am_gmm,
                          BaseFloat *auxf_change_out,
                          BaseFloat *count_out,
                          int32 *num_floored_out) {
    KALDI_ASSERT(num_stats.NumAccs() == den_stats.NumAccs()
                 && num_stats.NumAccs() == am_gmm->NumPdfs());
  
    if (auxf_change_out) *auxf_change_out = 0.0;
    if (count_out) *count_out = 0.0;
    if (num_floored_out) *num_floored_out = 0.0;
  
    for (int32 pdf = 0; pdf < num_stats.NumAccs(); pdf++)
      UpdateEbwDiagGmm(num_stats.GetAcc(pdf), den_stats.GetAcc(pdf), flags,
                       opts, &(am_gmm->GetPdf(pdf)), auxf_change_out,
                       count_out, num_floored_out);
  }                     
  
  
  void UpdateEbwWeightsAmDiagGmm(const AccumAmDiagGmm &num_stats, // with I-smoothing, if used.
                                 const AccumAmDiagGmm &den_stats,
                                 const EbwWeightOptions &opts,
                                 AmDiagGmm *am_gmm,
                                 BaseFloat *auxf_change_out,
                                 BaseFloat *count_out) {
    KALDI_ASSERT(num_stats.NumAccs() == den_stats.NumAccs()
                 && num_stats.NumAccs() == am_gmm->NumPdfs());
  
    if (auxf_change_out) *auxf_change_out = 0.0;
    if (count_out) *count_out = 0.0;
    
    for (int32 pdf = 0; pdf < num_stats.NumAccs(); pdf++)
      UpdateEbwWeightsDiagGmm(num_stats.GetAcc(pdf), den_stats.GetAcc(pdf),
                              opts, &(am_gmm->GetPdf(pdf)), auxf_change_out,
                              count_out);
  }                     
  
  void IsmoothStatsDiagGmm(const AccumDiagGmm &src_stats,
                           double tau,
                           AccumDiagGmm *dst_stats) {
    KALDI_ASSERT(src_stats.NumGauss() == dst_stats->NumGauss());
    int32 dim = src_stats.Dim(), num_gauss = src_stats.NumGauss();
    for (int32 g = 0; g < num_gauss; g++) {
      double occ = src_stats.occupancy()(g);
      if (occ != 0.0) { // can only do this for nonzero occupancies...
        Vector<double> x_stats(dim), x2_stats(dim);
        if (dst_stats->Flags() & kGmmMeans)
          x_stats.CopyFromVec(src_stats.mean_accumulator().Row(g));
        if (dst_stats->Flags() & kGmmVariances)
          x2_stats.CopyFromVec(src_stats.variance_accumulator().Row(g));
        x_stats.Scale(tau / occ);
        x2_stats.Scale(tau / occ);
        dst_stats->AddStatsForComponent(g, tau, x_stats, x2_stats);
      }
    }
  }
  
  /// Creates stats from the GMM.  Resizes them as needed.
  void DiagGmmToStats(const DiagGmm &gmm,
                      GmmFlagsType flags,
                      double state_occ,
                      AccumDiagGmm *dst_stats) {
    dst_stats->Resize(gmm, AugmentGmmFlags(flags));
    int32 num_gauss = gmm.NumGauss(), dim = gmm.Dim();
    DiagGmmNormal gmmnormal(gmm);
    Vector<double> x_stats(dim), x2_stats(dim);
    for (int32 g = 0; g < num_gauss; g++) {
      double occ = state_occ * gmmnormal.weights_(g);
      x_stats.SetZero();
      x_stats.AddVec(occ, gmmnormal.means_.Row(g));
      x2_stats.SetZero();
      x2_stats.AddVec2(occ, gmmnormal.means_.Row(g));
      x2_stats.AddVec(occ, gmmnormal.vars_.Row(g));
      dst_stats->AddStatsForComponent(g, occ, x_stats, x2_stats);
    }
  }
  
  void IsmoothStatsAmDiagGmm(const AccumAmDiagGmm &src_stats,
                             double tau,
                             AccumAmDiagGmm *dst_stats) {
    int num_pdfs = src_stats.NumAccs();
    KALDI_ASSERT(num_pdfs == dst_stats->NumAccs());
    for (int32 pdf = 0; pdf < num_pdfs; pdf++)
      IsmoothStatsDiagGmm(src_stats.GetAcc(pdf), tau, &(dst_stats->GetAcc(pdf)));
  }
  
  void IsmoothStatsAmDiagGmmFromModel(const AmDiagGmm &src_model,
                                      double tau,
                                      AccumAmDiagGmm *dst_stats) {
    int num_pdfs = src_model.NumPdfs();
    KALDI_ASSERT(num_pdfs == dst_stats->NumAccs());
    for (int32 pdf = 0; pdf < num_pdfs; pdf++) {
      AccumDiagGmm tmp_stats;
      double occ = 1.0; // its value doesn't matter.
      DiagGmmToStats(src_model.GetPdf(pdf), kGmmAll, occ, &tmp_stats);
      IsmoothStatsDiagGmm(tmp_stats, tau, &(dst_stats->GetAcc(pdf)));
    }
  }
  
  
  
  }  // End of namespace kaldi