// transform/regtree-mllr-diag-gmm.h

// Copyright 2009-2011  Saarland University;  Jan Silovsky

// 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_MLLR_DIAG_GMM_H_
#define KALDI_TRANSFORM_REGTREE_MLLR_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/common-utils.h"

namespace kaldi {


///  Configuration variables for FMLLR transforms
struct RegtreeMllrOptions {
  BaseFloat min_count;  ///< Minimum occupancy for computing a transform

  /// If 'true', find transforms to generate using regression tree.
  /// If 'false', generate transforms for each baseclass.
  bool use_regtree;

  RegtreeMllrOptions(): min_count(1000.0), use_regtree(true) { }

  void Register(OptionsItf *opts) {
    opts->Register("mllr-min-count", &min_count,
                   "Minimum count to estimate an MLLR transform.");
    opts->Register("mllr-use-regtree", &use_regtree,
                   "Use a regression-class tree for MLLR.");
  }
};

/// An MLLR mean transformation is an affine transformation of Gaussian means.
class RegtreeMllrDiagGmm {
 public:
  RegtreeMllrDiagGmm() {}

  /// Allocates memory for transform matrix & bias vector
  void Init(int32 num_xforms, int32 dim);

  /// Initialize transform matrix to identity and bias vector to zero
  void SetUnit();

  /// Apply the transform(s) to all the Gaussian means in the model
  void TransformModel(const RegressionTree &regtree, AmDiagGmm *am);

  /// Get all the transformed means for a given pdf.
  void GetTransformedMeans(const RegressionTree &regtree, const AmDiagGmm &am,
                           int32 pdf_index, MatrixBase<BaseFloat> *out) const;

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

  /// Mutators
  void SetParameters(const MatrixBase<BaseFloat> &mat, int32 regclass);
  void set_bclass2xforms(const std::vector<int32> &in) { bclass2xforms_ = in; }

  /// Accessors
  const std::vector< Matrix<BaseFloat> > xform_matrices() const {
    return xform_matrices_;
  }

 private:
  /// Transform matrices: size() = num_xforms_
  std::vector< Matrix<BaseFloat> > xform_matrices_;
  int32 num_xforms_;  ///< Number of transforms == xform_matrices_.size()
  /// For each baseclass index of which transform to use; -1 => no xform
  std::vector<int32> bclass2xforms_;
  int32 dim_;  ///< Dimension of feature vectors

  // Cannot have copy constructor and assigment operator
  KALDI_DISALLOW_COPY_AND_ASSIGN(RegtreeMllrDiagGmm);
};

inline void RegtreeMllrDiagGmm::SetParameters(const MatrixBase<BaseFloat> &mat,
                                              int32 regclass) {
  xform_matrices_[regclass].CopyFromMat(mat, kNoTrans);
}

/** Class for computing the maximum-likelihood estimates of the parameters of
 *  an acoustic model that uses diagonal Gaussian mixture models as emission
 *  densities.
 */
class RegtreeMllrDiagGmmAccs {
 public:
  RegtreeMllrDiagGmmAccs() {}
  ~RegtreeMllrDiagGmmAccs() { DeletePointers(&baseclass_stats_); }

  void Init(int32 num_bclass, int32 dim);
  void SetZero();

  /// Accumulate stats for a single GMM in the model; returns log likelihood.
  /// This does not work with multiple feature transforms.
  BaseFloat AccumulateForGmm(const RegressionTree &regtree,
                             const AmDiagGmm &am,
                             const VectorBase<BaseFloat> &data,
                             int32 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,
                             int32 pdf_index, int32 gauss_index,
                             BaseFloat weight);

  void Update(const RegressionTree &regtree, const RegtreeMllrOptions &opts,
              RegtreeMllrDiagGmm *out_mllr, BaseFloat *auxf_impr,
              BaseFloat *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_;
  int32 num_baseclasses_;    ///< Number of baseclasses
  int32 dim_;    ///< Dimension of feature vectors

  /// Returns the MLLR objective function for a given transform and baseclass.
  BaseFloat MllrObjFunction(const Matrix<BaseFloat> &xform,
                            int32 bclass_id) const;

  // Cannot have copy constructor and assigment operator
  KALDI_DISALLOW_COPY_AND_ASSIGN(RegtreeMllrDiagGmmAccs);
};

typedef TableWriter< KaldiObjectHolder<RegtreeMllrDiagGmm> >
            RegtreeMllrDiagGmmWriter;
typedef RandomAccessTableReader< KaldiObjectHolder<RegtreeMllrDiagGmm> >
            RandomAccessRegtreeMllrDiagGmmReader;
typedef RandomAccessTableReaderMapped< KaldiObjectHolder<RegtreeMllrDiagGmm> >
            RandomAccessRegtreeMllrDiagGmmReaderMapped;
typedef SequentialTableReader< KaldiObjectHolder<RegtreeMllrDiagGmm> >
            RegtreeMllrDiagGmmSeqReader;

}  // namespace kaldi

#endif  // KALDI_TRANSFORM_REGTREE_MLLR_DIAG_GMM_H_