// gmm/indirect-diff-diag-gmm.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 "gmm/indirect-diff-diag-gmm.h"

namespace kaldi {


void GetSingleStatsDerivative(
    double ml_count, double ml_x_stats, double ml_x2_stats,
    double disc_count, double disc_x_stats, double disc_x2_stats,
    double model_mean, double model_var, BaseFloat min_variance,
    double *ml_x_stats_deriv, double *ml_x2_stats_deriv) {
  
  double model_inv_var = 1.0/model_var,
      model_inv_var_sq = model_inv_var*model_inv_var,
      model_mean_sq = model_mean*model_mean;

  // First get derivative of discriminative objective function w.r.t. the
  // model mean and variance.
  // Below: eqs. 11 and 13 in 2005 ICASSP paper on fMPE.  Note: the factor of
  // kappa (in the fMPE case) is assumed to have been accounted for by
  // scaling the num and den accs at the command-line level.  We substituted
  // eq. 12 into 13 and rearranged to get the second expression.
  double diff_wrt_model_mean = (1.0/model_var) * (disc_x_stats - model_mean*disc_count),
      diff_wrt_model_var =
      0.5 * ((disc_x2_stats - 2*model_mean*disc_x_stats + disc_count*model_mean_sq)
             * model_inv_var_sq
             - disc_count*model_inv_var);

  double stats_mean = ml_x_stats / ml_count,
      stats_var = ml_x2_stats / ml_count - (ml_x_stats / ml_count)*(ml_x_stats / ml_count);

  // We assume the "rescaling" update will be as follows.  Apologies if this is
  // a bit confusing.  The idea is that if the mean and var from (stats versus
  // model) differ we assume that the model will be updated with
  // DoRescalingUpdate(), which takes two sets of ML accs (old and new).  The old ML
  // accs given to the update will be the current ml accumulators we have here in
  // this function, and the new ML accs will be affected by change in fMPE transform.
  // The update in DoRescalingUpdate() will preserve any current difference between
  // the ml stats and the model [represented as a shift in mean and factor in variance].
  // Concretely: the update in DoRescalingUpdate() will do:
  //
  // new_model_mean := old_model_mean + new_stats_mean - old_stats_mean   (eq. 1)
  // new_model_var := max(min_variance, old_model_var * new_stats_var / old_stats_var).  (eq. 2)
  //
  // We're differentiating back through this process to new_stats_mean.
  // If the model and the stats were actually the same (e.g. we had been doing ML updates),
  // then all this is equivalent to what was in the original fMPE paper.  It's just
  // extended to make sense outside of that scenario where you're doing ML.
  
  double diff_wrt_stats_mean = diff_wrt_model_mean; // This comes from eq. 1 above.
  double diff_wrt_stats_var;
  if (model_var <= min_variance*1.01) {
    diff_wrt_stats_var = 0.0; // model would be "pinned" at minimum variance.
    KALDI_VLOG(2) << "Variance derivative is zero (min variance)";
  } else {
    diff_wrt_stats_var = diff_wrt_model_var * model_var / stats_var; // note:
    // the factor "model_var / stats_var" comes from "old_model_var / old_stats_var" in eq. 2.
    // Also note: the {old_,new_} versions of variables are numerically the same here, at the
    // point where we're differentiating.
  }

  // The next equations don't appear in the paper but represent the backpropagation
  // of the derivative through the equations:
  // stats_mean := ml_x_stats / ml_count
  // stats_var := ml_x2_stats / ml_count - (ml_x_stats/ml_count)^2
  // [we use stats_mean = ml_x_stats/ml_count, here].
  *ml_x_stats_deriv = diff_wrt_stats_mean / ml_count - 2 * diff_wrt_stats_var * stats_mean / ml_count;
  *ml_x2_stats_deriv = diff_wrt_stats_var / ml_count;
}




// The function for just one GMM.  We don't export it as it's not currently
// needed outside this file.
void GetStatsDerivative(const DiagGmm &gmm,
                        const AccumDiagGmm &num_acc,
                        const AccumDiagGmm &den_acc,
                        const AccumDiagGmm &ml_acc,
                        BaseFloat min_variance,
                        BaseFloat min_gaussian_occupancy,
                        AccumDiagGmm *out_accs) {
  out_accs->Resize(gmm, kGmmAll);
  int32 num_gauss = gmm.NumGauss(), dim = gmm.Dim();
  KALDI_ASSERT(num_gauss == num_acc.NumGauss() && dim == num_acc.Dim());
  KALDI_ASSERT(num_gauss == den_acc.NumGauss()); // don't check den dim--
  // in the "compressed" form of stats (where num acc stores diff),
  // it could be zero.
  KALDI_ASSERT(num_gauss == ml_acc.NumGauss() && dim == ml_acc.Dim());

  KALDI_ASSERT((ml_acc.Flags() & (kGmmMeans|kGmmVariances)) ==
               (kGmmMeans|kGmmVariances));
  KALDI_ASSERT((num_acc.Flags() & (kGmmMeans|kGmmVariances)) ==
               (kGmmMeans|kGmmVariances));
  DiagGmmNormal gmm_normal(gmm);

  // if have_den_stats == false, we assume the num and den have been
  // "compressed" by putting the difference in mean and var stats in num.
  bool have_den_stats = ((den_acc.Flags() & (kGmmMeans|kGmmVariances)) != 0);

  for (int32 gauss = 0; gauss < num_gauss; gauss++) {
    Vector<double> x_stats_deriv(dim), x2_stats_deriv(dim);
    double num_count = num_acc.occupancy()(gauss),
        den_count = den_acc.occupancy()(gauss),
        ml_count = ml_acc.occupancy()(gauss);
    
    if (ml_count <= min_gaussian_occupancy) {
      // This Gaussian won't be updated since has small count
      KALDI_WARN << "Skipping Gaussian because very small ML count: (num,den,ml) = "
                 << num_count << ", " << den_count << ", " << ml_count;
    } else {
      double disc_count = num_count - den_count;
      for (int32 d = 0; d < dim; d++) {
        double disc_x_acc = num_acc.mean_accumulator()(gauss, d)
            - (have_den_stats ? den_acc.mean_accumulator()(gauss, d) : 0.0),
            disc_x2_acc = num_acc.variance_accumulator()(gauss, d)
            - (have_den_stats ? den_acc.variance_accumulator()(gauss, d) : 0.0),
            ml_x_acc = ml_acc.mean_accumulator()(gauss, d),
            ml_x2_acc = ml_acc.variance_accumulator()(gauss, d),
            model_mean = gmm_normal.means_(gauss, d),
            model_var = gmm_normal.vars_(gauss, d);

        double x_acc_deriv = 0.0, x2_acc_deriv = 0.0;
        GetSingleStatsDerivative(ml_count, ml_x_acc, ml_x2_acc,
                                 disc_count, disc_x_acc, disc_x2_acc,
                                 model_mean, model_var, min_variance,
                                 &x_acc_deriv, &x2_acc_deriv);

        x_stats_deriv(d) = x_acc_deriv;
        x2_stats_deriv(d) = x2_acc_deriv;
      }
      // set the stats to these quantities (we're adding, but the stats
      // are currently zero).
      out_accs->AddStatsForComponent(gauss, 0.0, x_stats_deriv, x2_stats_deriv);
    }
  }
}

void GetStatsDerivative(const AmDiagGmm &gmm,
                        const AccumAmDiagGmm &num_accs, // for MMI, would equal ml accs.
                        const AccumAmDiagGmm &den_accs,
                        const AccumAmDiagGmm &ml_accs,
                        BaseFloat min_variance,
                        BaseFloat min_gaussian_occupancy,
                        AccumAmDiagGmm *out_accs) {
  out_accs->Init(gmm, kGmmAll);
  int32 num_pdfs = gmm.NumPdfs();
  KALDI_ASSERT(num_accs.NumAccs() == num_pdfs);
  KALDI_ASSERT(den_accs.NumAccs() == num_pdfs);
  KALDI_ASSERT(ml_accs.NumAccs() == num_pdfs);
  for (int32 pdf = 0; pdf < num_pdfs; pdf++)
    GetStatsDerivative(gmm.GetPdf(pdf), num_accs.GetAcc(pdf), den_accs.GetAcc(pdf),
                       ml_accs.GetAcc(pdf), min_variance, min_gaussian_occupancy,
                       &(out_accs->GetAcc(pdf)));
  
}


void DoRescalingUpdate(const AccumDiagGmm &old_ml_acc,
                       const AccumDiagGmm &new_ml_acc,
                       BaseFloat min_variance,
                       BaseFloat min_gaussian_occupancy,
                       DiagGmm *gmm,
                       double *tot_count,
                       double *tot_divergence) {
  int32 num_gauss = gmm->NumGauss(), dim = gmm->Dim();
  KALDI_ASSERT(old_ml_acc.NumGauss() == num_gauss &&
               old_ml_acc.Dim() == dim);
  KALDI_ASSERT(new_ml_acc.NumGauss() == num_gauss &&
               new_ml_acc.Dim() == dim);
  KALDI_ASSERT((old_ml_acc.Flags() & (kGmmMeans|kGmmVariances)) ==
               (kGmmMeans|kGmmVariances));
  KALDI_ASSERT((new_ml_acc.Flags() & (kGmmMeans|kGmmVariances)) ==
               (kGmmMeans|kGmmVariances));

  DiagGmmNormal gmm_normal(*gmm);
  for (int32 gauss = 0; gauss < num_gauss; gauss++) {
    double old_ml_count = old_ml_acc.occupancy()(gauss),
        new_ml_count = new_ml_acc.occupancy()(gauss);
    if (old_ml_count <= min_gaussian_occupancy ||
        new_ml_count <= min_gaussian_occupancy) {
      KALDI_WARN << "Gaussian being skipped because it has small count: (old,new) = "
                 << old_ml_count << ", " << new_ml_count;
      continue;
    }
    *tot_count += new_ml_count;
    for (int32 d = 0; d < dim; d++) {
      double old_model_mean = gmm_normal.means_(gauss, d),
          old_model_var = gmm_normal.vars_(gauss, d),
          old_ml_mean = old_ml_acc.mean_accumulator()(gauss, d) / old_ml_count,
          old_ml_var = old_ml_acc.variance_accumulator()(gauss, d) / old_ml_count
          - old_ml_mean*old_ml_mean,
          new_ml_mean = new_ml_acc.mean_accumulator()(gauss, d) / new_ml_count,
          new_ml_var = new_ml_acc.variance_accumulator()(gauss, d) / new_ml_count
          - new_ml_mean*new_ml_mean,
          new_model_mean = old_model_mean + new_ml_mean - old_ml_mean,
          new_model_var = std::max(static_cast<double>(min_variance),
                                   old_model_var * new_ml_var / old_ml_var);
      double divergence = 
          0.5 *(((new_model_mean-old_model_mean)*(new_model_mean-old_model_mean) +
                 new_model_var - old_model_var)/old_model_var +
                Log(old_model_var / new_model_var));
      if (divergence < 0.0)
        KALDI_WARN << "Negative divergence " << divergence;
      *tot_divergence += divergence * new_ml_count;
      gmm_normal.means_(gauss, d) = new_model_mean;
      gmm_normal.vars_(gauss, d) = new_model_var;
    }
  }
  gmm_normal.CopyToDiagGmm(gmm);
}


void DoRescalingUpdate(const AccumAmDiagGmm &old_ml_accs,
                       const AccumAmDiagGmm &new_ml_accs,
                       BaseFloat min_variance,
                       BaseFloat min_gaussian_occupancy,
                       AmDiagGmm *am_gmm) {
  int32 num_pdfs = am_gmm->NumPdfs();
  KALDI_ASSERT(old_ml_accs.NumAccs() == num_pdfs);
  KALDI_ASSERT(new_ml_accs.NumAccs() == num_pdfs);
  double tot_count = 0.0, tot_divergence = 0.0;
  for (int32 pdf = 0; pdf < num_pdfs; pdf++)
    DoRescalingUpdate(old_ml_accs.GetAcc(pdf), new_ml_accs.GetAcc(pdf),
                      min_variance, min_gaussian_occupancy, &am_gmm->GetPdf(pdf),
                      &tot_count, &tot_divergence);
  KALDI_LOG << "K-L divergence from old to new model is "
            << (tot_divergence/tot_count) << " over "
            << tot_count << " frames.";
  am_gmm->ComputeGconsts();
}



}  // End of namespace kaldi