// transform/fmllr-diag-gmm.h

// Copyright 2009-2011  Microsoft Corporation;  Saarland University
//                2013  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_FMLLR_DIAG_GMM_H_
#define KALDI_TRANSFORM_FMLLR_DIAG_GMM_H_

#include <vector>

#include "base/kaldi-common.h"
#include "gmm/am-diag-gmm.h"
#include "gmm/mle-full-gmm.h"
#include "transform/transform-common.h"
#include "util/kaldi-table.h"
#include "util/kaldi-holder.h"

namespace kaldi {

/* This header contains routines for performing global CMLLR,
   without a regression tree (however, you can down-weight silence
   in training using the program weight-silence-post on the
   state-level posteriors).  For regression-tree CMLLR, see
   fmllr-diag-gmm.h
*/

struct FmllrOptions {
  std::string update_type;  ///< "full", "diag", "offset", "none"
  BaseFloat min_count;
  int32 num_iters;
  FmllrOptions(): update_type("full"), min_count(500.0), num_iters(40) { }
  void Register(OptionsItf *opts) {
    opts->Register("fmllr-update-type", &update_type,
                   "Update type for fMLLR (\"full\"|\"diag\"|\"offset\"|\"none\")");
    opts->Register("fmllr-min-count", &min_count,
                   "Minimum count required to update fMLLR");
    opts->Register("fmllr-num-iters", &num_iters,
                   "Number of iterations in fMLLR update phase.");
  }
};


/// This does not work with multiple feature transforms.

class FmllrDiagGmmAccs: public AffineXformStats {
 public:
  // If supplied, the "opts" will only be used to limit the
  // stats that are accumulated, to the parts we'll need in the
  // update.
  FmllrDiagGmmAccs(const FmllrOptions &opts = FmllrOptions()):
      opts_(opts) { }
  explicit FmllrDiagGmmAccs(const FmllrDiagGmmAccs &other):
      AffineXformStats(other), single_frame_stats_(other.single_frame_stats_),
      opts_(other.opts_) {}
  explicit FmllrDiagGmmAccs(int32 dim, const FmllrOptions &opts = FmllrOptions()):
      opts_(opts) { Init(dim); }
  
  // The following initializer gives us an efficient way to
  // compute these stats from full-cov Gaussian statistics
  // (accumulated from a *diagonal* model (e.g. use
  // AccumFullGmm::AccumulateFromPosteriors or
  // AccumulateFromDiag).
  FmllrDiagGmmAccs(const DiagGmm &gmm, const AccumFullGmm &fgmm_accs);
  
  void Init(size_t dim) {
    AffineXformStats::Init(dim, dim); single_frame_stats_.Init(dim);
  }
  void Read(std::istream &in, bool binary, bool add) {
      AffineXformStats::Read(in, binary, add);
      single_frame_stats_.Init(Dim());
  }
  /// Accumulate stats for a single GMM in the model; returns log likelihood.
  BaseFloat AccumulateForGmm(const DiagGmm &gmm,
                             const VectorBase<BaseFloat> &data,
                             BaseFloat weight);

  /// This is like AccumulateForGmm but when you have gselect
  /// (Gaussian selection) information
  BaseFloat AccumulateForGmmPreselect(const DiagGmm &gmm,
                                      const std::vector<int32> &gselect,
                                      const VectorBase<BaseFloat> &data,
                                      BaseFloat weight);
  
  /// Accumulate stats for a GMM, given supplied posteriors.
  void AccumulateFromPosteriors(const DiagGmm &gmm,
                                const VectorBase<BaseFloat> &data,
                                const VectorBase<BaseFloat> &posteriors);

  /// Accumulate stats for a GMM, given supplied posteriors.  The "posteriors"
  /// vector should be have the same size as "gselect".
  void AccumulateFromPosteriorsPreselect(
      const DiagGmm &gmm,
      const std::vector<int32> &gselect,
      const VectorBase<BaseFloat> &data,
      const VectorBase<BaseFloat> &posteriors);

  
  /// Update
  void Update(const FmllrOptions &opts,
              MatrixBase<BaseFloat> *fmllr_mat,
              BaseFloat *objf_impr,
              BaseFloat *count);

  // Note: we allow copy and assignment for this class.

  // Note: you can use the inherited AffineXformStats::Read 
  //       and AffineXformStats::Write methods for writing/reading
  //       of the object. It is not necessary to store the other 
  //       private variables of this class
       
 private:
  // The things below, added in 2013, relate to an optimization that lets us
  // speed up accumulation if there are multiple active pdfs per frame
  // (e.g. when accumulating from lattices), or if we don't anticipate
  // doing a "full" update.
  
  struct SingleFrameStats {
    Vector<BaseFloat> x; // dim-dimensional features.
    Vector<BaseFloat> a; // linear term in per-frame auxf; dim is model-dim.
    Vector<BaseFloat> b; // quadratic term in per-frame auxf; dim is model-dim.
    double count;
    SingleFrameStats(int32 dim = 0) { Init(dim); }
    SingleFrameStats(const SingleFrameStats &s): x(s.x), a(s.a), b(s.b),
                                                 count(s.count) {}
    void Init(int32 dim);
  };  

  void CommitSingleFrameStats();

  void InitSingleFrameStats(const VectorBase<BaseFloat> &data);
  
  bool DataHasChanged(const VectorBase<BaseFloat> &data) const; // compares it to the
  // data in single_frame_stats_, returns true if it's different.

  SingleFrameStats single_frame_stats_;
  
  // We only use the opts_ variable for its "update_type" data member,
  // which limits what parts of the G matrix we accumulate.
  FmllrOptions opts_;
  
};


// Initializes the FMLLR matrix to its default values.
inline void InitFmllr(int32 dim,
                      Matrix<BaseFloat> *out_fmllr) {
  out_fmllr->Resize(dim, dim+1);
  out_fmllr->SetUnit();  // sets diagonal elements to one.
}

// ComputeFmllr optimizes the FMLLR matrix, controlled by the options.
// It starts the optimization from the current value of the matrix (e.g. use
// InitFmllr to get this).
// Returns auxf improvement.
BaseFloat ComputeFmllrDiagGmm(const FmllrDiagGmmAccs &accs,
                              const FmllrOptions &opts,
                              Matrix<BaseFloat> *out_fmllr,
                              BaseFloat *logdet);  // add this to likelihoods

inline BaseFloat ComputeFmllrLogDet(const Matrix<BaseFloat> &fmllr_mat) {
  KALDI_ASSERT(fmllr_mat.NumRows() != 0 && fmllr_mat.NumCols() == fmllr_mat.NumRows()+1);
  SubMatrix<BaseFloat> tmp(fmllr_mat,
                           0, fmllr_mat.NumRows(),
                           0, fmllr_mat.NumRows());
  return tmp.LogDet();
}


/// Updates the FMLLR matrix using Mark Gales' row-by-row update.
/// Uses full fMLLR matrix (no structure).  Returns the
/// objective function improvement, not normalized by number of frames.
BaseFloat ComputeFmllrMatrixDiagGmmFull(const MatrixBase<BaseFloat> &in_xform,
                                        const AffineXformStats &stats,
                                        int32 num_iters,
                                        MatrixBase<BaseFloat> *out_xform);

/// This does diagonal fMLLR (i.e. only estimate an offset and scale per
/// dimension).  The format of the output is the same as for the full case.  Of
/// course, these statistics are unnecessarily large for this case.  Returns the
/// objective function improvement, not normalized by number of frames.
BaseFloat ComputeFmllrMatrixDiagGmmDiagonal(const MatrixBase<BaseFloat> &in_xform,
                                            const AffineXformStats &stats,
                                            MatrixBase<BaseFloat> *out_xform);
// Simpler implementation I am testing.
BaseFloat ComputeFmllrMatrixDiagGmmDiagonal2(const MatrixBase<BaseFloat> &in_xform,
                                             const AffineXformStats &stats,
                                             MatrixBase<BaseFloat> *out_xform);

/// This does offset-only fMLLR, i.e. it only estimates an offset.
BaseFloat ComputeFmllrMatrixDiagGmmOffset(const MatrixBase<BaseFloat> &in_xform,
                                          const AffineXformStats &stats,
                                          MatrixBase<BaseFloat> *out_xform);


/// This function internally calls ComputeFmllrMatrixDiagGmm{Full, Diagonal, Offset},
/// depending on "fmllr_type".
BaseFloat ComputeFmllrMatrixDiagGmm(const MatrixBase<BaseFloat> &in_xform,
                                    const AffineXformStats &stats,
                                    std::string fmllr_type,  // "none", "offset", "diag", "full"
                                    int32 num_iters,
                                    MatrixBase<BaseFloat> *out_xform);

/// Returns the (diagonal-GMM) FMLLR auxiliary function value given the transform
/// and the stats.
float FmllrAuxFuncDiagGmm(const MatrixBase<float> &xform,
                          const AffineXformStats &stats);
double FmllrAuxFuncDiagGmm(const MatrixBase<double> &xform,
                           const AffineXformStats &stats);



/// Returns the (diagonal-GMM) FMLLR auxiliary function value given the transform
/// and the stats.
BaseFloat FmllrAuxfGradient(const MatrixBase<BaseFloat> &xform,
                            const AffineXformStats &stats,
                            MatrixBase<BaseFloat> *grad_out);


/// This function applies a feature-level transform to stats (useful for
/// certain techniques based on fMLLR).  Assumes the stats are of the
/// standard diagonal sort.
/// The transform "xform" may be either dim x dim (linear),
/// dim x dim+1 (affine), or dim+1 x dim+1 (affine with the
/// last row equal to 0 0 0 .. 0 1).
void ApplyFeatureTransformToStats(const MatrixBase<BaseFloat> &xform,
                                  AffineXformStats *stats);

/// ApplyModelTransformToStats takes a transform "xform", which must be diagonal
/// (i.e. of the form T = [ D; b ] where D is diagonal), and applies it to the
/// stats as if we had made it a model-space transform (note that the transform
/// applied to the model means is the inverse transform of T).  Thus, if we are
/// estimating a transform T U, and we get stats valid for estimating T U and we
/// estimate T, we can then call this function (treating T as a model-space
/// transform) and will get stats valid for estimating U.  This only works if T is
/// diagonal, because otherwise the standard stats format is not valid.  xform must
/// be of dimension d x d+1
void ApplyModelTransformToStats(const MatrixBase<BaseFloat> &xform,
                                AffineXformStats *stats);


/// This function does one row of the inner-loop fMLLR transform update.
/// We export it because it's needed in the RawFmllr code.
/// Here, if inv_G is the inverse of the G matrix indexed by this row,
/// and k is the corresponding row of the K matrix.
void FmllrInnerUpdate(SpMatrix<double> &inv_G,
                      VectorBase<double> &k,
                      double beta,
                      int32 row,
                      MatrixBase<double> *transform);

                      
                      


} // namespace kaldi

#endif  // KALDI_TRANSFORM_FMLLR_DIAG_GMM_H_