// sgmm2/estimate-am-sgmm2-ebw.h
  
  // Copyright 2012  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_SGMM2_ESTIMATE_AM_SGMM2_EBW_H_
  #define KALDI_SGMM2_ESTIMATE_AM_SGMM2_EBW_H_ 1
  
  #include <string>
  #include <vector>
  
  #include "gmm/model-common.h"
  #include "itf/options-itf.h"
  #include "sgmm2/estimate-am-sgmm2.h"
  
  namespace kaldi {
  
  /**
     This header implements a form of Extended Baum-Welch training for SGMMs.
     If you are confused by this comment, see Dan Povey's thesis for an explanation of
     Extended Baum-Welch.
     A note on the EBW (Extended Baum-Welch) updates for the SGMMs... In general there is
     a parameter-specific value D that is similar to the D in EBW for GMMs.  The value of
     D is generally set to:
       E * (denominator-count for that parameter)   +   tau-value for that parameter
     where the tau-values are user-specified parameters that are specific to the type of
     the parameter (e.g. phonetic vector, subspace projection, etc.).  Things are a bit
     more complex for this update than for GMMs, because it's not just a question of picking
     a tau-value for smoothing: there is sometimes a scatter-matrix of some kind (e.g.
     an outer product of vectors, or something) that defines a quadratic objective function
     that we'll add as smoothing.  We have to pick where to get this scatter-matrix from.
     We feel that it's appropriate for the "E" part of the D to get its scatter-matrix from
     denominator stats, and the tau part of the D to get half its scatter-matrix from the
     both the numerator and denominator stats, assigned a weight proportional to how much
     stats there were.  When you see the auxiliary function written out, it's clear why this
     makes sense.
  
   */
  
  struct EbwAmSgmm2Options {
    BaseFloat tau_v; ///<  Smoothing constant for updates of sub-state vectors v_{jm}
    BaseFloat lrate_v; ///< Learning rate used in updating v-- default 0.5
    BaseFloat tau_M; ///<  Smoothing constant for the M quantities (phone-subspace projections)
    BaseFloat lrate_M; ///< Learning rate used in updating M-- default 0.5
    BaseFloat tau_N; ///<  Smoothing constant for the N quantities (speaker-subspace projections)
    BaseFloat lrate_N; ///< Learning rate used in updating N-- default 0.5
    BaseFloat tau_c;  ///< Tau value for smoothing substate weights (c)
    BaseFloat tau_w;  ///< Tau value for smoothing update of phonetic-subspace weight projectsions (w)
    BaseFloat lrate_w; ///< Learning rate used in updating w-- default 1.0
    BaseFloat tau_u;  ///< Tau value for smoothing update of speaker-subspace weight projectsions (u)
    BaseFloat lrate_u; ///< Learning rate used in updating u-- default 1.0
    BaseFloat max_impr_u; ///< Maximum improvement/frame allowed for u [0.25, carried over from ML update.]
    BaseFloat tau_Sigma; ///< Tau value for smoothing covariance-matrices Sigma.
    BaseFloat lrate_Sigma; ///< Learning rate used in updating Sigma-- default 0.5
    BaseFloat min_substate_weight; ///< Minimum allowed weight in a sub-state.
    
    BaseFloat cov_min_value; ///< E.g. 0.5-- the maximum any eigenvalue of a covariance
    /// is allowed to change.  [this is the minimum; the maximum is the inverse of this,
    /// i.e. 2.0 in this case.  For example, 0.9 would constrain the covariance quite tightly,
    /// 0.1 would be a loose setting.
    
    BaseFloat max_cond; ///< large value used in SolveQuadraticProblem.
    BaseFloat epsilon;  ///< very small value used in SolveQuadraticProblem; workaround
    /// for an issue in some implementations of SVD.
    
    EbwAmSgmm2Options() {
      tau_v = 50.0;
      lrate_v = 0.5;
      tau_M = 500.0;
      lrate_M = 0.5;
      tau_N = 500.0;
      lrate_N = 0.5;
      tau_c = 10.0;
      tau_w = 50.0;
      lrate_w = 1.0;
      tau_u = 50.0;
      lrate_u = 1.0;
      max_impr_u = 0.25;
      tau_Sigma = 500.0;
      lrate_Sigma = 0.5;
  
      min_substate_weight = 1.0e-05;
      cov_min_value = 0.5;
      
      max_cond = 1.0e+05;
      epsilon = 1.0e-40;
    }
  
    void Register(OptionsItf *opts) {
      std::string module = "EbwAmSgmm2Options: ";
      opts->Register("tau-v", &tau_v, module+
                     "Smoothing constant for phone vector estimation.");
      opts->Register("lrate-v", &lrate_v, module+
                     "Learning rate constant for phone vector estimation.");
      opts->Register("tau-m", &tau_M, module+
                     "Smoothing constant for estimation of phonetic-subspace projections (M).");
      opts->Register("lrate-m", &lrate_M, module+
                     "Learning rate constant for phonetic-subspace projections.");
      opts->Register("tau-n", &tau_N, module+
                     "Smoothing constant for estimation of speaker-subspace projections (N).");
      opts->Register("lrate-n", &lrate_N, module+
                     "Learning rate constant for speaker-subspace projections.");
      opts->Register("tau-c", &tau_c, module+
                     "Smoothing constant for estimation of substate weights (c)");
      opts->Register("tau-w", &tau_w, module+
                     "Smoothing constant for estimation of phonetic-space weight projections (w)");
      opts->Register("lrate-w", &lrate_w, module+
                     "Learning rate constant for phonetic-space weight-projections (w)");
      opts->Register("tau-u", &tau_u, module+
                     "Smoothing constant for estimation of speaker-space weight projections (u)");
      opts->Register("lrate-u", &lrate_u, module+
                     "Learning rate constant for speaker-space weight-projections (u)");
      opts->Register("tau-sigma", &tau_Sigma, module+
                     "Smoothing constant for estimation of within-class covariances (Sigma)");
      opts->Register("lrate-sigma", &lrate_Sigma, module+
                     "Constant that controls speed of learning for variances (larger->slower)");
      opts->Register("cov-min-value", &cov_min_value, module+
                     "Minimum value that an eigenvalue of the updated covariance matrix can take, "
                     "relative to its old value (maximum is inverse of this.)");
      opts->Register("min-substate-weight", &min_substate_weight, module+
                     "Floor for weights of sub-states.");
      opts->Register("max-cond", &max_cond, module+
                     "Value used in handling singular matrices during update.");
      opts->Register("epsilon", &max_cond, module+
                     "Value used in handling singular matrices during update.");
    }
  };
  
  
  /** \class EbwAmSgmmUpdater
   *  Contains the functions needed to update the SGMM parameters.
   */
  class EbwAmSgmm2Updater {
   public:
    explicit EbwAmSgmm2Updater(const EbwAmSgmm2Options &options):
        options_(options) {}
    
    void Update(const MleAmSgmm2Accs &num_accs,
                const MleAmSgmm2Accs &den_accs,
                AmSgmm2 *model,
                SgmmUpdateFlagsType flags,
                BaseFloat *auxf_change_out,
                BaseFloat *count_out);
      
   protected:
    // The following two classes relate to multi-core parallelization of some
    // phases of the update.
    friend class EbwUpdateWClass;
    friend class EbwUpdatePhoneVectorsClass;
   private:
    EbwAmSgmm2Options options_;
  
    Vector<double> gamma_j_;  ///< State occupancies
  
    double UpdatePhoneVectors(const MleAmSgmm2Accs &num_accs,
                              const MleAmSgmm2Accs &den_accs,
                              const std::vector< SpMatrix<double> > &H,
                              AmSgmm2 *model) const;
    
    // Called from UpdatePhoneVectors; updates a subset of states
    // (relates to multi-threading).
    void UpdatePhoneVectorsInternal(const MleAmSgmm2Accs &num_accs,
                                    const MleAmSgmm2Accs &den_accs,
                                    const std::vector<SpMatrix<double> > &H,
                                    AmSgmm2 *model,
                                    double *auxf_impr,
                                    int32 num_threads,
                                    int32 thread_id) const;
    // Called from UpdatePhoneVectorsInternal
    static void ComputePhoneVecStats(const MleAmSgmm2Accs &accs,
                                     const AmSgmm2 &model,
                                     const std::vector<SpMatrix<double> > &H,
                                     int32 j1,
                                     int32 m,
                                     const Vector<double> &w_jm,
                                     double gamma_jm,
                                     Vector<double> *g_jm,
                                     SpMatrix<double> *H_jm);
                                      
    double UpdateM(const MleAmSgmm2Accs &num_accs,
                   const MleAmSgmm2Accs &den_accs,
                   const std::vector< SpMatrix<double> > &Q_num,
                   const std::vector< SpMatrix<double> > &Q_den,
                   const Vector<double> &gamma_num,
                   const Vector<double> &gamma_den,
                   AmSgmm2 *model) const;
    
    double UpdateN(const MleAmSgmm2Accs &num_accs,
                   const MleAmSgmm2Accs &den_accs,
                   const Vector<double> &gamma_num,
                   const Vector<double> &gamma_den,
                   AmSgmm2 *model) const;
    
    double UpdateVars(const MleAmSgmm2Accs &num_accs,
                      const MleAmSgmm2Accs &den_accs,
                      const Vector<double> &gamma_num,
                      const Vector<double> &gamma_den,
                      const std::vector< SpMatrix<double> > &S_means,
                      AmSgmm2 *model) const;
  
    /// Note: in the discriminative case we do just one iteration of
    /// updating the w quantities.
    double UpdateW(const MleAmSgmm2Accs &num_accs,
                   const MleAmSgmm2Accs &den_accs,
                   const Vector<double> &gamma_num,
                   const Vector<double> &gamma_den,
                   AmSgmm2 *model);
  
  
    double UpdateU(const MleAmSgmm2Accs &num_accs,
                   const MleAmSgmm2Accs &den_accs,
                   const Vector<double> &gamma_num,
                   const Vector<double> &gamma_den,
                   AmSgmm2 *model);
    
    double UpdateSubstateWeights(const MleAmSgmm2Accs &num_accs,
                                 const MleAmSgmm2Accs &den_accs,
                                 AmSgmm2 *model);
  
    KALDI_DISALLOW_COPY_AND_ASSIGN(EbwAmSgmm2Updater);
    EbwAmSgmm2Updater() {}  // Prevent unconfigured updater.
  };
  
  
  }  // namespace kaldi
  
  
  #endif  // KALDI_SGMM2_ESTIMATE_AM_SGMM2_EBW_H_