// transform/fmpe.cc

// Copyright 2011-2012  Yanmin Qian  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 "transform/fmpe.h"
#include "util/text-utils.h"
#include "gmm/diag-gmm-normal.h"
#include "gmm/am-diag-gmm.h"
#include "hmm/transition-model.h"

namespace kaldi {

void Fmpe::SetContexts(std::string context_str) {
  // sets the contexts_ variable.
  using std::vector;
  using std::string;
  contexts_.clear();
  vector<string> ctx_vec; // splitting context_str on ":"
  SplitStringToVector(context_str, ":", false, &ctx_vec);
  contexts_.resize(ctx_vec.size());
  for (size_t i = 0; i < ctx_vec.size(); i++) {
    vector<string> pair_vec; // splitting ctx_vec[i] on ";"
    SplitStringToVector(ctx_vec[i], ";", false, &pair_vec);
    KALDI_ASSERT(pair_vec.size() != 0 && "empty context!");
    for (size_t j = 0; j < pair_vec.size(); j++) {
      vector<string> one_pair;
      SplitStringToVector(pair_vec[j], ",", false, &one_pair);
      KALDI_ASSERT(one_pair.size() == 2 &&
                   "Mal-formed context string: bad --context-expansion option?");
      int32 pos = 0;
      BaseFloat weight = BaseFloat(0);
      bool ok = ConvertStringToInteger(one_pair[0], &pos);
      ok = ConvertStringToReal(one_pair[1], &weight) && ok;
      if (!ok)
        KALDI_ERR << "Mal-formed context string: bad --context-expansion option?";
      contexts_[i].push_back(std::make_pair(pos, weight));
    }
  }
}

void Fmpe::ComputeC() {
  KALDI_ASSERT(gmm_.NumGauss() != 0.0);
  int32 dim = gmm_.Dim();

  // Getting stats from the GMM... assume the model is
  // correct.
  SpMatrix<double> x2_stats(dim);
  Vector<double> x_stats(dim);
  double tot_count = 0.0;
  DiagGmmNormal ngmm(gmm_);
  for (int32 pdf = 0; pdf < ngmm.NumGauss(); pdf++) {
    x2_stats.AddVec2(ngmm.weights_(pdf), ngmm.means_.Row(pdf));
    x2_stats.AddDiagVec(ngmm.weights_(pdf), ngmm.vars_.Row(pdf)); // add diagonal
    // covar to diagonal elements of x2_stats.
    x_stats.AddVec(ngmm.weights_(pdf), ngmm.means_.Row(pdf));
    tot_count += ngmm.weights_(pdf);
  }
  KALDI_ASSERT(tot_count != 0.0);
  x2_stats.Scale(1.0 / tot_count);
  x_stats.Scale(1.0 / tot_count);
  x2_stats.AddVec2(-1.0, x_stats); // subtract outer product of mean,
  // to get centered covariance.
  C_.Resize(dim);
  try {
    TpMatrix<double> Ctmp(dim); Ctmp.Cholesky(x2_stats);
    C_.CopyFromTp(Ctmp);
  } catch (...) {
    KALDI_ERR << "Error initializing fMPE object: cholesky of "
        "feature variance failed.  Probably code error, or NaN/inf in model";
  }
}

void Fmpe::ComputeStddevs() {
  const Matrix<BaseFloat> &inv_vars = gmm_.inv_vars();
  stddevs_.Resize(inv_vars.NumRows(), inv_vars.NumCols());
  stddevs_.CopyFromMat(inv_vars);
  stddevs_.ApplyPow(-0.5);
}


void Fmpe::ApplyContext(const MatrixBase<BaseFloat> &intermed_feat,
                        MatrixBase<BaseFloat> *feat_out) const {
  // Applies the temporal-context part of the transformation.
  int32 dim = FeatDim(), ncontexts = NumContexts(),
      T = intermed_feat.NumRows();
  KALDI_ASSERT(intermed_feat.NumCols() == dim * ncontexts &&
               intermed_feat.NumRows() == feat_out->NumRows()
               && feat_out->NumCols() == dim);
  // note: ncontexts == contexts_.size().
  for (int32 i = 0; i < ncontexts; i++) {
    // this_intermed_feat is the chunk of the "intermediate features"
    // that corresponds to this "context"
    SubMatrix<BaseFloat> this_intermed_feat(intermed_feat, 0, T,
                                            dim*i, dim);
    for (int32 j = 0; j < static_cast<int32>(contexts_[i].size()); j++) {
      int32 t_offset = contexts_[i][j].first;
      BaseFloat weight = contexts_[i][j].second;
      // Note: we could do this more efficiently using matrix operations,
      // but this doesn't dominate the computation and I think this is
      // clearer.
      for (int32 t_out = 0; t_out < T; t_out++) { // t_out indexes the output
        int32 t_in = t_out + t_offset; // t_in indexes the input.
        if (t_in >= 0 && t_in < T) // Discard frames outside range.
          feat_out->Row(t_out).AddVec(weight, this_intermed_feat.Row(t_in));
      }
    }
  }
}

void Fmpe::ApplyContextReverse(const MatrixBase<BaseFloat> &feat_deriv,
                               MatrixBase<BaseFloat> *intermed_feat_deriv)
    const {
  // Applies the temporal-context part of the transformation,
  // in reverse, for getting derivatives for training.
  int32 dim = FeatDim(), ncontexts = NumContexts(),
      T = feat_deriv.NumRows();
  KALDI_ASSERT(intermed_feat_deriv->NumCols() == dim * ncontexts &&
               intermed_feat_deriv->NumRows() == feat_deriv.NumRows()
               && feat_deriv.NumCols() == dim);
  // note: ncontexts == contexts_.size().
  for (int32 i = 0; i < ncontexts; i++) {
    // this_intermed_feat is the chunk of the derivative of
    // "intermediate features" that corresponds to this "context"
    // (this is output, in this routine).
    SubMatrix<BaseFloat> this_intermed_feat_deriv(*intermed_feat_deriv, 0, T,
                                                  dim*i, dim);
    for (int32 j = 0; j < static_cast<int32>(contexts_[i].size()); j++) {
      int32 t_offset = contexts_[i][j].first;
      BaseFloat weight = contexts_[i][j].second;
      // Note: we could do this more efficiently using matrix operations,
      // but this doesn't dominate the computation and I think this is
      // clearer.
      for (int32 t_out = 0; t_out < T; t_out++) { // t_out indexes the output
        int32 t_in = t_out + t_offset; // t_in indexes the input.
        if (t_in >= 0 && t_in < T) // Discard frames outside range.
          this_intermed_feat_deriv.Row(t_in).AddVec(weight,
                                                    feat_deriv.Row(t_out));
        // Note: the line above is where the work happens; it's the same
        // as in ApplyContext except reversing the input and output.
      }
    }
  }
}

void Fmpe::ApplyC(MatrixBase<BaseFloat> *feat_out, bool reverse) const {
  int32 T = feat_out->NumRows();
  Vector<BaseFloat> tmp(feat_out->NumCols());
  for (int32 t = 0; t < T; t++) {
    SubVector<BaseFloat> row(*feat_out, t);
    // Next line does: tmp = C_ * row
    tmp.AddTpVec(1.0, C_, (reverse ? kTrans : kNoTrans), row, 0.0);
    row.CopyFromVec(tmp);
  }
}

// Constructs the high-dim features and applies the main projection matrix
// projT_.  This projects from dimension ngauss*(dim+1) to dim*ncontexts.  Note:
// because the input vector of size ngauss*(dim+1) is sparse in a blocky way
// (i.e. each frame only has a couple of nonzero posteriors), we deal with
// sub-matrices of the projection matrix projT_.  We actually further optimize
// the code by taking all frames in a file that had nonzero posteriors for a
// particular Gaussian, and forming a matrix out of the corresponding
// high-dimensional features; we can then use a matrix-matrix multiply rather
// than using vector-matrix operations.

void Fmpe::ApplyProjection(const MatrixBase<BaseFloat> &feat_in,
                           const std::vector<std::vector<int32> > &gselect,
                           MatrixBase<BaseFloat> *intermed_feat) const {
  int32 dim = FeatDim(), ncontexts = NumContexts();  
  
  Vector<BaseFloat> post; // will be posteriors of selected Gaussians.
  Vector<BaseFloat> input_chunk(dim+1); // will be a segment of
  // the high-dimensional features.

  // "all_posts" is a vector of ((gauss-index, time-index), gaussian
  // posterior).
  // We'll compute the posterior information, sort it, and then
  // go through it in sorted order, which maintains memory locality
  // when accessing the projection matrix.
  // Note: if we really cared we could make this use level-3 BLAS
  // (matrix-matrix multiply), but we'd need to have a temporary
  // matrix for the output and input.
  std::vector<std::pair<std::pair<int32, int32>, BaseFloat> > all_posts;
  
  for (int32 t = 0; t < feat_in.NumRows(); t++) {
    SubVector<BaseFloat> this_feat(feat_in, t);
    gmm_.LogLikelihoodsPreselect(this_feat, gselect[t], &post);
    // At this point, post will contain log-likes of the selected
    // Gaussians.
    post.ApplySoftMax(); // Now they are posteriors (which sum to one).
    for (int32 i = 0; i < post.Dim(); i++) {
      int32 gauss = gselect[t][i];
      all_posts.push_back(std::make_pair(std::make_pair(gauss, t), post(i)));
    }
  }
  std::sort(all_posts.begin(), all_posts.end());
  
  bool optimize = true;

  if (!optimize) { // Why do we keep this un-optimized code around?
    // For clarity, so you can see what's going on, and for easier
    // comparision with ApplyProjectionReverse which is similar to this
    // un-optimized segment.  Both un-optimized and optimized versions
    // should give identical transforms (up to tiny roundoff differences).
    for (size_t i = 0; i < all_posts.size(); i++) {
      int32 gauss = all_posts[i].first.first, t = all_posts[i].first.second;
      SubVector<BaseFloat> this_feat(feat_in, t);
      SubVector<BaseFloat> this_intermed_feat(*intermed_feat, t);
      BaseFloat this_post = all_posts[i].second;
      SubVector<BaseFloat> this_stddev(stddevs_, gauss);

      // The next line is equivalent to setting input_chunk to
      // -this_post * the gaussian mean / (gaussian stddev).  Note: we use
      // the fact that mean * inv_var *  stddev == mean / stddev.
      input_chunk.Range(0, dim).AddVecVec(-this_post, gmm_.means_invvars().Row(gauss),
                                          this_stddev, 0.0);
      // The next line is equivalent to adding (feat / gaussian stddev) to
      // input_chunk, so now it contains (feat - mean) / stddev, which is
      // our "normalized" feature offset.
      input_chunk.Range(0, dim).AddVecDivVec(this_post, this_feat, this_stddev,
                                             1.0);
      // The last element of this input_chunk is the posterior itself
      // (between 0 and 1).
      input_chunk(dim) = this_post * config_.post_scale;

      // this_intermed_feat += [appropriate chjunk of projT_] * input_chunk.
      this_intermed_feat.AddMatVec(1.0, projT_.Range(gauss*(dim+1), dim+1,
                                                     0, dim*ncontexts),
                                   kTrans, input_chunk, 1.0);
    }
  } else {
    size_t i = 0;
    // We process the "posts" vector in chunks, where each chunk corresponds to
    // the same Gaussian index (but different times).
    while (i < all_posts.size()) {
      int32 gauss = all_posts[i].first.first;
      SubVector<BaseFloat> this_stddev(stddevs_, gauss),
          this_mean_invvar(gmm_.means_invvars(), gauss);
      SubMatrix<BaseFloat> this_projT_chunk(projT_, gauss*(dim+1), dim+1,
                                            0, dim*ncontexts);
      int32 batch_size; // number of posteriors with same Gaussian..
      for (batch_size = 0;
           batch_size+i < static_cast<int32>(all_posts.size()) &&
               all_posts[batch_size+i].first.first == gauss;
           batch_size++); // empty loop body.
      Matrix<BaseFloat> input_chunks(batch_size, dim+1);
      Matrix<BaseFloat> intermed_temp(batch_size, dim*ncontexts);
      for (int32 j = 0; j < batch_size; j++) { // set up "input_chunks".
        // To understand this code, first examine code and comments in "non-optimized"
        // code chunk above (the other branch of the if/else statement).
        int32 t = all_posts[i+j].first.second;
        SubVector<BaseFloat> this_feat(feat_in, t);
        SubVector<BaseFloat> this_input_chunk(input_chunks, j);
        BaseFloat this_post = all_posts[i+j].second;
        this_input_chunk.Range(0, dim).AddVecVec(-this_post,
                                                 this_mean_invvar,
                                                 this_stddev, 0.0);
        this_input_chunk.Range(0, dim).AddVecDivVec(this_post, this_feat,
                                                    this_stddev, 1.0);
        this_input_chunk(dim) = this_post * config_.post_scale;
      }
      // The next line is where most of the computation will happen,
      // during the feature computation phase.  We have rearranged
      // stuff so it's a matrix-matrix operation, for greater
      // efficiency (when using optimized libraries like ATLAS).
      intermed_temp.AddMatMat(1.0, input_chunks, kNoTrans,
                              this_projT_chunk, kNoTrans, 0.0);
      for (int32 j = 0; j < batch_size; j++) { // add data from
        // intermed_temp to the output "intermed_feat"
        int32 t = all_posts[i+j].first.second;
        SubVector<BaseFloat> this_intermed_feat(*intermed_feat, t);
        SubVector<BaseFloat> this_intermed_temp(intermed_temp, j);
        // this_intermed_feat += this_intermed_temp.
        this_intermed_feat.AddVec(1.0, this_intermed_temp);
      }
      i += batch_size;
    }
  }
}      



// This function does the reverse to ApplyProjection, for the case
// where we want the derivatives w.r.t. the projection matrix.
// It stores the positive and negative parts of this separately.
void Fmpe::ApplyProjectionReverse(const MatrixBase<BaseFloat> &feat_in,
                                  const std::vector<std::vector<int32> > &gselect,
                                  const MatrixBase<BaseFloat> &intermed_feat_deriv,
                                  MatrixBase<BaseFloat> *proj_deriv_plus,
                                  MatrixBase<BaseFloat> *proj_deriv_minus) const {
  int32 dim = FeatDim(), ncontexts = NumContexts();  
  
  Vector<BaseFloat> post; // will be posteriors of selected Gaussians.
  Vector<BaseFloat> input_chunk(dim+1); // will be a segment of
  // the high-dimensional features.

  // "all_posts" is a vector of ((gauss-index, time-index), gaussian
  // posterior).
  // We'll compute the posterior information, sort it, and then
  // go through it in sorted order, which maintains memory locality
  // when accessing the projection matrix.
  std::vector<std::pair<std::pair<int32, int32>, BaseFloat> > all_posts;
  
  for (int32 t = 0; t < feat_in.NumRows(); t++) {
    SubVector<BaseFloat> this_feat(feat_in, t);
    gmm_.LogLikelihoodsPreselect(this_feat, gselect[t], &post);
    // At this point, post will contain log-likes of the selected
    // Gaussians.
    post.ApplySoftMax(); // Now they are posteriors (which sum to one).
    for (int32 i = 0; i < post.Dim(); i++) {
      // The next few lines (where we set up "input_chunk") are identical
      // to ApplyProjection.
      int32 gauss = gselect[t][i];
      all_posts.push_back(std::make_pair(std::make_pair(gauss, t), post(i)));
    }
  }
  std::sort(all_posts.begin(), all_posts.end());
  for (size_t i = 0; i < all_posts.size(); i++) {
    int32 gauss = all_posts[i].first.first, t = all_posts[i].first.second;
    BaseFloat this_post = all_posts[i].second;
    SubVector<BaseFloat> this_feat(feat_in, t);    
    SubVector<BaseFloat> this_intermed_feat_deriv(intermed_feat_deriv, t);
    SubVector<BaseFloat> this_stddev(stddevs_, gauss);
    input_chunk.Range(0, dim).AddVecVec(-this_post, gmm_.means_invvars().Row(gauss),
                                        this_stddev, 0.0);
    input_chunk.Range(0, dim).AddVecDivVec(this_post, this_feat, this_stddev,
                                           1.0);
    input_chunk(dim) = this_post * config_.post_scale;

    // If not for accumulating the + and - parts separately, we would be
    // doing something like:
    // proj_deriv_.Range(0, dim*ncontexts, gauss*(dim+1), dim+1).AddVecVec(
    //                    1.0, this_intermed_feat_deriv, input_chunk);


    SubMatrix<BaseFloat> plus_chunk(*proj_deriv_plus, 
                                    gauss*(dim+1), dim+1,
                                    0, dim*ncontexts),
        minus_chunk(*proj_deriv_minus, 
                    gauss*(dim+1), dim+1,
                    0, dim*ncontexts);
          
    // This next function takes the rank-one matrix
    //  (input_chunk * this_intermed_deriv'), and adds the positive
    // part to proj_deriv_plus, and minus the negative part to
    // proj_deriv_minus.
    AddOuterProductPlusMinus(static_cast<BaseFloat>(1.0),
                             input_chunk,
                             this_intermed_feat_deriv,
                             &plus_chunk, &minus_chunk);
  }
}      

void Fmpe::ComputeFeatures(const MatrixBase<BaseFloat> &feat_in,
                           const std::vector<std::vector<int32> > &gselect,
                           Matrix<BaseFloat> *feat_out) const {
  int32 dim = FeatDim();
  KALDI_ASSERT(feat_in.NumRows() != 0 && feat_in.NumCols() == dim);
  KALDI_ASSERT(feat_in.NumRows() == static_cast<int32>(gselect.size()));
  feat_out->Resize(feat_in.NumRows(), feat_in.NumCols()); // will zero it.
  
  // Intermediate-dimension features
  Matrix<BaseFloat> intermed_feat(feat_in.NumRows(),
                                  dim * NumContexts());

  // Apply the main projection, from high-dim to intermediate
  // dimension (dim * NumContexts()).
  ApplyProjection(feat_in, gselect, &intermed_feat);

  // Apply the temporal context and reduces from
  // dimension dim*ncontexts to dim.
  ApplyContext(intermed_feat, feat_out);

  // Lastly, apply the the "C" matrix-- linear transform on the offsets.
  ApplyC(feat_out);
}


void Fmpe::AccStats(const MatrixBase<BaseFloat> &feat_in,
                    const std::vector<std::vector<int32> > &gselect,
                    const MatrixBase<BaseFloat> &direct_feat_deriv,
                    const MatrixBase<BaseFloat> *indirect_feat_deriv, // may be NULL
                    FmpeStats *fmpe_stats) const {
  SubMatrix<BaseFloat> stats_plus(fmpe_stats->DerivPlus());
  SubMatrix<BaseFloat> stats_minus(fmpe_stats->DerivMinus());
  int32 dim = FeatDim(), ncontexts = NumContexts();
  KALDI_ASSERT(feat_in.NumRows() != 0 && feat_in.NumCols() == dim);
  KALDI_ASSERT(feat_in.NumRows() == static_cast<int32>(gselect.size()));
  KALDI_ASSERT(SameDim(stats_plus, projT_) && SameDim(stats_minus, projT_) &&
               SameDim(feat_in, direct_feat_deriv));

  if (indirect_feat_deriv != NULL)
    fmpe_stats->AccumulateChecks(feat_in, direct_feat_deriv, *indirect_feat_deriv);
  
  Matrix<BaseFloat> feat_deriv(direct_feat_deriv); // "feat_deriv" is initially direct+indirect.
  if (indirect_feat_deriv != NULL)
    feat_deriv.AddMat(1.0, *indirect_feat_deriv);
  
  // We do the "*Reverse" version of each stage now, in reverse order.
  ApplyCReverse(&feat_deriv);
  
  Matrix<BaseFloat> intermed_feat_deriv(feat_in.NumRows(), dim*ncontexts);
  ApplyContextReverse(feat_deriv, &intermed_feat_deriv);
  
  ApplyProjectionReverse(feat_in, gselect, intermed_feat_deriv,
                         &stats_plus, &stats_minus);
}


void FmpeOptions::Write(std::ostream &os, bool binary) const {
  WriteToken(os, binary, context_expansion);
  WriteBasicType(os, binary, post_scale);
}
void FmpeOptions::Read(std::istream &is, bool binary) {
  ReadToken(is, binary, &context_expansion);
  ReadBasicType(is, binary, &post_scale);
}

Fmpe::Fmpe(const DiagGmm &gmm, const FmpeOptions &config): gmm_(gmm),
                                                          config_(config) {
  SetContexts(config.context_expansion);
  ComputeC();
  ComputeStddevs();
  projT_.Resize(NumGauss() * (FeatDim()+1), FeatDim() * NumContexts());
}

BaseFloat Fmpe::Update(const FmpeUpdateOptions &config,
                       const FmpeStats &stats) {
  SubMatrix<BaseFloat> proj_deriv_plus = stats.DerivPlus(),
      proj_deriv_minus = stats.DerivMinus();
  // tot_linear_objf_impr is the change in the actual
  // objective function if it were linear, i.e.
  //   objf-gradient . parameter-change
  // Note: none of this is normalized by the #frames (we don't have
  // this info here), so that is done at the script level.
  BaseFloat tot_linear_objf_impr = 0.0;
  int32 changed = 0; // Keep track of how many elements change sign.
  KALDI_ASSERT(SameDim(proj_deriv_plus, projT_) && SameDim(proj_deriv_minus, projT_));
  KALDI_ASSERT(proj_deriv_plus.Min() >= 0);
  KALDI_ASSERT(proj_deriv_minus.Min() >= 0);
  BaseFloat learning_rate = config.learning_rate,
      l2_weight = config.l2_weight;
  
  for (int32 i = 0; i < projT_.NumRows(); i++) {
    for (int32 j = 0; j < projT_.NumCols(); j++) {
      BaseFloat p = proj_deriv_plus(i, j), n = proj_deriv_minus(i, j),
          x = projT_(i, j);
      // Suppose the basic update (before regularization) is:
      // z <-- x  +   learning_rate * (p - n) / (p + n),
      // where z is the new parameter and x is the old one.
      // Here, we view (learning_rate / (p + n)) as a parameter-specific
      // learning rate.  In fact we view this update as the maximization
      // of an auxiliary function of the form:
      //  (z-x).(p-n)    - 0.5 (z - x)^2 (p+n)/learning_rate
      // and taking the derivative w.r.t z, we get:
      // Q'(z) =  (p-n) - (z - x) (p+n) / learning_rate
      // which we set to zero and solve for z, to get z = x + learning_rate.(p-n)/(p+n)
      // At this point we add regularization, a term of the form -l2_weight * z^2.
      // Our new auxiliary function derivative is:
      // Q(z) = -2.l2_weight.z + (p-n) - (z - x) (p+n) / learning_rate
      // We can write this as:
      // Q(z) = z . (-2.l2_weight - (p+n)/learning_rate)
      //        + (p-n) + x(p+n)/learning_rate
      // solving for z, we get:
      //      z = ((p-n) + x (p+n)/learning_rate) / (2.l2_weight + (p+n)/learning_rate)

      BaseFloat z = ((p-n) + x*(p+n)/learning_rate) / (2*l2_weight + (p+n)/learning_rate);
      // z is the new parameter value.

      tot_linear_objf_impr += (z-x) * (p-n); // objf impr based on linear assumption.
      projT_(i, j) = z;
      if (z*x < 0) changed++;
    }
  }
  KALDI_LOG << "Objf impr (assuming linear) is " << tot_linear_objf_impr;
  KALDI_LOG << ((100.0*changed)/(projT_.NumRows()*projT_.NumCols()))
            << "% of matrix elements changed sign.";
  return tot_linear_objf_impr;
}

// Note: we write the GMM first, without any other header.
// This way, the gselect code can treat the form on disk as
// a normal GMM object.
void Fmpe::Write(std::ostream &os, bool binary) const {
  if (gmm_.NumGauss() == 0)
    KALDI_ERR << "Fmpe::Write, object not initialized.";
  gmm_.Write(os, binary);
  config_.Write(os, binary);
  // stddevs_ are derived, don't write them.
  projT_.Write(os, binary);
  C_.Write(os, binary);
  // contexts_ are derived from config, don't write them.
}


void Fmpe::Read(std::istream &is, bool binary) {
  gmm_.Read(is, binary);
  config_.Read(is, binary);
  ComputeStddevs(); // computed from gmm.
  projT_.Read(is, binary);
  C_.Read(is, binary);
  SetContexts(config_.context_expansion);
}


BaseFloat ComputeAmGmmFeatureDeriv(const AmDiagGmm &am_gmm,
                                   const TransitionModel &trans_model,
                                   const Posterior &posterior,
                                   const MatrixBase<BaseFloat> &features,
                                   Matrix<BaseFloat> *direct_deriv,
                                   const AccumAmDiagGmm *model_diff,
                                   Matrix<BaseFloat> *indirect_deriv) {
  KALDI_ASSERT((model_diff != NULL) == (indirect_deriv != NULL));
  BaseFloat ans = 0.0;
  KALDI_ASSERT(posterior.size() == static_cast<size_t>(features.NumRows()));
  direct_deriv->Resize(features.NumRows(), features.NumCols());
  if (indirect_deriv != NULL)
    indirect_deriv->Resize(features.NumRows(), features.NumCols());
  Vector<BaseFloat> temp_vec(features.NumCols());
  Vector<double> temp_vec_dbl(features.NumCols());
  for (size_t i = 0; i < posterior.size(); i++) {
    for (size_t j = 0; j < posterior[i].size(); j++) {
      int32 tid = posterior[i][j].first,  // transition identifier.
          pdf_id = trans_model.TransitionIdToPdf(tid);
      BaseFloat weight = posterior[i][j].second;
      const DiagGmm &gmm = am_gmm.GetPdf(pdf_id);
      Vector<BaseFloat> gauss_posteriors;
      SubVector<BaseFloat> this_feat(features, i);
      SubVector<BaseFloat> this_direct_deriv(*direct_deriv, i);
      ans += weight * 
          gmm.ComponentPosteriors(this_feat, &gauss_posteriors);
      
      gauss_posteriors.Scale(weight);
      // The next line does: to i'th row of deriv, add
      // means_invvars^T * gauss_posteriors,
      // where each row of means_invvars is the mean times
      // diagonal inverse covariance... after transposing,
      // this becomes a weighted of these rows, weighted by
      // the posteriors.  This comes from the term
      //  feat^T * inv_var * mean
      // in the objective function.
      this_direct_deriv.AddMatVec(1.0, gmm.means_invvars(), kTrans,
                                  gauss_posteriors, 1.0);      

      // next line does temp_vec == inv_vars^T * gauss_posteriors,
      // which sets temp_vec to a weighted sum of the inv_vars,
      // weighed by Gaussian posterior.
      temp_vec.AddMatVec(1.0, gmm.inv_vars(), kTrans,
                         gauss_posteriors, 0.0);
      // Add to the derivative, -(this_feat .* temp_vec),
      // which is the term that comes from the -0.5 * inv_var^T feat_sq,
      // in the objective function (where inv_var is a vector, and feat_sq
      // is a vector of squares of the feature values).
      // Note: we have to do some messing about with double-precision here
      // because the stats only come in double precision.
      this_direct_deriv.AddVecVec(-1.0, this_feat, temp_vec, 1.0);
      if (model_diff != NULL && weight > 0.0) { // We need to get the indirect diff.
        // This "weight > 0.0" checks that this is the numerator stats, as the
        // fMPE indirect diff applies only to the ML stats-- CAUTION, this
        // code will only work as-is for fMMI (and the stats should not be
        // canceled), due to the assumption that ML stats == num stats.
        Vector<double> gauss_posteriors_dbl(gauss_posteriors);
        const AccumDiagGmm &deriv_acc = model_diff->GetAcc(pdf_id);
        // part of the derivative.  Note: we could just store the direct and
        // indirect derivatives together in one matrix, but it makes it easier
        // to accumulate certain diagnostics if we store them separately.
        SubVector<BaseFloat> this_indirect_deriv(*indirect_deriv, i);
        // note: deriv_acc.mean_accumulator() contains the derivative of
        // the objective function w.r.t. the "x stats" accumulated for
        // this GMM.  variance_accumulator() is the same for the "x^2 stats".
        temp_vec_dbl.AddMatVec(1.0, deriv_acc.mean_accumulator(), kTrans,
                               gauss_posteriors_dbl, 0.0);
        this_indirect_deriv.AddVec(1.0, temp_vec_dbl);
        temp_vec_dbl.AddMatVec(1.0, deriv_acc.variance_accumulator(), kTrans,
                               gauss_posteriors_dbl, 0.0);
        temp_vec.CopyFromVec(temp_vec_dbl); // convert to float.
        // next line because d(x^2 stats for Gaussian)/d(feature) =
        // 2 * (gaussian posterior) * feature.
        this_indirect_deriv.AddVecVec(2.0, this_feat, temp_vec, 1.0);
      }
    }
  }
  return ans;
}


SubMatrix<BaseFloat> FmpeStats::DerivPlus() const { // const-ness not preserved.
  KALDI_ASSERT(deriv.NumRows() != 0);
  int32 proj_num_rows = deriv.NumRows(),
      proj_num_cols = deriv.NumCols()/2;
  return SubMatrix<BaseFloat>(deriv, 0, proj_num_rows,
                              0, proj_num_cols);
}
SubMatrix<BaseFloat> FmpeStats::DerivMinus() const { // const-ness not preserved.
  KALDI_ASSERT(deriv.NumRows() != 0);
  int32 proj_num_rows = deriv.NumRows(),
      proj_num_cols = deriv.NumCols()/2;
  return SubMatrix<BaseFloat>(deriv, 0, proj_num_rows,
                              proj_num_cols, proj_num_cols);
}

void FmpeStats::Init(const Fmpe &fmpe) {
  int32 num_rows = fmpe.ProjectionTNumRows(),
      num_cols = fmpe.ProjectionTNumCols();
  deriv.Resize(num_rows, num_cols*2);

  int32 feat_dim = fmpe.FeatDim();
  checks.Resize(8, feat_dim);
}

void FmpeStats::AccumulateChecks(const MatrixBase<BaseFloat> &feats,
                                 const MatrixBase<BaseFloat> &direct_deriv,
                                 const MatrixBase<BaseFloat> &indirect_deriv) {
  int32 T = feats.NumRows(), dim = feats.NumCols();
  KALDI_ASSERT(direct_deriv.NumRows() == T && direct_deriv.NumCols() == dim &&
               indirect_deriv.NumRows() == T && indirect_deriv.NumCols() == dim);
  KALDI_ASSERT(checks.NumRows() == 8 && checks.NumCols() == dim);
  for (int32 t = 0; t < T; t++) {
    for (int32 d = 0; d < dim; d++) {
      BaseFloat zero = 0.0;
      checks(0, d) += std::max(zero, direct_deriv(t, d));
      checks(1, d) += std::max(zero, -direct_deriv(t, d));
      checks(2, d) += std::max(zero, indirect_deriv(t, d));
      checks(3, d) += std::max(zero, -indirect_deriv(t, d));
      checks(4, d) += std::max(zero, feats(t, d)*direct_deriv(t, d));
      checks(5, d) += std::max(zero, -feats(t, d)*direct_deriv(t, d));
      checks(6, d) += std::max(zero, feats(t, d)*indirect_deriv(t, d));
      checks(7, d) += std::max(zero, -feats(t, d)*indirect_deriv(t, d));
    }
  }
}

void FmpeStats::DoChecks() {
  if (checks.IsZero()) {
    KALDI_LOG << "No checks will be done, probably indirect derivative was not used.";
    return;
  }
  int32 dim = checks.NumCols();
  Vector<double> shift_check(dim), shift_check2(dim), scale_check(dim), scale_check2(dim);
  for (int32 d = 0; d < dim; d++) {
    // shiftnumerator = direct+indirect deriv-- should be zero.
    double shift_num = checks(0, d) - checks(1, d) + checks(2, d) - checks(3, d),
        shift_den = checks(0, d) + checks(1, d) + checks(2, d) + checks(3, d),
        shift_den2 = fabs(checks(0, d) - checks(1, d)) + fabs(checks(2, d) - checks(3, d));
    shift_check(d) = shift_num / shift_den;
    shift_check2(d) = shift_num / shift_den2;
    double scale_num = checks(4, d) - checks(5, d) + checks(6, d) - checks(7, d),
        scale_den = checks(4, d) + checks(5, d) + checks(6, d) + checks(7, d),
        scale_den2 = fabs(checks(4, d) - checks(5, d)) + fabs(checks(6, d) - checks(7, d));
    scale_check(d) = scale_num / scale_den;
    scale_check2(d) = scale_num / scale_den2;
  }

  KALDI_LOG << "Shift-check is as follows (should be in range +- 0.01 or less)."
            << shift_check;
  KALDI_LOG << "Scale-check is as follows (should be in range +- 0.01 or less)."
            << scale_check;
  KALDI_LOG << "Shift-check(2) is as follows: most elements should be in range +-0.1: "
            << shift_check2;
  KALDI_LOG << "Scale-check(2) is as follows: most elements should be in range +-0.1: "
            << scale_check2;
}

void FmpeStats::Write(std::ostream &os, bool binary) const {
  deriv.Write(os, binary);
  checks.Write(os, binary);
}

void FmpeStats::Read(std::istream &is, bool binary, bool add) {
  deriv.Read(is, binary, add);
  checks.Read(is, binary, add);
}


}  // End of namespace kaldi