// sgmm2/estimate-am-sgmm2.h

// Copyright 2009-2011  Microsoft Corporation;  Lukas Burget;
//                      Saarland University (Author: Arnab Ghoshal);
//                      Ondrej Glembek;  Yanmin Qian;
// Copyright 2012-2013  Johns Hopkins University (Author: Daniel Povey)
//                      Liang Lu;  Arnab Ghoshal

// 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_SGMM2_ESTIMATE_AM_SGMM2_H_
#define KALDI_SGMM2_ESTIMATE_AM_SGMM2_H_ 1

#include <string>
#include <vector>

#include "sgmm2/am-sgmm2.h"
#include "gmm/model-common.h"
#include "itf/options-itf.h"
#include "util/kaldi-thread.h"

namespace kaldi {

/** \struct MleAmSgmm2Options
 *  Configuration variables needed in the SGMM estimation process.
 */
struct MleAmSgmm2Options {
  /// Smoothing constant for sub-state weights [count to add to each one].
  BaseFloat tau_c;
  /// Floor covariance matrices Sigma_i to this times average cov.
  BaseFloat cov_floor;
  /// ratio to dim below which we use diagonal. default 2, set to inf for diag.
  BaseFloat cov_diag_ratio;
  /// Max on condition of matrices in update beyond which we do not update.
  /// Should probably be related to numerical properties of machine
  /// or BaseFloat type.
  BaseFloat max_cond;

  bool renormalize_V;  // Renormalize the phonetic space.
  bool renormalize_N;  // Renormalize the speaker space.

  /// Number of iters when re-estimating weight projections "w".
  int weight_projections_iters;

  BaseFloat epsilon;  ///< very small value used to prevent SVD crashing.
  BaseFloat max_impr_u; ///< max improvement per frame allowed in update of u.

  BaseFloat tau_map_M;  ///< For MAP update of the phonetic subspace M
  int map_M_prior_iters;  ///< num of iterations to update the prior of M
  bool full_row_cov;  ///< Estimate row covariance instead of using I
  bool full_col_cov;  ///< Estimate col covariance instead of using I

  MleAmSgmm2Options() {
    cov_floor = 0.025;
    tau_c  = 2.0;
    cov_diag_ratio = 2.0;  // set this to very large to get diagonal-cov models.
    max_cond = 1.0e+05;
    epsilon = 1.0e-40;
    renormalize_V = true;
    renormalize_N = false;  // default to false since will invalidate spk vectors
    // on disk.
    weight_projections_iters = 3;
    max_impr_u = 0.25;

    map_M_prior_iters = 5;
    tau_map_M = 0.0;  // No MAP update by default (~500-1000 depending on prior)
    full_row_cov = false;
    full_col_cov = false;
  }

  void Register(OptionsItf *opts) {
    std::string module = "MleAmSgmm2Options: ";
    opts->Register("tau-c", &tau_c, module+
                   "Count for smoothing weight update.");
    opts->Register("cov-floor", &cov_floor, module+
                   "Covariance floor (fraction of average covariance).");
    opts->Register("cov-diag-ratio", &cov_diag_ratio, module+
                   "Minimum occ/dim ratio below which use diagonal covariances.");
    opts->Register("max-cond", &max_cond, module+"Maximum condition number used to "
                   "regularize the solution of certain quadratic auxiliary functions.");
    opts->Register("weight-projections-iters", &weight_projections_iters, module+
                   "Number for iterations for weight projection estimation.");
    opts->Register("renormalize-v", &renormalize_V, module+"If true, renormalize "
                   "the phonetic-subspace vectors to have meaningful sizes.");
    opts->Register("renormalize-n", &renormalize_N, module+"If true, renormalize "
                   "the speaker subspace to have meaningful sizes.");
    opts->Register("max-impr-u", &max_impr_u, module+"Maximum objective function "
                   "improvement per frame allowed in update of u (to "
                   "maintain stability.");

    opts->Register("tau-map-M", &tau_map_M, module+"Smoothing for MAP estimate "
                   "of M (0 means ML update).");
    opts->Register("map-M-prior-iters", &map_M_prior_iters, module+
                   "Number of iterations to estimate prior covariances for M.");
    opts->Register("full-row-cov", &full_row_cov, module+
                   "Estimate row covariance instead of using I.");
    opts->Register("full-col-cov", &full_col_cov, module+
                   "Estimate column covariance instead of using I.");
  }
};

/** \class MleAmSgmm2Accs
 *  Class for the accumulators associated with the phonetic-subspace model
 *  parameters
 */
class MleAmSgmm2Accs {
 public:
  explicit MleAmSgmm2Accs(BaseFloat rand_prune = 1.0e-05)
      : total_frames_(0.0), total_like_(0.0), feature_dim_(0),
        phn_space_dim_(0), spk_space_dim_(0), num_gaussians_(0),
        num_pdfs_(0), num_groups_(0), rand_prune_(rand_prune) {}

  MleAmSgmm2Accs(const AmSgmm2 &model, SgmmUpdateFlagsType flags,
                 bool have_spk_vecs,
                 BaseFloat rand_prune = 1.0e-05)
      : total_frames_(0.0), total_like_(0.0), rand_prune_(rand_prune) {
    ResizeAccumulators(model, flags, have_spk_vecs);
  }

  ~MleAmSgmm2Accs();

  void Read(std::istream &in_stream, bool binary, bool add);
  void Write(std::ostream &out_stream, bool binary) const;

  /// Checks the various accumulators for correct sizes given a model. With
  /// wrong sizes, assertion failure occurs. When the show_properties argument
  /// is set to true, dimensions and presence/absence of the various
  /// accumulators are printed. For use when accumulators are read from file.
  void Check(const AmSgmm2 &model, bool show_properties = true) const;

  /// Resizes the accumulators to the correct sizes given the model. The flags
  /// argument controls which accumulators to resize.
  void ResizeAccumulators(const AmSgmm2 &model, SgmmUpdateFlagsType flags,
                          bool have_spk_vecs);

  /// Returns likelihood.
  BaseFloat Accumulate(const AmSgmm2 &model,
                       const Sgmm2PerFrameDerivedVars &frame_vars,
                       int32 pdf_index, // == j2.
                       BaseFloat weight,
                       Sgmm2PerSpkDerivedVars *spk_vars);

  /// Returns count accumulated (may differ from posteriors.Sum()
  /// due to weight pruning).
  BaseFloat AccumulateFromPosteriors(const AmSgmm2 &model,
                                     const Sgmm2PerFrameDerivedVars &frame_vars,
                                     const Matrix<BaseFloat> &posteriors,
                                     int32 pdf_index, // == j2.
                                     Sgmm2PerSpkDerivedVars *spk_vars);

  /// Accumulates global stats for the current speaker (if applicable).  If
  /// flags contains kSgmmSpeakerProjections (N), or
  /// kSgmmSpeakerWeightProjections (u), must call this after finishing the
  /// speaker's data.
  void CommitStatsForSpk(const AmSgmm2 &model,
                         const Sgmm2PerSpkDerivedVars &spk_vars);

  /// Accessors
  void GetStateOccupancies(Vector<BaseFloat> *occs) const;
  int32 FeatureDim() const { return feature_dim_; }
  int32 PhoneSpaceDim() const { return phn_space_dim_; }
  int32 NumPdfs() const { return num_pdfs_; } // returns J2
  int32 NumGroups() const { return num_groups_; } // returns J1
  int32 NumGauss() const { return num_gaussians_; }

 private:
  /// The stats which are not tied to any state.
  /// Stats Y_{i} for phonetic-subspace projections M; Dim is [I][D][S].
  std::vector< Matrix<double> > Y_;
  /// Stats Z_{i} for speaker-subspace projections N. Dim is [I][D][T].
  std::vector< Matrix<double> > Z_;
  /// R_{i}, quadratic term for speaker subspace estimation. Dim is [I][T][T]
  std::vector< SpMatrix<double> > R_;
  /// S_{i}^{-}, scatter of adapted feature vectors x_{i}(t). Dim is [I][D][D].
  std::vector< SpMatrix<double> > S_;

  /// The SGMM state specific stats.
  /// Statistics y_{jm} for state vectors v_{jm}. dimension is [J1][#mix][S].
  std::vector< Matrix<double> > y_;
  /// Gaussian occupancies gamma_{jmi} for each substate and Gaussian index,
  /// pooled over groups. Dim is [J1][#mix][I].
  std::vector< Matrix<double> > gamma_;

  /// [SSGMM] These a_{jmi} quantities are dimensionally the same
  /// as the gamma quantities.  They're needed to estimate the v_{jm}
  /// and w_i quantities in the symmetric SGMM.  Dimension is [J1][#mix][S]
  std::vector< Matrix<double> > a_;

  /// [SSGMM] each row is one of the t_i quantities in the less-exact
  /// version of the SSGMM update for the speaker weight projections.
  /// Dimension is [I][T]
  Matrix<double> t_;

  /// [SSGMM], this is a per-speaker variable storing the a_i^{(s)}
  /// quantities that we will use in order to compute the non-speaker-
  /// specific quantities [see eqs. 53 and 54 in techreport].  Note:
  /// there is a separate variable a_s_ in class MleSgmm2SpeakerAccs,
  /// which is the same thing but for purposes of computing
  /// the speaker-vector v^{(s)}.
  Vector<double> a_s_;

  /// the U_i quantities from the less-exact version of the SSGMM update for the
  /// speaker weight projections.  Dimension is [I][T][T]
  std::vector<SpMatrix<double> > U_;

  /// Sub-state occupancies gamma_{jm}^{(c)} for each sub-state.  In the
  /// SCTM version of the SGMM, for compactness we store two separate
  /// sets of gamma statistics, one to estimate the v_{jm} quantities
  /// and one to estimate the sub-state weights c_{jm}.
  std::vector< Vector<double> > gamma_c_;

  /// gamma_{i}^{(s)}.  Per-speaker counts for each Gaussian. Dimension is [I]
  /// Needed for stats R_.  This can be viewed as a temporary variable; it
  /// does not form part of the stats that we eventually dump to disk.
  Vector<double> gamma_s_;

  double total_frames_, total_like_;

  /// Dimensionality of various subspaces
  int32 feature_dim_, phn_space_dim_, spk_space_dim_;
  int32 num_gaussians_, num_pdfs_, num_groups_;  ///< Other model specifications

  BaseFloat rand_prune_;

  KALDI_DISALLOW_COPY_AND_ASSIGN(MleAmSgmm2Accs);
  friend class MleAmSgmm2Updater;
  friend class EbwAmSgmm2Updater;
};

/** \class MleAmSgmmUpdater
 *  Contains the functions needed to update the SGMM parameters.
 */
class MleAmSgmm2Updater {
 public:
  explicit MleAmSgmm2Updater(const MleAmSgmm2Options &options)
      : options_(options) {}
  void Reconfigure(const MleAmSgmm2Options &options) {
    options_ = options;
  }

  void Update(const MleAmSgmm2Accs &accs,
              AmSgmm2 *model,
              SgmmUpdateFlagsType flags);

 private:
  friend class UpdateWClass;
  friend class UpdatePhoneVectorsClass;
  friend class EbwEstimateAmSgmm2;

  ///  Compute the Q_i quantities (Eq. 64).
  static void ComputeQ(const MleAmSgmm2Accs &accs,
                       const AmSgmm2 &model,
                       std::vector< SpMatrix<double> > *Q);

  /// Compute the S_means quantities, minus sum: (Y_i M_i^T + M_i Y_I^T).
  static void ComputeSMeans(const MleAmSgmm2Accs &accs,
                            const AmSgmm2 &model,
                            std::vector< SpMatrix<double> > *S_means);
  friend class EbwAmSgmm2Updater;

  MleAmSgmm2Options options_;

  // Called from UpdatePhoneVectors; updates a subset of states
  // (relates to multi-threading).
  void UpdatePhoneVectorsInternal(const MleAmSgmm2Accs &accs,
                                  const std::vector<SpMatrix<double> > &H,
                                  const std::vector<Matrix<double> > &log_a,
                                  AmSgmm2 *model,
                                  double *auxf_impr,
                                  int32 num_threads,
                                  int32 thread_id) const;

  double UpdatePhoneVectors(const MleAmSgmm2Accs &accs,
                            const std::vector<SpMatrix<double> > &H,
                            const std::vector<Matrix<double> > &log_a,
                            AmSgmm2 *model) const;

  double UpdateM(const MleAmSgmm2Accs &accs,
                 const std::vector< SpMatrix<double> > &Q,
                 const Vector<double> &gamma_i,
                 AmSgmm2 *model);

  void RenormalizeV(const MleAmSgmm2Accs &accs, AmSgmm2 *model,
                    const Vector<double> &gamma_i,
                    const std::vector<SpMatrix<double> > &H);

  double UpdateN(const MleAmSgmm2Accs &accs, const Vector<double> &gamma_i,
                 AmSgmm2 *model);
  void RenormalizeN(const MleAmSgmm2Accs &accs, const Vector<double> &gamma_i,
                    AmSgmm2 *model);
  double UpdateVars(const MleAmSgmm2Accs &accs,
                    const std::vector< SpMatrix<double> > &S_means,
                    const Vector<double> &gamma_i,
                    AmSgmm2 *model);
  // Update for the phonetic-subspace weight projections w_i
  double UpdateW(const MleAmSgmm2Accs &accs,
                 const std::vector<Matrix<double> > &log_a,
                 const Vector<double> &gamma_i,
                 AmSgmm2 *model);
  // Update for the speaker-subspace weight projections u_i [SSGMM]
  double UpdateU(const MleAmSgmm2Accs &accs, const Vector<double> &gamma_i,
                 AmSgmm2 *model);

  /// Called, multithreaded, inside UpdateW
  static
  void UpdateWGetStats(const MleAmSgmm2Accs &accs,
                       const AmSgmm2 &model,
                       const Matrix<double> &w,
                       const std::vector<Matrix<double> > &log_a,
                       Matrix<double> *F_i,
                       Matrix<double> *g_i,
                       double *tot_like,
                       int32 num_threads,
                       int32 thread_id);

  double UpdateSubstateWeights(const MleAmSgmm2Accs &accs,
                               AmSgmm2 *model);

  static void ComputeLogA(const MleAmSgmm2Accs &accs,
                          std::vector<Matrix<double> > *log_a); // [SSGMM]

  void ComputeMPrior(AmSgmm2 *model);  // TODO(arnab): Maybe make this static?
  double MapUpdateM(const MleAmSgmm2Accs &accs,
                    const std::vector< SpMatrix<double> > &Q,
                    const Vector<double> &gamma_i, AmSgmm2 *model);

  KALDI_DISALLOW_COPY_AND_ASSIGN(MleAmSgmm2Updater);
  MleAmSgmm2Updater() {}  // Prevent unconfigured updater.
};


/** \class MleSgmm2SpeakerAccs
 *  Class for the accumulators required to update the speaker
 *  vectors v_s.
 *  Note: if you have multiple speakers you will want to initialize
 *  this just once and call Clear() after you're done with each speaker,
 *  rather than creating a new object for each speaker, since the
 *  initialization function does nontrivial work.
 */

class MleSgmm2SpeakerAccs {
 public:
  /// Initialize the object.  Error if speaker subspace not set up.
  MleSgmm2SpeakerAccs(const AmSgmm2 &model,
                      BaseFloat rand_prune_ = 1.0e-05);

  /// Clear the statistics.
  void Clear();

  /// Accumulate statistics.  Returns per-frame log-likelihood.
  BaseFloat Accumulate(const AmSgmm2 &model,
                       const Sgmm2PerFrameDerivedVars &frame_vars,
                       int32 pdf_index,
                       BaseFloat weight,
                       Sgmm2PerSpkDerivedVars *spk_vars);

  /// Accumulate statistics, given posteriors.  Returns total
  /// count accumulated, which may differ from posteriors.Sum()
  /// due to randomized pruning.
  BaseFloat AccumulateFromPosteriors(const AmSgmm2 &model,
                                     const Sgmm2PerFrameDerivedVars &frame_vars,
                                     const Matrix<BaseFloat> &posteriors,
                                     int32 pdf_index,
                                     Sgmm2PerSpkDerivedVars *spk_vars);

  /// Update speaker vector.  If v_s was empty, will assume it started as zero
  /// and will resize it to the speaker-subspace size.
  void Update(const AmSgmm2 &model,
              BaseFloat min_count,  // e.g. 100
              Vector<BaseFloat> *v_s,
              BaseFloat *objf_impr_out,
              BaseFloat *count_out);

 private:
  // Update without speaker-dependent weights (vectors u_i),
  // i.e. not symmetric SGMM (SSGMM)
  void UpdateNoU(Vector<BaseFloat> *v_s,
                 BaseFloat *objf_impr_out,
                 BaseFloat *count_out);
  // Update for SSGMM
  void UpdateWithU(const AmSgmm2 &model,
                   Vector<BaseFloat> *v_s,
                   BaseFloat *objf_impr_out,
                   BaseFloat *count_out);


  /// Statistics for speaker adaptation (vectors), stored per-speaker.
  /// Per-speaker stats for vectors, y^{(s)}. Dimension [T].
  Vector<double> y_s_;
  /// gamma_{i}^{(s)}.  Per-speaker counts for each Gaussian. Dimension is [I]
  Vector<double> gamma_s_;
  /// a_i^{(s)}.  For SSGMM.
  Vector<double> a_s_;

  /// The following variable does not change per speaker, it just
  /// relates to the speaker subspace.
  /// Eq. (82): H_{i}^{spk} = N_{i}^T \Sigma_{i}^{-1} N_{i}
  std::vector< SpMatrix<double> > H_spk_;

  /// N_i^T \Sigma_{i}^{-1}. Needed for y^{(s)}
  std::vector< Matrix<double> > NtransSigmaInv_;

  /// small constant to randomly prune tiny posteriors
  BaseFloat rand_prune_;
};

// This class, used in multi-core implementation of the updates of the "w_i"
// quantities, was previously in estimate-am-sgmm.cc, but is being moved to the
// header so it can be used in estimate-am-sgmm-ebw.cc.  It is responsible for
// computing, in parallel, the F_i and g_i quantities used in the updates of
// w_i.
class UpdateWClass: public MultiThreadable {
 public:
  UpdateWClass(const MleAmSgmm2Accs &accs,
               const AmSgmm2 &model,
               const Matrix<double> &w,
               const std::vector<Matrix<double> > &log_a,
               Matrix<double> *F_i,
               Matrix<double> *g_i,
               double *tot_like):
      accs_(accs), model_(model), w_(w), log_a_(log_a),
      F_i_ptr_(F_i), g_i_ptr_(g_i), tot_like_ptr_(tot_like) {
    tot_like_ = 0.0;
    F_i_.Resize(F_i->NumRows(), F_i->NumCols());
    g_i_.Resize(g_i->NumRows(), g_i->NumCols());
  }

  UpdateWClass(const UpdateWClass &other) :
      MultiThreadable(other),
      accs_(other.accs_), model_(other.model_), w_(other.w_),
      log_a_(other.log_a_), F_i_ptr_(other.F_i_ptr_), g_i_ptr_(other.g_i_ptr_),
      F_i_(other.F_i_), g_i_(other.g_i_), tot_like_ptr_(other.tot_like_ptr_),
      tot_like_(0.0) { }

  ~UpdateWClass() {
    F_i_ptr_->AddMat(1.0, F_i_, kNoTrans);
    g_i_ptr_->AddMat(1.0, g_i_, kNoTrans);
    *tot_like_ptr_ += tot_like_;
  }

  inline void operator() () {
    // Note: give them local copy of the sums we're computing,
    // which will be propagated to the total sums in the destructor.
    MleAmSgmm2Updater::UpdateWGetStats(accs_, model_, w_, log_a_,
                                      &F_i_, &g_i_, &tot_like_,
                                      num_threads_, thread_id_);
  }
 private:
  const MleAmSgmm2Accs &accs_;
  const AmSgmm2 &model_;
  const Matrix<double> &w_;
  const std::vector<Matrix<double> > &log_a_;
  Matrix<double> *F_i_ptr_;
  Matrix<double> *g_i_ptr_;
  Matrix<double> F_i_;
  Matrix<double> g_i_;
  double *tot_like_ptr_;
  double tot_like_;
};


}  // namespace kaldi


#endif  // KALDI_SGMM2_ESTIMATE_AM_SGMM2_H_