// transform/regtree-fmllr-diag-gmm.cc
  
  // 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.
  
  #include <utility>
  #include <vector>
  using std::vector;
  
  #include "itf/optimizable-itf.h"
  #include "transform/fmllr-diag-gmm.h"
  #include "transform/regtree-fmllr-diag-gmm.h"
  
  namespace kaldi {
  
  void RegtreeFmllrDiagGmm::Init(size_t num_xforms, size_t dim) {
    if (num_xforms == 0) {  // empty transform
      xform_matrices_.clear();
      logdet_.Resize(0);
      valid_logdet_ = false;
      dim_ = 0;  // non-zero dimension is meaningless with empty transform
      num_xforms_ = 0;
    } else {
      KALDI_ASSERT(dim != 0);  // if not empty, dim = 0 is meaningless
      dim_ = dim;
      num_xforms_ = num_xforms;
      xform_matrices_.resize(num_xforms);
      logdet_.Resize(num_xforms);
      vector< Matrix<BaseFloat> >::iterator xform_itr = xform_matrices_.begin(),
          xform_itr_end = xform_matrices_.end();
      for (; xform_itr != xform_itr_end; ++xform_itr) {
        xform_itr->Resize(dim, dim+1);
        xform_itr->SetUnit();
      }
      valid_logdet_ = true;
    }
  }
  
  void RegtreeFmllrDiagGmm::SetUnit() {
    KALDI_ASSERT(num_xforms_ > 0 && dim_ > 0);
    vector< Matrix<BaseFloat> >::iterator xform_itr = xform_matrices_.begin(),
        xform_itr_end = xform_matrices_.end();
    for (; xform_itr != xform_itr_end; ++xform_itr) {
      xform_itr->SetUnit();
    }
  }
  
  void RegtreeFmllrDiagGmm::Validate() {
    if (dim_ < 0 || num_xforms_ < 0) {  // uninitialized case
      KALDI_ERR <<"Do not call Validate() with an uninitialized object (dim = "
                << (dim_) << ", # transforms = " << (num_xforms_);
    } else if (dim_ * num_xforms_ == 0) {  // empty case
      KALDI_ASSERT(num_xforms_ == 0 && dim_ == 0);
      if (xform_matrices_.size() != 0 || logdet_.Dim() != 0) {
        KALDI_ERR << "Number of transforms = " << (xform_matrices_.size())
                  << ", number of log-determinant terms = " << (logdet_.Dim())
                  << ". Expected number = 0";
      }
      return;
    }
  
    // non-empty case: typical usage scenario
    if (xform_matrices_.size() != static_cast<size_t>(num_xforms_)
        || logdet_.Dim() != num_xforms_) {
      KALDI_ERR << "Number of transforms = " << (xform_matrices_.size())
                << ", number of log-determinant terms = " << (logdet_.Dim())
                << ". `Expected number = " << (num_xforms_);
    }
    for (int32 i = 0; i < num_xforms_; i++) {
      if (xform_matrices_[i].NumRows() != dim_ ||
          xform_matrices_[i].NumCols() != (dim_+1)) {
        KALDI_ERR << "For transform " << (i) << ": inconsistent size: rows = "
                  << (xform_matrices_[i].NumRows()) << ", cols = "
                  << xform_matrices_[i].NumCols() << ", dim = " << (dim_);
      }
    }
    if (bclass2xforms_.size() > 0) {
      for (int32 i = 0, maxi = bclass2xforms_.size(); i < maxi; i++) {
        if (bclass2xforms_[i] >= num_xforms_) {
          KALDI_ERR << "For baseclass " << (i) << ", transform index "
                    << (bclass2xforms_[i]) << " exceeds total transforms "
                    << (num_xforms_);
        }
      }
    } else {
      if (num_xforms_ > 1) {
        KALDI_WARN << "Multiple FMLLR transforms found without baseclass info.";
      }
    }
  }
  
  void RegtreeFmllrDiagGmm::ComputeLogDets() {
    logdet_.Resize(num_xforms_);
    for (int32 r = 0; r < num_xforms_; r++) {
      SubMatrix<BaseFloat> tmp_a(xform_matrices_[r], 0, dim_, 0,
                                 dim_);
      logdet_(r) = tmp_a.LogDet();
      KALDI_ASSERT(!KALDI_ISNAN(logdet_(r)));
    }
    valid_logdet_ = true;
  }
  
  void RegtreeFmllrDiagGmm::TransformFeature(const VectorBase<BaseFloat> &in,
                                      vector<Vector<BaseFloat> > *out) const {
    KALDI_ASSERT(out != NULL);
  
    if (xform_matrices_.size() == 0) {  // empty transform
      KALDI_ASSERT(num_xforms_ == 0 && dim_ == 0 && logdet_.Dim() == 0);
      KALDI_WARN << "Asked to apply empty feature transform. Copying instead.";
      out->resize(1);
      (*out)[0].Resize(in.Dim());
      (*out)[0].CopyFromVec(in);
      return;
    } else {
      KALDI_ASSERT(in.Dim() == dim_);
      // if (!valid_logdet_)
      // KALDI_ERR << "Must call ComputeLogDets() before transforming data.";
      // [no need for this check].
      Vector<BaseFloat> extended_feat(dim_ + 1);
      extended_feat.Range(0, dim_).CopyFromVec(in);
      extended_feat(dim_) = 1.0;
      KALDI_ASSERT(num_xforms_ > 0);
      out->resize(num_xforms_);
      for (int32 xform_index = 0; xform_index < num_xforms_;
           ++xform_index) {
        (*out)[xform_index].Resize(dim_);
        (*out)[xform_index].AddMatVec(1.0, xform_matrices_[xform_index],
                                      kNoTrans, extended_feat, 0.0);
      }
    }
  }
  
  void RegtreeFmllrDiagGmm::Write(std::ostream &out, bool binary) const {
    WriteToken(out, binary, "<FMLLRXFORM>");
    WriteToken(out, binary, "<NUMXFORMS>");
    WriteBasicType(out, binary, num_xforms_);
    WriteToken(out, binary, "<DIMENSION>");
    WriteBasicType(out, binary, dim_);
  
    vector< Matrix<BaseFloat> >::const_iterator xform_itr =
        xform_matrices_.begin(), xform_itr_end = xform_matrices_.end();
    for (; xform_itr != xform_itr_end; ++xform_itr) {
      WriteToken(out, binary, "<XFORM>");
      xform_itr->Write(out, binary);
    }
  
    WriteToken(out, binary, "<BCLASS2XFORMS>");
    WriteIntegerVector(out, binary, bclass2xforms_);
    WriteToken(out, binary, "</FMLLRXFORM>");
  }
  
  
  void RegtreeFmllrDiagGmm::Read(std::istream &in, bool binary) {
    ExpectToken(in, binary, "<FMLLRXFORM>");
    ExpectToken(in, binary, "<NUMXFORMS>");
    ReadBasicType(in, binary, &num_xforms_);
    ExpectToken(in, binary, "<DIMENSION>");
    ReadBasicType(in, binary, &dim_);
    KALDI_ASSERT(num_xforms_ >= 0 && dim_ >= 0);  // can be 0 for empty xform
  
    xform_matrices_.resize(num_xforms_);
    vector< Matrix<BaseFloat> >::iterator xform_itr = xform_matrices_.begin(),
        xform_itr_end = xform_matrices_.end();
    for (; xform_itr != xform_itr_end; ++xform_itr) {
      ExpectToken(in, binary, "<XFORM>");
      xform_itr->Read(in, binary);
      KALDI_ASSERT(xform_itr->NumRows() == (xform_itr->NumCols() - 1)
             && xform_itr->NumRows() == dim_);
    }
  
    ExpectToken(in, binary, "<BCLASS2XFORMS>");
    ReadIntegerVector(in, binary, &bclass2xforms_);
    ExpectToken(in, binary, "</FMLLRXFORM>");
    ComputeLogDets();  // so that the transforms can be used.
  }
  
  // ************************************************************************
  
  
  
  
  void RegtreeFmllrDiagGmmAccs::Init(size_t num_bclass, size_t dim) {
    if (num_bclass == 0) {  // empty stats
      DeletePointers(&baseclass_stats_);
      baseclass_stats_.clear();
      num_baseclasses_ = 0;
      dim_ = 0;  // non-zero dimension is meaningless in empty stats
    } else {
      KALDI_ASSERT(dim != 0);  // if not empty, dim = 0 is meaningless
      num_baseclasses_ = num_bclass;
      dim_ = dim;
      DeletePointers(&baseclass_stats_);
      baseclass_stats_.resize(num_bclass);
      for (vector<AffineXformStats*>::iterator it = baseclass_stats_.begin(),
               end = baseclass_stats_.end(); it != end; ++it) {
        *it = new AffineXformStats();
        (*it)->Init(dim, dim);
      }
    }
  }
  
  void RegtreeFmllrDiagGmmAccs::SetZero() {
    for (vector<AffineXformStats*>::iterator it = baseclass_stats_.begin(),
             end = baseclass_stats_.end(); it != end; ++it) {
      (*it)->SetZero();
    }
  }
  
  BaseFloat RegtreeFmllrDiagGmmAccs::AccumulateForGmm(
      const RegressionTree &regtree, const AmDiagGmm &am,
      const VectorBase<BaseFloat> &data, size_t pdf_index, BaseFloat weight) {
    const DiagGmm &pdf = am.GetPdf(pdf_index);
    int32 num_comp = pdf.NumGauss();
    Vector<BaseFloat> posterior(num_comp);
    BaseFloat loglike = pdf.ComponentPosteriors(data, &posterior);
    posterior.Scale(weight);
    Vector<double> posterior_d(posterior);
  
    Vector<double> extended_data(dim_+1);
    extended_data.Range(0, dim_).CopyFromVec(data);
    extended_data(dim_) = 1.0;
    SpMatrix<double> scatter(dim_+1);
    scatter.AddVec2(1.0, extended_data);
  
    Vector<double> inv_var_mean(dim_);
    Matrix<double> g_scale(baseclass_stats_.size(), dim_);  // scale on "scatter" for each dim.
    for (int32 m = 0; m < num_comp; m++) {
      inv_var_mean.CopyRowFromMat(pdf.means_invvars(), m);
      int32 bclass = regtree.Gauss2BaseclassId(pdf_index, m);
  
      baseclass_stats_[bclass]->beta_ += posterior_d(m);
      baseclass_stats_[bclass]->K_.AddVecVec(posterior_d(m), inv_var_mean,
                                             extended_data);
      for (int32 d = 0; d < dim_; d++)
        g_scale(bclass, d) +=  posterior(m) * pdf.inv_vars()(m, d);
    }
    for (size_t bclass = 0; bclass < baseclass_stats_.size(); bclass++) {
      vector< SpMatrix<double> > &G = baseclass_stats_[bclass]->G_;
      for (int32 d = 0; d < dim_; d++)
        if (g_scale(bclass, d) != 0.0)
          G[d].AddSp(g_scale(bclass, d), scatter);
    }
    return loglike;
  }
  
  void RegtreeFmllrDiagGmmAccs::AccumulateForGaussian(
      const RegressionTree &regtree, const AmDiagGmm &am,
      const VectorBase<BaseFloat> &data, size_t pdf_index, size_t gauss_index,
      BaseFloat weight) {
    const DiagGmm &pdf = am.GetPdf(pdf_index);
    size_t dim = static_cast<size_t>(dim_);
    Vector<double> extended_data(dim+1);
    extended_data.Range(0, dim).CopyFromVec(data);
    extended_data(dim) = 1.0;
    SpMatrix<double> scatter(dim+1);
    scatter.AddVec2(1.0, extended_data);
    double weight_d = static_cast<double>(weight);
  
    unsigned bclass = regtree.Gauss2BaseclassId(pdf_index, gauss_index);
    Vector<double> inv_var_mean(dim_);
    inv_var_mean.CopyRowFromMat(pdf.means_invvars(), gauss_index);
  
    baseclass_stats_[bclass]->beta_ += weight_d;
    baseclass_stats_[bclass]->K_.AddVecVec(weight_d, inv_var_mean, extended_data);
    vector< SpMatrix<double> > &G = baseclass_stats_[bclass]->G_;
    for (size_t d = 0; d < dim; d++)
      G[d].AddSp((weight_d * pdf.inv_vars()(gauss_index, d)), scatter);
  }
  
  void RegtreeFmllrDiagGmmAccs::Write(std::ostream &out, bool binary) const {
    WriteToken(out, binary, "<FMLLRACCS>");
    WriteToken(out, binary, "<NUMBASECLASSES>");
    WriteBasicType(out, binary, num_baseclasses_);
    WriteToken(out, binary, "<DIMENSION>");
    WriteBasicType(out, binary, dim_);
    WriteToken(out, binary, "<STATS>");
    vector<AffineXformStats*>::const_iterator itr = baseclass_stats_.begin(),
        end = baseclass_stats_.end();
    for ( ; itr != end; ++itr)
      (*itr)->Write(out, binary);
    WriteToken(out, binary, "</FMLLRACCS>");
  }
  
  void RegtreeFmllrDiagGmmAccs::Read(std::istream &in, bool binary, bool add) {
    ExpectToken(in, binary, "<FMLLRACCS>");
    ExpectToken(in, binary, "<NUMBASECLASSES>");
    ReadBasicType(in, binary, &num_baseclasses_);
    ExpectToken(in, binary, "<DIMENSION>");
    ReadBasicType(in, binary, &dim_);
    KALDI_ASSERT(num_baseclasses_ > 0 && dim_ > 0);
    baseclass_stats_.resize(num_baseclasses_);
    ExpectToken(in, binary, "<STATS>");
    vector<AffineXformStats*>::iterator itr = baseclass_stats_.begin(),
        end = baseclass_stats_.end();
    for ( ; itr != end; ++itr) {
      *itr = new AffineXformStats();
      (*itr)->Init(dim_, dim_);
      (*itr)->Read(in, binary, add);
    }
    ExpectToken(in, binary, "</FMLLRACCS>");
  }
  
  
  void RegtreeFmllrDiagGmmAccs::Update(const RegressionTree &regtree,
                                const RegtreeFmllrOptions &opts,
                                RegtreeFmllrDiagGmm *out_fmllr,
                                BaseFloat *auxf_impr_out,
                                BaseFloat *tot_t_out) const {
    BaseFloat tot_auxf_impr = 0.0, tot_t = 0.0;
    Matrix<BaseFloat> xform_mat(dim_, dim_+1);
    if (opts.use_regtree) {  // estimate transforms using a regression tree
      vector<AffineXformStats*> regclass_stats;
      vector<int32> base2regclass;
      bool update_xforms = regtree.GatherStats(baseclass_stats_, opts.min_count,
                                               &base2regclass, &regclass_stats);
      out_fmllr->set_bclass2xforms(base2regclass);
      // If update_xforms == true, none should be negative, else all should be -1
      if (update_xforms) {
        out_fmllr->Init(regclass_stats.size(), dim_);
        size_t num_rclass = regclass_stats.size();
        for (size_t rclass_index = 0;
             rclass_index < num_rclass; ++rclass_index) {
          KALDI_ASSERT(regclass_stats[rclass_index]->beta_ >= opts.min_count);
          xform_mat.SetUnit();
          tot_t += regclass_stats[rclass_index]->beta_;
  
          tot_auxf_impr +=
              ComputeFmllrMatrixDiagGmmFull(xform_mat, *(regclass_stats[rclass_index]),
                                            opts.num_iters, &xform_mat);
          
          out_fmllr->SetParameters(xform_mat, rclass_index);
        }
        KALDI_LOG << "Estimated " << num_rclass << " regression classes.";
      } else {
        out_fmllr->Init(1, dim_);  // Use a unit transform at the root.
      }
      DeletePointers(&regclass_stats);
      // end of estimation using regression tree
    } else {  // No regtree: estimate 1 transform per baseclass (if enough count)
      for (int32 bclass_index = 0; bclass_index < num_baseclasses_;
           ++bclass_index) {
        tot_t += baseclass_stats_[bclass_index]->beta_;
      }
  
      out_fmllr->Init(num_baseclasses_, dim_);
      vector<int32> base2regclass(num_baseclasses_);
      for (int32 bclass_index = 0; bclass_index < num_baseclasses_;
           ++bclass_index) {
        if (baseclass_stats_[bclass_index]->beta_ >= opts.min_count) {
          xform_mat.SetUnit();
  
          if (opts.update_type == "full") {
            tot_auxf_impr +=
                ComputeFmllrMatrixDiagGmmFull(xform_mat,
                                              *(baseclass_stats_[bclass_index]),
                                              opts.num_iters, &xform_mat);
          } else if (opts.update_type == "diag")
            tot_auxf_impr +=
                ComputeFmllrMatrixDiagGmmDiagonal(xform_mat,
                                                  *(baseclass_stats_[bclass_index]),
                                                  &xform_mat);
          else if (opts.update_type == "offset")
            tot_auxf_impr +=
                ComputeFmllrMatrixDiagGmmOffset(xform_mat,
                                                *(baseclass_stats_[bclass_index]),
                                                &xform_mat);
          else if (opts.update_type == "none")
            tot_auxf_impr = 0.0;
          else
            KALDI_ERR << "Unknown fMLLR update type " << opts.update_type
                      << ", fmllr-update-type must be one of \"full\"|\"diag\"|\"offset\"|\"none\"";
  
          out_fmllr->SetParameters(xform_mat, bclass_index);
          base2regclass[bclass_index] = bclass_index;
        } else {
          KALDI_WARN << "For baseclass " << (bclass_index) << " count = "
                     << (baseclass_stats_[bclass_index]->beta_) << " < "
                     << opts.min_count << ": not updating FMLLR";
          base2regclass[bclass_index] = -1;
        }
        out_fmllr->set_bclass2xforms(base2regclass);
      }  // end looping over all baseclasses
    }  // end of estimating one transform per baseclass without regtree
    if (auxf_impr_out) *auxf_impr_out = tot_auxf_impr;
    if (tot_t_out) *tot_t_out = tot_t;
  }
  
  
  
  
  }  // namespace kaldi