// gmm/mle-diag-gmm.h
  
  // Copyright 2009-2012  Saarland University;  Georg Stemmer;
  //                      Microsoft Corporation;  Jan Silovsky; Yanmin Qian
  //                      Johns Hopkins University (author: Daniel Povey)
  //                      Cisco Systems (author: Neha Agrawal)
  
  // 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_GMM_MLE_DIAG_GMM_H_
  #define KALDI_GMM_MLE_DIAG_GMM_H_ 1
  
  #include "gmm/diag-gmm.h"
  #include "gmm/diag-gmm-normal.h"
  #include "gmm/model-common.h"
  #include "itf/options-itf.h"
  
  namespace kaldi {
  
  /** \struct MleDiagGmmOptions
   *  Configuration variables like variance floor, minimum occupancy, etc.
   *  needed in the estimation process.
   */
  struct MleDiagGmmOptions {
    /// Variance floor for each dimension [empty if not supplied].
    /// It is in double since the variance is computed in double precision.
    Vector<double> variance_floor_vector;
    /// Minimum weight below which a Gaussian is not updated (and is
    /// removed, if remove_low_count_gaussians == true);
    BaseFloat min_gaussian_weight;
    /// Minimum count below which a Gaussian is not updated (and is
    /// removed, if remove_low_count_gaussians == true).
    BaseFloat min_gaussian_occupancy;
    /// Minimum allowed variance in any dimension (if no variance floor)
    /// It is in double since the variance is computed in double precision.
    double min_variance;
    bool remove_low_count_gaussians;
    MleDiagGmmOptions() {
      // don't set var floor vector by default.
      min_gaussian_weight     = 1.0e-05;
      min_gaussian_occupancy  = 10.0;
      min_variance            = 0.001;
      remove_low_count_gaussians = true;
    }
    void Register(OptionsItf *opts) {
      std::string module = "MleDiagGmmOptions: ";
      opts->Register("min-gaussian-weight", &min_gaussian_weight,
                   module+"Min Gaussian weight before we remove it.");
      opts->Register("min-gaussian-occupancy", &min_gaussian_occupancy,
                   module+"Minimum occupancy to update a Gaussian.");
      opts->Register("min-variance", &min_variance,
                   module+"Variance floor (absolute variance).");
      opts->Register("remove-low-count-gaussians", &remove_low_count_gaussians,
                   module+"If true, remove Gaussians that fall below the floors.");
    }
  };
  
  
  /** \struct MapDiagGmmOptions
   *  Configuration variables for Maximum A Posteriori (MAP) update.
   */
  struct MapDiagGmmOptions {
    /// Tau value for the means.
    BaseFloat mean_tau;
  
    /// Tau value for the variances.  (Note:
    /// whether or not the variances are updated at all will
    /// be controlled by flags.)
    BaseFloat variance_tau;
  
    /// Tau value for the weights-- this tau value is applied
    /// per state, not per Gaussian.
    BaseFloat weight_tau;
  
    MapDiagGmmOptions(): mean_tau(10.0),
                               variance_tau(50.0),
                               weight_tau(10.0) { }
  
    void Register(OptionsItf *opts) {
      opts->Register("mean-tau", &mean_tau,
                     "Tau value for updating means.");
      opts->Register("variance-tau", &mean_tau,
                     "Tau value for updating variances (note: only relevant if "
                     "update-flags contains \"v\".");
      opts->Register("weight-tau", &weight_tau,
                     "Tau value for updating weights.");
    }
  };
  
  
  
  class AccumDiagGmm {
   public:
    AccumDiagGmm(): dim_(0), num_comp_(0), flags_(0) { }
    explicit AccumDiagGmm(const DiagGmm &gmm, GmmFlagsType flags) {
      Resize(gmm, flags);
    }
    // provide copy constructor.
    explicit AccumDiagGmm(const AccumDiagGmm &other);
  
    void Read(std::istream &in_stream, bool binary, bool add);
    void Write(std::ostream &out_stream, bool binary) const;
  
    /// Allocates memory for accumulators
    void Resize(int32 num_gauss, int32 dim, GmmFlagsType flags);
    /// Calls ResizeAccumulators with arguments based on gmm
    void Resize(const DiagGmm &gmm, GmmFlagsType flags);
  
    /// Returns the number of mixture components
    int32 NumGauss() const { return num_comp_; }
    /// Returns the dimensionality of the feature vectors
    int32 Dim() const { return dim_; }
  
    void SetZero(GmmFlagsType flags);
    void Scale(BaseFloat f, GmmFlagsType flags);
  
    /// Accumulate for a single component, given the posterior
    void AccumulateForComponent(const VectorBase<BaseFloat> &data,
                                int32 comp_index, BaseFloat weight);
  
    /// Accumulate for all components, given the posteriors.
    void AccumulateFromPosteriors(const VectorBase<BaseFloat> &data,
                                  const VectorBase<BaseFloat> &gauss_posteriors);
  
    /// Accumulate for all components given a diagonal-covariance GMM.
    /// Computes posteriors and returns log-likelihood
    BaseFloat AccumulateFromDiag(const DiagGmm &gmm,
                                 const VectorBase<BaseFloat> &data,
                                 BaseFloat frame_posterior);
  
    /// This does the same job as AccumulateFromDiag, but using
    /// multiple threads.  Returns sum of (log-likelihood times
    /// frame weight) over all frames.
    BaseFloat AccumulateFromDiagMultiThreaded(
        const DiagGmm &gmm,
        const MatrixBase<BaseFloat> &data,
        const VectorBase<BaseFloat> &frame_weights,
        int32 num_threads);
  
  
    /// Increment the stats for this component by the specified amount
    /// (not all parts may be taken, depending on flags).
    /// Note: x_stats and x2_stats are assumed to already be multiplied by "occ"
    void AddStatsForComponent(int32 comp_id,
                              double occ,
                              const VectorBase<double> &x_stats,
                              const VectorBase<double> &x2_stats);
  
    /// Increment with stats from this other accumulator (times scale)
    void Add(double scale, const AccumDiagGmm &acc);
  
    /// Smooths the accumulated counts by adding 'tau' extra frames. An example
    /// use for this is I-smoothing for MMIE.   Calls SmoothWithAccum.
    void SmoothStats(BaseFloat tau);
  
    /// Smooths the accumulated counts using some other accumulator. Performs a
    /// weighted sum of the current accumulator with the given one. An example use
    /// for this is I-smoothing for MMI and MPE. Both accumulators must have the
    /// same dimension and number of components.
    void SmoothWithAccum(BaseFloat tau, const AccumDiagGmm &src_acc);
  
    /// Smooths the accumulated counts using the parameters of a given model.
    /// An example use of this is MAP-adaptation. The model must have the
    /// same dimension and number of components as the current accumulator.
    void SmoothWithModel(BaseFloat tau, const DiagGmm &src_gmm);
  
    // Const accessors
    GmmFlagsType Flags() const { return flags_; }
    const VectorBase<double> &occupancy() const { return occupancy_; }
    const MatrixBase<double> &mean_accumulator() const { return mean_accumulator_; }
    const MatrixBase<double> &variance_accumulator() const { return variance_accumulator_; }
  
    // used in testing.
    void AssertEqual(const AccumDiagGmm &other);
   private:
    int32 dim_;
    int32 num_comp_;
    /// Flags corresponding to the accumulators that are stored.
    GmmFlagsType flags_;
  
    Vector<double> occupancy_;
    Matrix<double> mean_accumulator_;
    Matrix<double> variance_accumulator_;
  };
  
  
  /// Returns "augmented" version of flags: e.g. if just updating means, need
  /// weights too.
  GmmFlagsType AugmentGmmFlags(GmmFlagsType f);
  
  
  inline void AccumDiagGmm::Resize(const DiagGmm &gmm, GmmFlagsType flags) {
    Resize(gmm.NumGauss(), gmm.Dim(), flags);
  }
  
  /// for computing the maximum-likelihood estimates of the parameters of
  /// a Gaussian mixture model.
  /// Update using the DiagGmm: exponential form.  Sets, does not increment,
  /// objf_change_out, floored_elements_out and floored_gauss_out.
  void MleDiagGmmUpdate(const MleDiagGmmOptions &config,
                        const AccumDiagGmm &diag_gmm_acc,
                        GmmFlagsType flags,
                        DiagGmm *gmm,
                        BaseFloat *obj_change_out,
                        BaseFloat *count_out,
                        int32 *floored_elements_out = NULL,
                        int32 *floored_gauss_out = NULL,
                        int32 *removed_gauss_out = NULL);
  
  /// Maximum A Posteriori estimation of the model.
  void MapDiagGmmUpdate(const MapDiagGmmOptions &config,
                        const AccumDiagGmm &diag_gmm_acc,
                        GmmFlagsType flags,
                        DiagGmm *gmm,
                        BaseFloat *obj_change_out,
                        BaseFloat *count_out);
  
  /// Calc using the DiagGMM exponential form
  BaseFloat MlObjective(const DiagGmm &gmm,
                        const AccumDiagGmm &diaggmm_acc);
  
  }  // End namespace kaldi
  
  
  #endif  // KALDI_GMM_MLE_DIAG_GMM_H_