// transform/basis-fmllr-diag-gmm.h
  
  // Copyright 2012  Carnegie Mellon University (author: Yajie Miao)
  //           2014  Johns Hopkins University (author: Daniel Povey)
  //           2014  IMSL, PKU-HKUST (Author: Wei Shi)
  
  // 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_BASIS_FMLLR_DIAG_GMM_H_
  #define KALDI_TRANSFORM_BASIS_FMLLR_DIAG_GMM_H_
  
  #include <vector>
  #include <string>
  
  #include "base/kaldi-common.h"
  #include "gmm/am-diag-gmm.h"
  #include "gmm/mle-full-gmm.h"
  #include "gmm/mle-am-diag-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 subspace CMLLR
     (without a regression tree) for diagonal GMM acoustic model.
  
     Refer to Dan Povey's paper for derivations:
     Daniel Povey, Kaisheng Yao. A basis representation of constrained
     MLLR transforms for robust adaptation. Computer Speech and Language,
     volume 26:35–51, 2012.
  */
  
  struct BasisFmllrOptions {
    int32 num_iters;
    BaseFloat size_scale; // how many basis elements we add for each new frame.
    BaseFloat min_count;
    int32 step_size_iters;
    BasisFmllrOptions(): num_iters(10), size_scale(0.2), min_count(50.0), step_size_iters(3) { }
    void Register(OptionsItf *opts) {
      opts->Register("num-iters", &num_iters,
                     "Number of iterations in basis fMLLR update during testing");
      opts->Register("size-scale", &size_scale,
                     "Scale (< 1.0) on speaker occupancy that gives number of "
                     "basis elements.");
      opts->Register("fmllr-min-count", &min_count,
                     "Minimum count required to update fMLLR");
      opts->Register("step-size-iters", &step_size_iters,
                     "Number of iterations in computing step size");
    }
  };
  
  /** \class BasisFmllrAccus
   *  Stats for fMLLR subspace estimation.  This class is only to estimate
   *  the "basis", which is done in training time.  Class BasisFmllrEstimate
   *  contains the functions that are used in test time.  (see the
   *  function BasisFmllrCoefficients()).
   */
  class BasisFmllrAccus {
  
   public:
    BasisFmllrAccus() { }
    explicit BasisFmllrAccus(int32 dim) {
        dim_ = dim;
        beta_ = 0;
        ResizeAccus(dim);
    }
  
    void ResizeAccus(int32 dim);
  
    /// Routines for reading and writing stats
    void Write(std::ostream &out_stream, bool binary) const;
    void Read(std::istream &in_stream, bool binary, bool add = false);
  
    /// Accumulate gradient scatter for one (training) speaker.
    /// To finish the process, we need to traverse the whole training
    /// set. Parallelization works if the speaker list is splitted, and
    /// stats are summed up by setting add=true in BasisFmllrEstimate::
    /// ReadBasis. See section 5.2 of the paper.
    void AccuGradientScatter(const AffineXformStats &spk_stats);
  
    /// Gradient scatter. Dim is [(D+1)*D] [(D+1)*D]
    SpMatrix<BaseFloat> grad_scatter_;
    /// Feature dimension
    int32 dim_;
    /// Occupancy count
    double beta_;
  };
  
  /** \class BasisFmllrEstimate
   *  Estimation functions for basis fMLLR.
   */
  class BasisFmllrEstimate {
  
   public:
    BasisFmllrEstimate(): dim_(0), basis_size_(0) { }
    explicit BasisFmllrEstimate(int32 dim) {
        dim_ = dim; basis_size_ = dim * (dim + 1);
    }
  
    /// Routines for reading and writing fMLLR basis matrices
    void Write(std::ostream &out_stream, bool binary) const;
    void Read(std::istream &in_stream, bool binary);
  
  
    /// Estimate the base matrices efficiently in a Maximum Likelihood manner.
    /// It takes diagonal GMM as argument, which will be used for preconditioner
    /// computation. The total number of bases is fixed to
    /// N = (dim + 1) * dim
    /// Note that SVD is performed in the normalized space. The base matrices
    /// are finally converted back to the unnormalized space.
    void EstimateFmllrBasis(const AmDiagGmm &am_gmm,
                            const BasisFmllrAccus &basis_accus);
  
    /// This function computes the preconditioner matrix, prior to base matrices
    /// estimation. Since the expected values of G statistics are used, it
    /// takes the acoustic model as the argument, rather than the actual
    /// accumulations AffineXformStats
    /// See section 5.1 of the paper.
    void ComputeAmDiagPrecond(const AmDiagGmm &am_gmm,
                              SpMatrix<double> *pre_cond);
  
    int32 Dim() const { return dim_; }
  
    int32 BasisSize() const { return basis_size_; }
  
    /// This function performs speaker adaptation, computing the fMLLR matrix
    /// based on speaker statistics. It takes fMLLR stats as argument.
    /// The basis weights (d_{1}, d_{2}, ..., d_{N}) are also optimized
    /// explicitly. Finally, it returns objective function improvement over
    /// all the iterations, compared with the value at the initial value of
    /// "out_xform" (or the unit transform if not provided).
    /// The coefficients are output to "coefficients" only if the vector is
    /// provided.
    /// See section 5.3 of the paper for more details.
    double ComputeTransform(const AffineXformStats &spk_stats,
                            Matrix<BaseFloat> *out_xform,
                            Vector<BaseFloat> *coefficients,
                            BasisFmllrOptions options) const;
  
   private:
  
    /// Basis matrices. Dim is [T] [D] [D+1]
    /// T is the number of bases
    std::vector< Matrix<BaseFloat> > fmllr_basis_;
    /// Feature dimension
    int32 dim_;
    /// Number of bases D*(D+1)
    int32 basis_size_;
  };
  
  
  } // namespace kaldi
  
  #endif  // KALDI_TRANSFORM_BASIS_FMLLR_DIAG_GMM_H_