// transform/fmpe.h
  
  // 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.
  
  
  #ifndef KALDI_TRANSFORM_FMPE_H_
  #define KALDI_TRANSFORM_FMPE_H_ 1
  
  #include <vector>
  
  #include "gmm/am-diag-gmm.h"
  #include "gmm/mle-am-diag-gmm.h"
  #include "hmm/transition-model.h"
  #include "hmm/posterior.h"
  
  namespace kaldi {
  
  
  struct FmpeOptions {
    // Probably the easiest place to start, to understand fMPE, is the
    // paper "Improvements to fMPE for discriminative training of features".
    // We are simplifying a few things here.  We are getting rid of the
    // "indirect differential"; we are adding a linear transform after the
    // high->low dimension projection whose function is to "un-whiten" the
    // transformed features (i.e. project from a nominally Gaussian-distributed
    // space into our actual feature space), in order to make it unnecessary to
    // take into account the per-dim variance during the update phase of fMPE;
    // and the update equations are rather simpler than described in
    // the paper; we take away some stuff, but add in the capability to
    // do l2 regularization during the update phase.
    
    std::string context_expansion; // This string describes the various contexts...
    // the easiest way to think of it is, we first generate the high-dimensional
    // features without context expansion, and we then append the left and right
    // frames, and also weighted averages of further-out frames, as specified by
    // this string.  Suppose there are 1024 Gaussians and the feature dimension is
    // 40.  In the simple way to describe it, supposing there are 9 contexts (the
    // central frame, the left and right frames, and 6 averages of more distant
    // frames), we generate the "offset features" of dimension (1024 * 41), then
    // add left and right temporal context to the high-dim features so the
    // dimension is (1024 * 41 * 9), and then project down to 40, so we train a
    // matrix of 40 x (1024 * 41 * 9).  As described in the paper, though, we
    // reorganize the computation for efficiency (it has to do with preserving
    // sparsity), and we train a matrix of dimension (40 * 9) x (1024 * 41).  The
    // (40 x 9) -> 40 transformation, which involves time as well, is dictated by
    // these contexts.
  
    // You probably won't want to mess with this "context_expansion" string.
    // The most important parameter to tune is the number of Gaussians in
    // the UBM.  Typically this will be in the range 300 to 1000.
  
    BaseFloat post_scale; // Scale on the posterior component of the high-dim
    // features (1 of these for every [feat-dim] of the offset features).
    // Typically 5.0-- this just gives a bit more emphasis to these posteriors
    // during training, like a faster learning rate.
    
    FmpeOptions(): context_expansion("0,1.0:-1,1.0:1,1.0:-2,0.5;-3,0.5:2,0.5;3,0.5:-4,0.5;-5,0.5:4,0.5;5,0.5:-6,0.333;-7,0.333;-8,0.333:6,0.333;7,0.333;8,0.333"),
                  post_scale(5.0) { }
  
    void Register(OptionsItf *opts) {
      opts->Register("post-scale", &post_scale, "Scaling constant on posterior "
                     "element of offset features, to give it a faster learning "
                     "rate.");
      opts->Register("context-expansion", &context_expansion, "Specifies the "
                     "temporal context-splicing of high-dimensional features.");
    }
    // We include write and read functions, since this
    // object is included as a member of the fMPE object.
    void Write(std::ostream &os, bool binary) const;
    void Read(std::istream &is, bool binary);
  };
  
  struct FmpeUpdateOptions {
    BaseFloat learning_rate; // Learning rate constant.  Like inverse of E
    // in the papers.
    BaseFloat l2_weight; // Weight on l2 regularization term
    
    FmpeUpdateOptions(): learning_rate(0.1), l2_weight(100.0) { }
  
    void Register(OptionsItf *opts) {
      opts->Register("learning-rate", &learning_rate,
                     "Learning rate constant (like inverse of E in fMPE papers)");
      opts->Register("l2-weight", &l2_weight,
                     "Weight on l2 regularization term in objective function.");
    }  
  };
  
  class Fmpe;
  
  struct FmpeStats {
    FmpeStats() { };
    void Init(const Fmpe &fmpe);
    FmpeStats(const Fmpe &fmpe) { Init(fmpe); }
  
    void Write(std::ostream &os, bool binary) const;
    void Read(std::istream &is, bool binary, bool add = false);
    
    SubMatrix<BaseFloat> DerivPlus() const;
    SubMatrix<BaseFloat> DerivMinus() const;
  
    /// If we're using the indirect differential, accumulates certain quantities
    /// that will be used in the update phase to verify that the computation
    /// of the indirect differential was done correctly
    void AccumulateChecks(const MatrixBase<BaseFloat> &feats,
                          const MatrixBase<BaseFloat> &direct_deriv,
                          const MatrixBase<BaseFloat> &indirect_deriv);
    void DoChecks(); // Will check that stuff cancels.  Just prints
    // messages for now.
   private:
    Matrix<BaseFloat> deriv; // contains positive and negative parts of derivatives
    // separately as sub-parts of the matrix, to ensure memory locality.
  
    // checks() is an 8 x fmpe.FeatDim() matrix that stores:
    //  (0-1) summed-deriv from direct, +ve and -ve part.
    //  (2-3) summed-deriv from indirect, +ve and -ve part.
    //  (4-5) (summed-deriv from direct * features), +ve and -ve part.
    //  (6-7) (summed-deriv from indirect * features), +ve and -ve part.
    Matrix<double> checks; // contains quantities we use to check the
    // indirect and direct derivatives are canceling as they should.
  
  };
  
  class Fmpe {
   public:
    Fmpe() {}
    Fmpe(const DiagGmm &gmm, const FmpeOptions &config);
  
    int32 FeatDim() const { return gmm_.Dim(); }
    int32 NumGauss() const { return gmm_.NumGauss(); }
    int32 NumContexts() const { return static_cast<int32>(contexts_.size()); }
  
    // Note: this returns the number of rows and columns in projT_,
    // which is the transpose of the high->intermediate dimensional
    // projection matrix.  This is the dimension we want for the
    // stats.
    int32 ProjectionTNumRows() const { return (FeatDim()+1) * NumGauss(); }
    int32 ProjectionTNumCols() const { return FeatDim() * NumContexts(); }
  
    
    // Computes the fMPE feature offsets and outputs them.
    // You can add feat_in to this afterwards, if you want.
    // Requires the Gaussian-selection info, which would normally
    // be computed by a separate program-- this consists of
    // lists of the top-scoring Gaussians for these features.
    void ComputeFeatures(const MatrixBase<BaseFloat> &feat_in,
                         const std::vector<std::vector<int32> > &gselect,
                         Matrix<BaseFloat> *feat_out) const;
  
    // For training-- compute the derivative w.r.t the projection matrix
    // (we keep the positive and negative parts separately to help
    // set the learning rates).
    void 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 *stats) const;
    
    // Note: the form on disk starts with the GMM; that way,
    // the gselect program can treat the fMPE object as if it
    // is a GMM.
    void Write(std::ostream &os, bool binary) const;
    void Read(std::istream &is, bool binary);
  
    // Returns total objf improvement, based on linear assumption.
    BaseFloat Update(const FmpeUpdateOptions &config,
                     const FmpeStats &stats);
    
   private:
    void SetContexts(std::string context_str);
    void ComputeC(); // Computes the Cholesky factor C, from the GMM.
    void ComputeStddevs();
  
    // Constructs the high-dim features and applies the main projection matrix proj_.
    void ApplyProjection(const MatrixBase<BaseFloat> &feat_in,
                         const std::vector<std::vector<int32> > &gselect,
                         MatrixBase<BaseFloat> *intermed_feat) const;
  
    // The same in reverse, for computing derivatives.
    void 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;
  
    // Applies the temporal context splicing from the intermediate
    // features-- adds the result to feat_out which at this point
    // will typically be zero.
    void ApplyContext(const MatrixBase<BaseFloat> &intermed_feat,
                      MatrixBase<BaseFloat> *feat_out) const;
  
    // This is as ApplyContext but for back-propagating the derivative.
    // Result is added to intermediate_feat_deriv which at this point will
    // typically be zero.
    void ApplyContextReverse(const MatrixBase<BaseFloat> &feat_deriv,
                             MatrixBase<BaseFloat> *intermed_feat_deriv) const;
  
    // Multiplies the feature offsets by the Cholesky matrix C.
    void ApplyC(MatrixBase<BaseFloat> *feat_out, bool reverse = false) const;
  
    // For computing derivatives-- multiply the derivatives by C^T,
    // which is the "reverse" of the forward pass of multiplying
    // by C (this is how derivatives behave...)
    void ApplyCReverse(MatrixBase<BaseFloat> *deriv) const { ApplyC(deriv, true); }
  
    
    
    DiagGmm gmm_; // The GMM used to get posteriors.
    FmpeOptions config_;
    Matrix<BaseFloat> stddevs_; // The standard deviations of the
    // variances of the GMM -- computed to avoid taking a square root
    // in the fMPE computation.   Derived variable-- not stored on
    // disk.
    Matrix<BaseFloat> projT_; // The transpose of the projection matrix;
    // this is of dimension
    // (NumGauss() * (FeatDim()+1)) * (FeatDim() * NumContexts()).
    
    TpMatrix<BaseFloat> C_; // Cholesky factor of the variance Sigma of
    // features around their mean (as estimated from GMM)... applied
    // to fMPE offset just before we add it to the features.  This allows
    // us to simplify the fMPE update and not have to worry about
    // the features having non-unit variance, and what effect this should
    // have on the learning rate..
    
    // The following variable dictates how we use temporal context.
    // e.g. contexts = { { (0, 1.0) }, { (-1, 1.0) }, { (1, 1.0) },
    //                   { (-2, 0.5 ), (-3, 0.5) }, ...  }
    std::vector<std::vector<std::pair<int32, BaseFloat> > > contexts_;
    
  };
  
  /// Computes derivatives of the likelihood of these states (weighted),
  /// w.r.t. the feature values.  Used in fMPE training.  Note, the
  /// weights "posterior" may be positive or negative-- for MMI, MPE,
  /// etc., they will typically be of both signs.  Will resize "deriv".
  /// Returns the sum of (GMM likelihood * weight), which may be used
  /// as an approximation to the objective function.
  /// Last two parameters are optional.  See GetStatsDerivative() for
  /// or fMPE paper (ICASSP, 2005) more info on indirect derivative.
  /// Caution: if you supply the last two parameters, this function only
  /// works in the MMI case as it assumes the stats with positive weight
  /// are numerator == ml stats-- this is only the same thing in the MMI
  /// case, not fMPE.
  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 = NULL,
                                     Matrix<BaseFloat> *indirect_deriv = NULL);
  
  
  
  }  // End namespace kaldi
  
  
  #endif