// sgmm2/fmllr-sgmm2.h

// Copyright 2009-2012     Saarland University (author: Arnab Ghoshal)
//                         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_FMLLR_SGMM2_H_
#define KALDI_SGMM2_FMLLR_SGMM2_H_

#include <string>
#include <vector>

#include "base/kaldi-common.h"
#include "sgmm2/am-sgmm2.h"
#include "transform/transform-common.h"
#include "util/kaldi-table.h"
#include "util/kaldi-holder.h"
#include "itf/options-itf.h"

namespace kaldi {

/** \struct Sgmm2FmllrConfig
 *  Configuration variables needed in the estimation of FMLLR for SGMMs.
 */
struct Sgmm2FmllrConfig {
  int32 fmllr_iters;  ///< Number of iterations in FMLLR estimation.
  int32 step_iters;  ///< Iterations to find optimal FMLLR step size.
  /// Minimum occupancy count to estimate FMLLR using basis matrices.
  BaseFloat fmllr_min_count_basis;
  /// Minimum occupancy count to estimate FMLLR without basis matrices.
  BaseFloat fmllr_min_count;
  /// Minimum occupancy count to stop using FMLLR bases and switch to
  /// regular FMLLR estimation.
  BaseFloat fmllr_min_count_full;
  /// Number of basis matrices to use for FMLLR estimation. Can only *reduce*
  /// the number of bases present. Overridden by the 'bases_occ_scale' option.
  int32 num_fmllr_bases;
  /// Scale per-speaker count to determine number of CMLLR bases.
  BaseFloat bases_occ_scale;

  Sgmm2FmllrConfig() {
    fmllr_iters = 5;
    step_iters = 10;
    fmllr_min_count_basis = 100.0;
    fmllr_min_count = 1000.0;
    fmllr_min_count_full = 5000.0;
    num_fmllr_bases = 50;
    bases_occ_scale = 0.2;
  }

  void Register(OptionsItf *opts);
};

inline void Sgmm2FmllrConfig::Register(OptionsItf *opts) {
  std::string module = "Sgmm2FmllrConfig: ";
  opts->Register("fmllr-iters", &fmllr_iters, module+
                 "Number of iterations in FMLLR estimation.");
  opts->Register("fmllr-step-iters", &step_iters, module+
                 "Number of iterations to find optimal FMLLR step size.");
  opts->Register("fmllr-min-count-bases", &fmllr_min_count_basis, module+
                 "Minimum occupancy count to estimate FMLLR using basis matrices.");
  opts->Register("fmllr-min-count", &fmllr_min_count, module+
                 "Minimum occupancy count to estimate FMLLR (without bases).");
  opts->Register("fmllr-min-count-full", &fmllr_min_count_full, module+
                 "Minimum occupancy count to stop using basis matrices for FMLLR.");
  opts->Register("fmllr-num-bases", &num_fmllr_bases, module+
                 "Number of FMLLR basis matrices.");
  opts->Register("fmllr-bases-occ-scale", &bases_occ_scale, module+
                 "Scale per-speaker count to determine number of CMLLR bases.");
}


/** \class Sgmm2FmllrGlobalParams
 *  Global adaptation parameters.
 */
class Sgmm2FmllrGlobalParams {
 public:
  void Init(const AmSgmm2 &sgmm, const Vector<BaseFloat> &state_occs);
  void Write(std::ostream &out_stream, bool binary) const;
  void Read(std::istream &in_stream, bool binary);
  bool IsEmpty() const {
    return (pre_xform_.NumRows() == 0 || inv_xform_.NumRows() == 0 ||
            mean_scatter_.Dim() == 0);
  }
  bool HasBasis() const { return fmllr_bases_.size() != 0; }

  /// Pre-transform matrix. Dim is [D][D+1].
  Matrix<BaseFloat> pre_xform_;
  /// Inverse of pre-transform. Dim is [D][D+1].
  Matrix<BaseFloat> inv_xform_;
  /// Diagonal of mean-scatter matrix. Dim is [D]
  Vector<BaseFloat> mean_scatter_;
  /// \tilde{W}_b.  [b][d][d], dim is [B][D][D+1].
  std::vector< Matrix<BaseFloat> > fmllr_bases_;
};

inline void Sgmm2FmllrGlobalParams::Init(const AmSgmm2 &sgmm,
                                        const Vector<BaseFloat> &state_occs) {
  sgmm.ComputeFmllrPreXform(state_occs, &pre_xform_, &inv_xform_,
                            &mean_scatter_);
}

/** \class FmllrSgmm2Accs
 *  Class for computing the accumulators needed for the maximum-likelihood
 *  estimate of FMLLR transforms for a subspace GMM acoustic model.
 */
class FmllrSgmm2Accs {
 public:
  FmllrSgmm2Accs() : dim_(-1) {}
  ~FmllrSgmm2Accs() {}

  void Init(int32 dim, int32 num_gaussians);
  void SetZero() { stats_.SetZero(); }

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

  /// Accumulation routine that computes the Gaussian posteriors and calls
  /// the AccumulateFromPosteriors function with the computed posteriors.
  /// The 'data' argument is not FMLLR-transformed and is needed in addition
  /// to the the 'frame_vars' since the latter only contains a copy of the
  /// transformed feature vector.
  BaseFloat Accumulate(const AmSgmm2 &sgmm,                       
                       const VectorBase<BaseFloat> &data,
                       const Sgmm2PerFrameDerivedVars &frame_vars,
                       int32 state_index,
                       BaseFloat weight,
                       Sgmm2PerSpkDerivedVars *spk);

  void AccumulateFromPosteriors(const AmSgmm2 &sgmm,
                                const Sgmm2PerSpkDerivedVars &spk,
                                const VectorBase<BaseFloat> &data,
                                const std::vector<int32> &gauss_select,
                                const Matrix<BaseFloat> &posteriors,
                                int32 state_index);

  void AccumulateForFmllrSubspace(const AmSgmm2 &sgmm,
                                  const Sgmm2FmllrGlobalParams &fmllr_globals,
                                  SpMatrix<double> *grad_scatter);

  BaseFloat FmllrObjGradient(const AmSgmm2 &sgmm,
                             const Matrix<BaseFloat> &xform,
                             Matrix<BaseFloat> *grad_out,
                             Matrix<BaseFloat> *G_out) const;

  /// Computes the FMLLR transform from the accumulated stats, using the
  /// pre-transforms in fmllr_globals. Expects the transform matrix out_xform
  /// to be initialized to the correct size. Returns true if the transform was
  /// updated (i.e. had enough counts).
  bool Update(const AmSgmm2 &model,
              const Sgmm2FmllrGlobalParams &fmllr_globals,
              const Sgmm2FmllrConfig &opts, Matrix<BaseFloat> *out_xform,
              BaseFloat *frame_count, BaseFloat *auxf_improv) const;

  /// Accessors
  int32 Dim() const { return dim_; }
  const AffineXformStats &stats() const { return stats_; }

 private:
  AffineXformStats stats_;  ///< Accumulated stats
  int32 dim_;  ///< Dimension of feature vectors

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

/// Computes the fMLLR basis matrices given the scatter of the vectorized
/// gradients (eq: B.10). The result is stored in 'fmllr_globals'.
/// The actual number of bases may be less than 'num_fmllr_bases' depending
/// on the feature dimension and number of eigenvalues greater than 'min_eig'.
void EstimateSgmm2FmllrSubspace(const SpMatrix<double> &fmllr_grad_scatter,
                               int32 num_fmllr_bases, int32 feat_dim,
                               Sgmm2FmllrGlobalParams *fmllr_globals,
                               double min_eig = 0.0);

}  // namespace kaldi

#endif  // KALDI_SGMM2_FMLLR_SGMM2_H_