// sgmm2/estimate-am-sgmm2-ebw.cc

// Copyright 2012  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 "base/kaldi-common.h"
#include "sgmm2/estimate-am-sgmm2-ebw.h"
#include "util/kaldi-thread.h"
using std::vector;

namespace kaldi {

void EbwAmSgmm2Updater::Update(const MleAmSgmm2Accs &num_accs,
                               const MleAmSgmm2Accs &den_accs,
                               AmSgmm2 *model,
                               SgmmUpdateFlagsType flags,
                               BaseFloat *auxf_change_out,
                               BaseFloat *count_out) {

  // Various quantities need to be computed at the start, before we
  // change any of the model parameters.
  std::vector< SpMatrix<double> > Q_num, Q_den, H, S_means;

  if (flags & kSgmmPhoneProjections) {
    MleAmSgmm2Updater::ComputeQ(num_accs, *model, &Q_num);
    MleAmSgmm2Updater::ComputeQ(den_accs, *model, &Q_den);
  }
  if (flags & kSgmmCovarianceMatrix) { // compute the difference between
    // the num and den S_means matrices... this is what we will need.
    MleAmSgmm2Updater::ComputeSMeans(num_accs, *model, &S_means);
    std::vector< SpMatrix<double> > S_means_tmp;
    MleAmSgmm2Updater::ComputeSMeans(den_accs, *model, &S_means_tmp);
    for (size_t i = 0; i < S_means.size(); i++)
      S_means[i].AddSp(-1.0, S_means_tmp[i]);
  }
  if (flags & (kSgmmPhoneVectors | kSgmmPhoneWeightProjections))
    model->ComputeH(&H);

  Vector<double> gamma_num(num_accs.num_gaussians_);
  for (int32 j1 = 0; j1 < num_accs.num_groups_; j1++)
    gamma_num.AddRowSumMat(1.0, num_accs.gamma_[j1]);
  Vector<double> gamma_den(den_accs.num_gaussians_);
  for (int32 j1 = 0; j1 < den_accs.num_groups_; j1++)
    gamma_den.AddRowSumMat(1.0, den_accs.gamma_[j1]);

  BaseFloat tot_impr = 0.0;

  if (flags & kSgmmPhoneVectors)
    tot_impr += UpdatePhoneVectors(num_accs, den_accs, H, model);

  if (flags & kSgmmPhoneProjections)
    tot_impr += UpdateM(num_accs, den_accs, Q_num, Q_den,
                        gamma_num, gamma_den, model);

  if (flags & kSgmmPhoneWeightProjections)
    tot_impr += UpdateW(num_accs, den_accs, gamma_num, gamma_den, model);

  if (flags & kSgmmSpeakerWeightProjections)
    tot_impr += UpdateU(num_accs, den_accs, gamma_num, gamma_den, model);

  if (flags & kSgmmCovarianceMatrix)
    tot_impr += UpdateVars(num_accs, den_accs,
                           gamma_num, gamma_den, S_means, model);

  if (flags & kSgmmSubstateWeights)
    tot_impr += UpdateSubstateWeights(num_accs, den_accs, model);

  if (flags & kSgmmSpeakerProjections)
    tot_impr += UpdateN(num_accs, den_accs, gamma_num, gamma_den, model);


  if (auxf_change_out) *auxf_change_out = tot_impr * num_accs.total_frames_;
  if (count_out) *count_out = num_accs.total_frames_;

  if (fabs(num_accs.total_frames_ - den_accs.total_frames_) >
      0.01*(num_accs.total_frames_ + den_accs.total_frames_))
    KALDI_WARN << "Num and den frame counts differ, "
               << num_accs.total_frames_ << " vs. " << den_accs.total_frames_;

  BaseFloat like_diff = num_accs.total_like_ - den_accs.total_like_;

  KALDI_LOG << "***Averaged differenced likelihood per frame is "
            << (like_diff/num_accs.total_frames_)
            << " over " << (num_accs.total_frames_) << " frames.";
  KALDI_LOG << "***Note: for this to be at all meaningful, if you use "
            << "\"canceled\" stats you will have to renormalize this over "
            << "the \"real\" frame count.";
  KALDI_ASSERT(num_accs.total_frames_ > 0 && den_accs.total_frames_ > 0);

  model->ComputeNormalizers();
}


class EbwUpdatePhoneVectorsClass: public MultiThreadable { // For multi-threaded.
 public:
  EbwUpdatePhoneVectorsClass(const EbwAmSgmm2Updater *updater,
                             const MleAmSgmm2Accs &num_accs,
                             const MleAmSgmm2Accs &den_accs,
                             const std::vector<SpMatrix<double> > &H,
                             AmSgmm2 *model,
                             double *auxf_impr):
      updater_(updater), num_accs_(num_accs), den_accs_(den_accs),
      model_(model), H_(H), auxf_impr_ptr_(auxf_impr), auxf_impr_(0.0) { }

  EbwUpdatePhoneVectorsClass(const EbwUpdatePhoneVectorsClass &other) :
      MultiThreadable(other),
      updater_(other.updater_), num_accs_(other.num_accs_),
      den_accs_(other.den_accs_), model_(other.model_),
      H_(other.H_), auxf_impr_ptr_(other.auxf_impr_ptr_), auxf_impr_(0.0) { }

  ~EbwUpdatePhoneVectorsClass() {
    *auxf_impr_ptr_ += auxf_impr_;
  }

  inline void operator() () {
    // Note: give them local copy of the sums we're computing,
    // which will be propagated to the total sums in the destructor.
    updater_->UpdatePhoneVectorsInternal(num_accs_, den_accs_, H_, model_,
                                         &auxf_impr_, num_threads_, thread_id_);
  }
 private:
  const EbwAmSgmm2Updater *updater_;
  const MleAmSgmm2Accs &num_accs_;
  const MleAmSgmm2Accs &den_accs_;
  AmSgmm2 *model_;
  const std::vector<SpMatrix<double> > &H_;
  double *auxf_impr_ptr_;
  double auxf_impr_;
};


void EbwAmSgmm2Updater::ComputePhoneVecStats(
    const MleAmSgmm2Accs &accs,
    const AmSgmm2 &model,
    const std::vector<SpMatrix<double> > &H,
    int32 j1,
    int32 m,
    const Vector<double> &w_jm_in,
    double gamma_jm,
    Vector<double> *g_jm,
    SpMatrix<double> *H_jm) {
  Vector<double> w_jm(w_jm_in);
  if (!accs.a_.empty() && accs.a_[j1](m, 0) != 0) { // [SSGMM]
    w_jm.MulElements(accs.a_[j1].Row(m)); // multiply by "a" quantities..
    w_jm.Scale(1.0 / w_jm.Sum()); // renormalize.
  }
  g_jm->CopyFromVec(accs.y_[j1].Row(m));
  for (int32 i = 0; i < accs.num_gaussians_; i++) {
    double gamma_jmi = accs.gamma_[j1](m, i);
    double quadratic_term = std::max(gamma_jmi, gamma_jm * w_jm(i));
    double scalar = gamma_jmi - gamma_jm * w_jm(i) + quadratic_term
        * VecVec(model.w_.Row(i), model.v_[j1].Row(m));
    g_jm->AddVec(scalar, model.w_.Row(i));
    if (gamma_jmi != 0.0)
      H_jm->AddSp(gamma_jmi, H[i]);  // The most important term..
    if (quadratic_term > 1.0e-10)
      H_jm->AddVec2(static_cast<BaseFloat>(quadratic_term), model.w_.Row(i));
  }
}


// Runs the phone vectors update for a subset of states (called
// multi-threaded).
void EbwAmSgmm2Updater::UpdatePhoneVectorsInternal(
    const MleAmSgmm2Accs &num_accs,
    const MleAmSgmm2Accs &den_accs,
    const std::vector<SpMatrix<double> > &H,
    AmSgmm2 *model,
    double *auxf_impr,
    int32 num_threads,
    int32 thread_id) const {

  int32 block_size = (num_accs.num_groups_ + (num_threads-1)) / num_threads,
      j1_start = block_size * thread_id,
      j1_end = std::min(num_accs.num_groups_, j1_start + block_size);

  int32 S = num_accs.phn_space_dim_, I = num_accs.num_gaussians_;

  for (int32 j1 = j1_start; j1 < j1_end; j1++) {
    double num_state_count = 0.0,
        state_auxf_impr = 0.0;
    Vector<double> w_jm(I);
    for (int32 m = 0; m < model->NumSubstatesForGroup(j1); m++) {
      double gamma_jm_num = num_accs.gamma_[j1].Row(m).Sum();
      double gamma_jm_den = den_accs.gamma_[j1].Row(m).Sum();
      num_state_count += gamma_jm_num;
      Vector<double> g_jm_num(S);  // computed using eq. 58 of SGMM paper [for numerator stats]
      SpMatrix<double> H_jm_num(S);  // computed using eq. 59 of SGMM paper [for numerator stats]
      Vector<double> g_jm_den(S); // same, but for denominator stats.
      SpMatrix<double> H_jm_den(S);

      // Compute the weights for this sub-state.
      // w_jm = softmax([w_{k1}^T ... w_{kD}^T] * v_{jkm})  eq.(7)
      w_jm.AddMatVec(1.0, Matrix<double>(model->w_), kNoTrans,
                     Vector<double>(model->v_[j1].Row(m)), 0.0);
      w_jm.ApplySoftMax();
      // Note: in the ML code, in the SSGMM case, at this point the w_jm would
      // be modified with the "a" quantities to get the "\tilde{w}_{jm}" of the
      // SSGMM techreport.  But in this code, it gets done inside ComputePhoneVecStats.

      ComputePhoneVecStats(num_accs, *model, H, j1, m, w_jm, gamma_jm_num,
                           &g_jm_num, &H_jm_num);
      ComputePhoneVecStats(den_accs, *model, H, j1, m, w_jm, gamma_jm_den,
                           &g_jm_den, &H_jm_den);

      Vector<double> v_jm(model->v_[j1].Row(m));
      Vector<double> local_derivative(S); // difference of derivative of numerator
      // and denominator objetive function.
      local_derivative.AddVec(1.0, g_jm_num);
      local_derivative.AddSpVec(-1.0, H_jm_num, v_jm, 1.0);
      local_derivative.AddVec(-1.0, g_jm_den);
      local_derivative.AddSpVec(-1.0 * -1.0, H_jm_den, v_jm, 1.0);

      SpMatrix<double> quadratic_term(H_jm_num);
      quadratic_term.AddSp(1.0, H_jm_den);
      double substate_count = 1.0e-10 + gamma_jm_num + gamma_jm_den;
      quadratic_term.Scale( (substate_count + options_.tau_v) / substate_count);
      quadratic_term.Scale(1.0 / (options_.lrate_v + 1.0e-10) );

      Vector<double> delta_v_jm(S);

      SolverOptions opts;
      opts.name = "v";
      opts.K = options_.max_cond;
      opts.eps = options_.epsilon;

      double auxf_impr =
          ((gamma_jm_num + gamma_jm_den == 0) ? 0.0 :
           SolveQuadraticProblem(quadratic_term,
                                 local_derivative,
                                 opts, &delta_v_jm));

      v_jm.AddVec(1.0, delta_v_jm);
      model->v_[j1].Row(m).CopyFromVec(v_jm);
      state_auxf_impr += auxf_impr;
    }

    *auxf_impr += state_auxf_impr;
    if (j1 < 10 && thread_id == 0) {
      KALDI_LOG << "Objf impr for group j = " << j1 << "  is "
                << (state_auxf_impr / (num_state_count + 1.0e-10))
                << " over " << num_state_count << " frames";
    }
  }
}

double EbwAmSgmm2Updater::UpdatePhoneVectors(const MleAmSgmm2Accs &num_accs,
                                             const MleAmSgmm2Accs &den_accs,
                                             const vector< SpMatrix<double> > &H,
                                             AmSgmm2 *model) const {
  KALDI_LOG << "Updating phone vectors.";

  double count = 0.0, auxf_impr = 0.0;

  int32 J1 = num_accs.num_groups_;
  for (int32 j1 = 0; j1 < J1; j1++) count += num_accs.gamma_[j1].Sum();

  EbwUpdatePhoneVectorsClass c(this, num_accs, den_accs, H, model, &auxf_impr);
  RunMultiThreaded(c);

  auxf_impr /= count;

  KALDI_LOG << "**Overall auxf improvement for v is " << auxf_impr
            << " over " << count << " frames";
  return auxf_impr;
}


double EbwAmSgmm2Updater::UpdateM(const MleAmSgmm2Accs &num_accs,
                                  const MleAmSgmm2Accs &den_accs,
                                  const std::vector< SpMatrix<double> > &Q_num,
                                  const std::vector< SpMatrix<double> > &Q_den,
                                  const Vector<double> &gamma_num,
                                  const Vector<double> &gamma_den,
                                  AmSgmm2 *model) const {
  int32 S = model->PhoneSpaceDim(),
      D = model->FeatureDim(),
      I = model->NumGauss();

  Vector<double> impr_vec(I);

  for (int32 i = 0; i < I; i++) {
    double gamma_i_num = gamma_num(i), gamma_i_den = gamma_den(i);

    if (gamma_i_num + gamma_i_den == 0.0) {
      KALDI_WARN << "Not updating phonetic basis for i = " << i
                 << " because count is zero. ";
      continue;
    }

    Matrix<double> Mi(model->M_[i]);
    Matrix<double> L(D, S); // this is something like the Y quantity, which
    // represents the linear term in the objf on M-- except that we make it the local
    // derivative about the current value, instead of the derivative around zero.
    // But it's not exactly the derivative w.r.t. M, due to the factor of Sigma_i.
    // The auxiliary function is Q(x) = tr(M^T P Y) - 0.5 tr(P M Q M^T),
    // where P is Y^{-1}.  The quantity L we define here will be Y - M Q,
    // and you can think of this as like the local derivative, except there is
    // a term P in there.
    L.AddMat(1.0, num_accs.Y_[i]);
    L.AddMatSp(-1.0, Mi, kNoTrans, Q_num[i], 1.0);
    L.AddMat(-1.0, den_accs.Y_[i]);
    L.AddMatSp(-1.0*-1.0, Mi, kNoTrans, Q_den[i], 1.0);

    SpMatrix<double> Q(S); // This is a combination of the Q's for the numerator and denominator.
    Q.AddSp(1.0, Q_num[i]);
    Q.AddSp(1.0, Q_den[i]);

    double state_count = 1.0e-10 + gamma_i_num + gamma_i_den; // the count
    // represented by the quadratic part of the stats.
    Q.Scale( (state_count + options_.tau_M) / state_count );
    Q.Scale( 1.0 / (options_.lrate_M + 1.0e-10) );


    SolverOptions opts;
    opts.name = "M";
    opts.K = options_.max_cond;
    opts.eps = options_.epsilon;

    Matrix<double> deltaM(D, S);
    double impr =
        SolveQuadraticMatrixProblem(Q, L,
                                    SpMatrix<double>(model->SigmaInv_[i]),
                                    opts, &deltaM);

    impr_vec(i) = impr;
    Mi.AddMat(1.0, deltaM);
    model->M_[i].CopyFromMat(Mi);
    if (i < 10 || impr / state_count > 3.0) {
      KALDI_VLOG(2) << "Objf impr for projection M for i = " << i << ", is "
                    << (impr/(gamma_i_num + 1.0e-20)) << " over " << gamma_i_num
                    << " frames";
    }
  }
  BaseFloat tot_count = gamma_num.Sum(), tot_impr = impr_vec.Sum();

  tot_impr /= (tot_count + 1.0e-20);
  KALDI_LOG << "Overall auxiliary function improvement for model projections "
            << "M is " << tot_impr << " over " << tot_count << " frames";

  KALDI_VLOG(1) << "Updating M: num-count is " << gamma_num;
  KALDI_VLOG(1) << "Updating M: den-count is " << gamma_den;
  KALDI_VLOG(1) << "Updating M: objf-impr is " << impr_vec;

  return tot_impr;
}


// Note: we do just one iteration of the weight-projection update here.  The
// weak-sense auxiliary functions used don't really make sense if we do it for
// multiple iterations.  It would be possible to use a similar auxiliary
// function to the one on my (D. Povey)'s thesis for the Gaussian mixture
// weights, which would make sense for multiple iterations, but this would be a
// bit more complex to implement and probably would not give much improvement
// over this approach.
double EbwAmSgmm2Updater::UpdateW(const MleAmSgmm2Accs &num_accs,
                                  const MleAmSgmm2Accs &den_accs,
                                  const Vector<double> &gamma_num,
                                  const Vector<double> &gamma_den,
                                  AmSgmm2 *model) {
  KALDI_LOG << "Updating weight projections";

  int32 I = num_accs.num_gaussians_, S = num_accs.phn_space_dim_;

  Matrix<double> g_i_num(I, S), g_i_den(I, S);

  // View F_i_{num,den} as vectors of SpMatrix [i.e. symmetric matrices,
  // linearized into vectors]
  Matrix<double> F_i_num(I, (S*(S+1))/2), F_i_den(I, (S*(S+1))/2);

  Vector<double> impr_vec(I);

  // Get the F_i and g_i quantities-- this is done in parallel (multi-core),
  // using the same code we use in the ML update [except we get it for
  // numerator and denominator separately.]
  Matrix<double> w(model->w_);
  {
    std::vector<Matrix<double> > log_a_num;
    if (model->HasSpeakerDependentWeights())
      MleAmSgmm2Updater::ComputeLogA(num_accs, &log_a_num);
    double garbage;
    UpdateWClass c_num(num_accs, *model, w, log_a_num, &F_i_num, &g_i_num, &garbage);
    RunMultiThreaded(c_num);
  }
  {
    std::vector<Matrix<double> > log_a_den;
    if (model->HasSpeakerDependentWeights())
      MleAmSgmm2Updater::ComputeLogA(den_accs, &log_a_den);
    double garbage;
    UpdateWClass c_den(den_accs, *model, w, log_a_den, &F_i_den, &g_i_den, &garbage);
    RunMultiThreaded(c_den);
  }

  for (int32 i = 0; i < I; i++) {

    // auxf was originally formulated in terms of the change in w (i.e. the
    // g quantities are the local derivatives), so there is less hassle than
    // with some of the other updates, in changing it to be discriminative.
    // we essentially just difference the linear terms and add the quadratic
    // terms.

    Vector<double> derivative(g_i_num.Row(i));
    derivative.AddVec(-1.0, g_i_den.Row(i));
    // F_i_num quadratic_term is a bit like the negated 2nd derivative
    // of the numerator stats-- actually it's not the actual 2nd deriv,
    // but an upper bound on it.
    SpMatrix<double> quadratic_term(S), tmp_F(S);
    quadratic_term.CopyFromVec(F_i_num.Row(i));
    tmp_F.CopyFromVec(F_i_den.Row(i)); // tmp_F is used for Vector->SpMatrix conversion.
    quadratic_term.AddSp(1.0, tmp_F);

    double state_count = gamma_num(i) + gamma_den(i);

    quadratic_term.Scale((state_count + options_.tau_w) / (state_count + 1.0e-10));
    quadratic_term.Scale(1.0 / (options_.lrate_w + 1.0e-10) );

    Vector<double> delta_w(S);

    SolverOptions opts;
    opts.name = "w";
    opts.K = options_.max_cond;
    opts.eps = options_.epsilon;

    double objf_impr =
        SolveQuadraticProblem(quadratic_term, derivative, opts, &delta_w);

    impr_vec(i) = objf_impr;
    if (i < 10 || objf_impr / (gamma_num(i) + 1.0e-10) > 2.0) {
      KALDI_LOG << "Predicted objf impr for w per frame is "
                << (objf_impr / (gamma_num(i) + 1.0e-10))
                << " over " << gamma_num(i) << " frames.";
    }
    model->w_.Row(i).AddVec(1.0, delta_w);
  }
  KALDI_VLOG(1) << "Updating w: numerator count is " << gamma_num;
  KALDI_VLOG(1) << "Updating w: denominator count is " << gamma_den;
  KALDI_VLOG(1) << "Updating w: objf-impr is " << impr_vec;

  double tot_num_count = gamma_num.Sum(), tot_impr = impr_vec.Sum();
  tot_impr /= tot_num_count;

  KALDI_LOG << "**Overall objf impr for w per frame is "
            << tot_impr << " over " << tot_num_count
            << " frames.";
  return tot_impr;
}


double EbwAmSgmm2Updater::UpdateU(const MleAmSgmm2Accs &num_accs,
                                  const MleAmSgmm2Accs &den_accs,
                                  const Vector<double> &gamma_num,
                                  const Vector<double> &gamma_den,
                                  AmSgmm2 *model) {
  int32 T = num_accs.spk_space_dim_;
  double tot_impr = 0.0;
  for (int32 i = 0; i < num_accs.num_gaussians_; i++) {
    if (gamma_num(i) < 200.0) {
      KALDI_LOG << "Numerator count is small " << gamma_num(i) << " for gaussian "
                << i << ", not updating u_i.";
      continue;
    }
    Vector<double> u_i(model->u_.Row(i));
    Vector<double> delta_u(T);
    Vector<double> t(T); // derivative.
    t.AddVec(1.0, num_accs.t_.Row(i));
    t.AddVec(-1.0, den_accs.t_.Row(i));
    SpMatrix<double> U(T); // quadratic term.
    U.AddSp(1.0, num_accs.U_[i]);
    U.AddSp(1.0, den_accs.U_[i]);

    double state_count = gamma_num(i) + gamma_den(i);
    U.Scale((state_count + options_.tau_u) / (state_count + 1.0e-10));
    U.Scale(1.0 / (options_.lrate_u + 1.0e-10) );

    SolverOptions opts;
    opts.name = "u";
    opts.K = options_.max_cond;
    opts.eps = options_.epsilon;

    double impr = SolveQuadraticProblem(U, t, opts, &delta_u);
    double impr_per_frame = impr / gamma_num(i);
    if (impr_per_frame > options_.max_impr_u) {
      KALDI_WARN << "Updating speaker weight projections u, for Gaussian index "
                 << i << ", impr/frame is " << impr_per_frame << " over "
                 << gamma_num(i) << " frames, scaling back to not exceed "
                 << options_.max_impr_u;
      double scale = options_.max_impr_u / impr_per_frame;
      impr *= scale;
      delta_u.Scale(scale);
      // Note: a linear scaling of "impr" with "scale" is not quite accurate
      // in depicting how the quadratic auxiliary function varies as we change
      // the scale on "delta", but this does not really matter-- the goal is
      // to limit the auxiliary-function change to not be too large.
    }
    if (i < 10) {
      KALDI_LOG << "Objf impr for spk weight-projection u for i = " << (i)
                << ", is " << (impr / (gamma_num(i) + 1.0e-20)) << " over "
                << gamma_num(i) << " frames";
    }
    u_i.AddVec(1.0, delta_u);
    model->u_.Row(i).CopyFromVec(u_i);
    tot_impr += impr;
  }
  KALDI_LOG << "**Overall objf impr for u is " << (tot_impr/gamma_num.Sum())
            << ", over " << gamma_num.Sum() << " frames";
  return tot_impr;
}


double EbwAmSgmm2Updater::UpdateN(const MleAmSgmm2Accs &num_accs,
                                  const MleAmSgmm2Accs &den_accs,
                                  const Vector<double> &gamma_num,
                                  const Vector<double> &gamma_den,
                                  AmSgmm2 *model) const {
  if (num_accs.spk_space_dim_ == 0 || num_accs.R_.size() == 0 ||
      num_accs.Z_.size() == 0) {
    KALDI_ERR << "Speaker subspace dim is zero or no stats accumulated";
  }

  int32 I = num_accs.num_gaussians_, D = num_accs.feature_dim_,
      T = num_accs.spk_space_dim_;

  Vector<double> impr_vec(I);

  for (int32 i = 0; i < I; i++) {
    double gamma_i_num = gamma_num(i), gamma_i_den = gamma_den(i);
    if (gamma_i_num + gamma_i_den == 0.0) {
      KALDI_WARN << "Not updating speaker basis for i = " << i
                 << " because count is zero. ";
      continue;
    }
    Matrix<double> Ni(model->N_[i]);
    // See comment near declaration of L in UpdateM().  This update is the
    // same, but change M->N, Y->Z and Q->R.

    Matrix<double> L(D, T);
    L.AddMat(1.0, num_accs.Z_[i]);
    L.AddMatSp(-1.0, Ni, kNoTrans, num_accs.R_[i], 1.0);
    L.AddMat(-1.0, den_accs.Z_[i]);
    L.AddMatSp(-1.0*-1.0, Ni, kNoTrans, den_accs.R_[i], 1.0);

    SpMatrix<double> R(T); // combination of the numerator and denominator R's.
    R.AddSp(1.0, num_accs.R_[i]);
    R.AddSp(1.0, den_accs.R_[i]);

    double state_count = 1.0e-10 + gamma_i_num + gamma_i_den; // the count
    // represented by the quadratic part of the stats.
    R.Scale( (state_count + options_.tau_N) / state_count );
    R.Scale( 1.0 / (options_.lrate_N + 1.0e-10) );

    Matrix<double> deltaN(D, T);

    SolverOptions opts;
    opts.name = "N";
    opts.K = options_.max_cond;
    opts.eps = options_.epsilon;

    double impr =
        SolveQuadraticMatrixProblem(R, L,
                                    SpMatrix<double>(model->SigmaInv_[i]),
                                    opts, &deltaN);
    impr_vec(i) = impr;
    Ni.AddMat(1.0, deltaN);
    model->N_[i].CopyFromMat(Ni);
    if (i < 10 || impr / (state_count+1.0e-20) > 3.0) {
      KALDI_LOG << "Objf impr for spk projection N for i = " << (i)
                << ", is " << (impr / (gamma_i_num + 1.0e-20)) << " over "
                << gamma_i_num << " frames";
    }
  }

  KALDI_VLOG(1) << "Updating N: numerator count is " << gamma_num;
  KALDI_VLOG(1) << "Updating N: denominator count is " << gamma_den;
  KALDI_VLOG(1) << "Updating N: objf-impr is " << impr_vec;

  double tot_count = gamma_num.Sum(), tot_impr = impr_vec.Sum();
  tot_impr /= (tot_count + 1.0e-20);
  KALDI_LOG << "**Overall auxf impr for N is " << tot_impr
            << " over " << tot_count << " frames";
  return tot_impr;
}

double EbwAmSgmm2Updater::UpdateVars(const MleAmSgmm2Accs &num_accs,
                                     const MleAmSgmm2Accs &den_accs,
                                     const Vector<double> &gamma_num,
                                     const Vector<double> &gamma_den,
                                     const std::vector< SpMatrix<double> > &S_means,
                                     AmSgmm2 *model) const {
  // Note: S_means contains not only the quantity S_means in the paper,
  // but also has a term - (Y_i M_i^T + M_i Y_i^T).  Plus, it is differenced
  // between numerator and denominator.  We don't calculate it here,
  // because it had to be computed with the original model, before we
  // changed the M quantities.
  int32 I = num_accs.num_gaussians_;
  KALDI_ASSERT(S_means.size() == I);
  Vector<double> impr_vec(I);

  for (int32 i = 0; i < I; i++) {
    double num_count = gamma_num(i), den_count = gamma_den(i);

    SpMatrix<double> SigmaStats(S_means[i]);
    SigmaStats.AddSp(1.0, num_accs.S_[i]);
    SigmaStats.AddSp(-1.0, den_accs.S_[i]);
    // SigmaStats now contain the stats for estimating Sigma (as in the main SGMM paper),
    // differenced between num and den.
    SpMatrix<double> SigmaInvOld(model->SigmaInv_[i]), SigmaOld(model->SigmaInv_[i]);
    SigmaOld.Invert();
    double count = num_count - den_count;
    KALDI_ASSERT(options_.lrate_Sigma <= 1.0);
    double inv_lrate = 1.0 / options_.lrate_Sigma;
    // These formulas assure that the objective function behaves in
    // a roughly symmetric way w.r.t. num and den counts.
    double E_den = 1.0 + inv_lrate, E_num = inv_lrate - 1.0;

    double smoothing_count =
        (options_.tau_Sigma * inv_lrate) + // multiply tau_Sigma by inverse-lrate
        (E_den * den_count) +              // for compatibility with other updates.
        (E_num * num_count) +
        1.0e-10;
    SigmaStats.AddSp(smoothing_count, SigmaOld);
    count += smoothing_count;
    SigmaStats.Scale(1.0 / count);
    SpMatrix<double> SigmaInv(SigmaStats); // before floor and ceiling.  Currently sigma,
    // not its inverse.
    bool verbose = false;
    int n_floor = SigmaInv.ApplyFloor(SigmaOld, options_.cov_min_value, verbose);
    SigmaInv.Invert(); // make it inverse variance.
    int n_ceiling = SigmaInv.ApplyFloor(SigmaInvOld, options_.cov_min_value, verbose);

    // this auxf_change.
    double auxf_change = -0.5 * count *(TraceSpSp(SigmaInv, SigmaStats)
                                        - TraceSpSp(SigmaInvOld, SigmaStats)
                                        - SigmaInv.LogDet()
                                        + SigmaInvOld.LogDet());

    model->SigmaInv_[i].CopyFromSp(SigmaInv);
    impr_vec(i) = auxf_change;
    if (i < 10 || auxf_change / (num_count+den_count+1.0e-10) > 2.0
        || n_floor+n_ceiling > 0) {
      KALDI_LOG << "Updating variance: Auxf change per frame for Gaussian "
                << i << " is " << (auxf_change / num_count) << " over "
                << num_count << " frames " << "(den count was " << den_count
                << "), #floor,ceil was " << n_floor << ", " << n_ceiling;
    }
  }
  KALDI_VLOG(1) << "Updating Sigma: numerator count is " << gamma_num;
  KALDI_VLOG(1) << "Updating Sigma: denominator count is " << gamma_den;
  KALDI_VLOG(1) << "Updating Sigma: objf-impr is " << impr_vec;

  double tot_count = gamma_num.Sum(), tot_impr = impr_vec.Sum();
  tot_impr /= tot_count+1.0e-20;
  KALDI_LOG << "**Overall auxf impr for Sigma is " << tot_impr
            << " over " << tot_count << " frames";
  return tot_impr;
}


double EbwAmSgmm2Updater::UpdateSubstateWeights(
    const MleAmSgmm2Accs &num_accs,
    const MleAmSgmm2Accs &den_accs,
    AmSgmm2 *model) {
  KALDI_LOG << "Updating substate mixture weights";

  double tot_count = 0.0, tot_impr = 0.0;
  for (int32 j2 = 0; j2 < num_accs.num_pdfs_; j2++) {
    int32 M = model->NumSubstatesForPdf(j2);
    Vector<double> num_occs(M), den_occs(M),
        orig_weights(model->c_[j2]), weights(model->c_[j2]);

    for (int32 m = 0; m < M; m++) {
      num_occs(m) = num_accs.gamma_c_[j2](m)
          + options_.tau_c * weights(m);
      den_occs(m) = den_accs.gamma_c_[j2](m);
    }

    if (weights.Dim() > 1) {
      double begin_auxf = 0.0, end_auxf = 0.0;
      for (int32 m = 0; m < M; m++) {  // see eq. 4.32, Dan Povey's PhD thesis.
        begin_auxf += num_occs(m) * log (weights(m))
            - den_occs(m) * weights(m) / orig_weights(m);
      }
      for (int32 iter = 0; iter < 50; iter++) {
        Vector<double> k_jm(M);
        double max_m = 0.0;
        for (int32 m = 0; m < M; m++)
          max_m = std::max(max_m, den_occs(m)/orig_weights(m));
        for (int32 m = 0; m < M; m++)
          k_jm(m) = max_m - den_occs(m)/orig_weights(m);
        for (int32 m = 0; m < M; m++)
          weights(m) = num_occs(m) + k_jm(m)*weights(m);
        weights.Scale(1.0 / weights.Sum());
      }
      for (int32 m = 0; m < M; m++)
        weights(m) = std::max(weights(m),
                              static_cast<double>(options_.min_substate_weight));
      weights.Scale(1.0 / weights.Sum()); // renormalize.

      for (int32 m = 0; m < M; m++) {
        end_auxf += num_occs(m) * log (weights(m))
            - den_occs(m) * weights(m) / orig_weights(m);
      }
      tot_impr += end_auxf - begin_auxf;
      double this_impr = ((end_auxf - begin_auxf) / num_occs.Sum());
      if (j2 < 10 || this_impr > 0.5) {
        KALDI_LOG << "Updating substate weights: auxf impr for pdf " << j2
                  << " is " << this_impr << " per frame over " << num_occs.Sum()
                  << " frames (den count is " << den_occs.Sum() << ")";
      }
    }
    model->c_[j2].CopyFromVec(weights);
    tot_count += den_occs.Sum(); // Note: num and den occs should be the
    // same, except num occs are smoothed, so this is what we want.
  }

  tot_impr /= (tot_count + 1.0e-20);

  KALDI_LOG << "**Overall auxf impr for c is " << tot_impr
            << " over " << tot_count << " frames";
  return tot_impr;
}

}  // namespace kaldi