// transform/regtree-fmllr-diag-gmm.h
  
  // Copyright 2009-2011  Saarland University;  Georg Stemmer;
  //                      Microsoft Corporation
  
  // 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_REGTREE_FMLLR_DIAG_GMM_H_
  #define KALDI_TRANSFORM_REGTREE_FMLLR_DIAG_GMM_H_
  
  #include <vector>
  
  #include "base/kaldi-common.h"
  #include "gmm/am-diag-gmm.h"
  #include "transform/transform-common.h"
  #include "transform/regression-tree.h"
  #include "util/kaldi-table.h"
  #include "util/kaldi-holder.h"
  
  namespace kaldi {
  
  
  ///  Configuration variables for FMLLR transforms
  struct RegtreeFmllrOptions {
    std::string update_type;  ///< "full", "diag", "offset", "none"
    BaseFloat min_count;  ///< Minimum occupancy for computing a transform
    int32 num_iters;      ///< Number of iterations (if using an iterative update)
    bool use_regtree;     ///< If 'true', find transforms to generate using regression tree.
                          ///< If 'false', generate transforms for each baseclass.
  
    RegtreeFmllrOptions(): update_type("full"), min_count(1000.0),
                           num_iters(10), use_regtree(true) { }
  
    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 to estimate an fMLLR transform.");
      opts->Register("fmllr-num-iters", &num_iters,
                     "Number of fMLLR iterations (if using an iterative update).");
      opts->Register("fmllr-use-regtree", &use_regtree,
                     "Use a regression-class tree for fMLLR.");
    }
  };
  
  
  /** An FMLLR (feature-space MLLR) transformation, also called CMLLR
   *  (constrained MLLR) is an affine transformation of the feature vectors.
   *  This class supports multiple transforms, and a regression tree.
   *  For a single, feature-level transformation see fmllr-diag-gmm-global.h
   *  Note: the "regression classes" are the classes after tree-clustering,
   *  which are smaller in number than the "base classes"  (these correspond
   *  to the leaves of the tree).
   */
  class RegtreeFmllrDiagGmm {
   public:
    RegtreeFmllrDiagGmm() : dim_(-1), num_xforms_(-1), valid_logdet_(false) {}
    explicit RegtreeFmllrDiagGmm(const RegtreeFmllrDiagGmm &other)
        : dim_(other.dim_), num_xforms_(other.num_xforms_),
          xform_matrices_(other.xform_matrices_), logdet_(other.logdet_),
          valid_logdet_(other.valid_logdet_),
          bclass2xforms_(other.bclass2xforms_) {}
    ~RegtreeFmllrDiagGmm() {}
    /// Allocates memory for transform matrix & bias vector
    void Init(size_t num_xforms, size_t dim);
    void Validate();  ///< Checks whether the various parameters are consistent
    /// Sets transform matrix to identity and bias vector to zero
    void SetUnit();
    /// Computes the log-determinant of the Jacobians for each transform
    void ComputeLogDets();
    /// Get the transformed features for each of the transforms.
    void TransformFeature(const VectorBase<BaseFloat> &in,
                          std::vector< Vector<BaseFloat> > *out) const;
    void Write(std::ostream &out_stream, bool binary) const;
    void Read(std::istream &in_stream, bool binary);
  
    /// Accessors
    int32 Dim() const { return dim_; }
    int32 NumBaseClasses() const { return bclass2xforms_.size(); }
    int32 NumRegClasses() const { return num_xforms_; }
    void GetXformMatrix(int32 xform_index, Matrix<BaseFloat> *out) const;
    void GetLogDets(VectorBase<BaseFloat> *out) const;
    int32 Base2RegClass(int32 bclass) const { return bclass2xforms_[bclass]; }
  
    /// Mutators
    void SetParameters(const MatrixBase<BaseFloat> &mat, size_t regclass);
    void set_bclass2xforms(const std::vector<int32> &in) { bclass2xforms_ = in; }
  
   private:
    int32 dim_;             ///< Dimension of feature vectors
    int32 num_xforms_;            ///< Number of transform matrices
    std::vector< Matrix<BaseFloat> > xform_matrices_;  ///< Transform matrices
    Vector<BaseFloat> logdet_;    ///< Log-determinants of the Jacobians
    bool valid_logdet_;           ///< Whether logdets are for current transforms
    /// For each baseclass index of which transform to use; -1 => no xform
    std::vector<int32> bclass2xforms_;
  
    void operator = (const RegtreeFmllrDiagGmm&);  // Disallow assignment operator
  };
  
  inline void RegtreeFmllrDiagGmm::GetXformMatrix(int32 xform_index,
                                                Matrix<BaseFloat> *out) const {
    if (xform_index >= num_xforms_) {
      KALDI_ERR << "Index (" << xform_index << ") out of range [0, "
          << num_xforms_ << "]";
    }
    out->Resize(dim_, dim_ + 1);
    out->CopyFromMat(xform_matrices_[xform_index], kNoTrans);
  }
  
  inline void RegtreeFmllrDiagGmm::SetParameters(const MatrixBase<BaseFloat> &mat,
                                          size_t regclass) {
    xform_matrices_[regclass].CopyFromMat(mat, kNoTrans);
    valid_logdet_ = false;
  }
  
  inline void RegtreeFmllrDiagGmm::GetLogDets(VectorBase<BaseFloat> *out) const {
    KALDI_ASSERT(valid_logdet_ && out->Dim() == logdet_.Dim());
    out->CopyFromVec(logdet_);
  }
  
  typedef TableWriter< KaldiObjectHolder<RegtreeFmllrDiagGmm> >  RegtreeFmllrDiagGmmWriter;
  typedef RandomAccessTableReader< KaldiObjectHolder<RegtreeFmllrDiagGmm> >
              RandomAccessRegtreeFmllrDiagGmmReader;
  typedef RandomAccessTableReaderMapped< KaldiObjectHolder<RegtreeFmllrDiagGmm> >
              RandomAccessRegtreeFmllrDiagGmmReaderMapped;
  typedef SequentialTableReader< KaldiObjectHolder<RegtreeFmllrDiagGmm> >  RegtreeFmllrDiagGmmSeqReader;
  
  /** \class RegtreeFmllrDiagGmmAccs
   *  Class for computing the accumulators needed for the maximum-likelihood
   *  estimate of FMLLR transforms for an acoustic model that uses diagonal
   *  Gaussian mixture models as emission densities.
   */
  class RegtreeFmllrDiagGmmAccs {
   public:
    RegtreeFmllrDiagGmmAccs() : num_baseclasses_(-1), dim_(-1) {}
    ~RegtreeFmllrDiagGmmAccs() { DeletePointers(&baseclass_stats_); }
  
    void Init(size_t num_bclass, size_t dim);
    void SetZero();
  
    /// Accumulate stats for a single GMM in the model; returns log likelihood.
    /// This does not work if the features have already been transformed
    /// with multiple feature transforms (so you can't use use this to
    /// do a 2nd pass of regression-tree fMLLR estimation, which as I write
    /// (Dan, 2016) I'm not sure that this framework even supports.
    BaseFloat AccumulateForGmm(const RegressionTree &regtree,
                               const AmDiagGmm &am,
                               const VectorBase<BaseFloat> &data,
                               size_t pdf_index, BaseFloat weight);
  
    /// Accumulate stats for a single Gaussian component in the model.
    void AccumulateForGaussian(const RegressionTree &regtree,
                               const AmDiagGmm &am,
                               const VectorBase<BaseFloat> &data,
                               size_t pdf_index, size_t gauss_index,
                               BaseFloat weight);
  
    void Update(const RegressionTree &regtree, const RegtreeFmllrOptions &opts,
                RegtreeFmllrDiagGmm *out_fmllr, BaseFloat *auxf_impr,
                BaseFloat *tot_t) const;
  
    void Write(std::ostream &out_stream, bool binary) const;
    void Read(std::istream &in_stream, bool binary, bool add);
  
    /// Accessors
    int32 Dim() const { return dim_; }
    int32 NumBaseClasses() const { return num_baseclasses_; }
    const std::vector<AffineXformStats*> &baseclass_stats() const {
      return baseclass_stats_;
    }
  
   private:
    /// Per-baseclass stats; used for accumulation
    std::vector<AffineXformStats*> baseclass_stats_;
    /// Number of baseclasses
    int32 num_baseclasses_;
    /// Dimension of feature vectors
    int32 dim_;
  
    // Cannot have copy constructor and assigment operator
    KALDI_DISALLOW_COPY_AND_ASSIGN(RegtreeFmllrDiagGmmAccs);
  };
  
  
  
  
  }  // namespace kaldi
  
  #endif  // KALDI_TRANSFORM_REGTREE_FMLLR_DIAG_GMM_H_