// transform/regtree-mllr-diag-gmm.cc
  
  // 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.
  
  #include <utility>
  using std::pair;
  #include <vector>
  using std::vector;
  
  #include "transform/regtree-mllr-diag-gmm.h"
  
  namespace kaldi {
  
  void RegtreeMllrDiagGmm::Init(int32 num_xforms, int32 dim) {
    if (num_xforms == 0) {  // empty transform
      xform_matrices_.clear();
      dim_ = 0;  // non-zero dimension is meaningless with empty transform
      num_xforms_ = 0;
      bclass2xforms_.clear();
    } else {
      KALDI_ASSERT(dim != 0);  // if not empty, dim = 0 is meaningless
      dim_ = dim;
      num_xforms_ = num_xforms;
      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) {
        xform_itr->Resize(dim, dim+1);
        xform_itr->SetUnit();
      }
    }
  }
  
  void RegtreeMllrDiagGmm::SetUnit() {
    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 RegtreeMllrDiagGmm::TransformModel(const RegressionTree &regtree,
                                          AmDiagGmm *am) {
    KALDI_ASSERT(static_cast<int32>(bclass2xforms_.size()) ==
                 regtree.NumBaseclasses());
    Vector<BaseFloat> extended_mean(dim_+1), xformed_mean(dim_);
    for (int32 bclass_index = 0, num_bclasses = regtree.NumBaseclasses();
         bclass_index < num_bclasses; ++bclass_index) {
      int32 xform_index;
      if ((xform_index = bclass2xforms_[bclass_index]) > -1) {
        KALDI_ASSERT(xform_index < num_xforms_);
        const vector< pair<int32, int32> > &bclass =
            regtree.GetBaseclass(bclass_index);
        for (vector< pair<int32, int32> >::const_iterator itr = bclass.begin(),
            end = bclass.end(); itr != end; ++itr) {
          SubVector<BaseFloat> tmp_mean(extended_mean.Range(0, dim_));
          am->GetGaussianMean(itr->first, itr->second, &tmp_mean);
          extended_mean(dim_) = 1.0;
          xformed_mean.AddMatVec(1.0, xform_matrices_[xform_index], kNoTrans,
                                 extended_mean, 0.0);
          am->SetGaussianMean(itr->first, itr->second, xformed_mean);
        }  // end iterating over Gaussians in baseclass
      }  // else keep the means untransformed
    }  // end iterating over all baseclasses
    am->ComputeGconsts();
  }
  
  
  void RegtreeMllrDiagGmm::GetTransformedMeans(const RegressionTree &regtree,
                                               const AmDiagGmm &am,
                                               int32 pdf_index,
                                               MatrixBase<BaseFloat> *out) const {
    KALDI_ASSERT(static_cast<int32>(bclass2xforms_.size()) ==
                 regtree.NumBaseclasses());
    int32 num_gauss = am.GetPdf(pdf_index).NumGauss();
    KALDI_ASSERT(out->NumRows() == num_gauss && out->NumCols() == dim_);
  
    Vector<BaseFloat> extended_mean(dim_+1);
    extended_mean(dim_) = 1.0;
  
    for (int32 gauss_index = 0; gauss_index < num_gauss; gauss_index++) {
      int32 bclass_index = regtree.Gauss2BaseclassId(pdf_index, gauss_index);
      int32 xform_index = bclass2xforms_[bclass_index];
      if (xform_index > -1) {  // use a transform
        KALDI_ASSERT(xform_index < num_xforms_);
        SubVector<BaseFloat> tmp_mean(extended_mean.Range(0, dim_));
        am.GetGaussianMean(pdf_index, gauss_index, &tmp_mean);
        SubVector<BaseFloat> out_row(out->Row(gauss_index));
        out_row.AddMatVec(1.0, xform_matrices_[xform_index], kNoTrans,
                          extended_mean, 0.0);
      } else {  // Copy untransformed mean
        SubVector<BaseFloat> out_row(out->Row(gauss_index));
        am.GetGaussianMean(pdf_index, gauss_index, &out_row);
      }
    }
  }
  
  
  void RegtreeMllrDiagGmm::Write(std::ostream &out, bool binary) const {
    WriteToken(out, binary, "<MLLRXFORM>");
    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, "</MLLRXFORM>");
  }
  
  
  void RegtreeMllrDiagGmm::Read(std::istream &in, bool binary) {
    ExpectToken(in, binary, "<MLLRXFORM>");
    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, "</MLLRXFORM>");
  }
  
  // ************************************************************************
  
  void RegtreeMllrDiagGmmAccs::Init(int32 num_bclass, int32 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;
      baseclass_stats_.resize(num_baseclasses_);
      for (vector<AffineXformStats*>::iterator it = baseclass_stats_.begin(),
          end = baseclass_stats_.end(); it != end; ++it) {
        *it = new AffineXformStats();
        (*it)->Init(dim_, dim_);
      }
    }
  }
  
  void RegtreeMllrDiagGmmAccs::SetZero() {
    for (vector<AffineXformStats*>::iterator it = baseclass_stats_.begin(),
        end = baseclass_stats_.end(); it != end; ++it) {
      (*it)->SetZero();
    }
  }
  
  BaseFloat RegtreeMllrDiagGmmAccs::AccumulateForGmm(
      const RegressionTree &regtree, const AmDiagGmm &am,
      const VectorBase<BaseFloat> &data, int32 pdf_index, BaseFloat weight) {
    const DiagGmm &pdf = am.GetPdf(pdf_index);
    int32 num_comp = static_cast<int32>(pdf.NumGauss());
    Vector<BaseFloat> posterior(num_comp);
    BaseFloat loglike = pdf.ComponentPosteriors(data, &posterior);
    posterior.Scale(weight);
    Vector<double> posterior_d(posterior);
  
    Vector<double> data_d(data);
    Vector<double> inv_var_x(dim_);
    Vector<double> extended_mean(dim_+1);
    SpMatrix<double> mean_scatter(dim_+1);
  
    for (int32 m = 0; m < num_comp; m++) {
      unsigned bclass = regtree.Gauss2BaseclassId(pdf_index, m);
      inv_var_x.CopyFromVec(pdf.inv_vars().Row(m));
      inv_var_x.MulElements(data_d);
  
      // Using SubVector to stop compiler warning
      SubVector<double> tmp_mean(extended_mean, 0, dim_);
      pdf.GetComponentMean(m, &tmp_mean);  // modifies extended_mean
      extended_mean(dim_) = 1.0;
      mean_scatter.SetZero();
      mean_scatter.AddVec2(1.0, extended_mean);
  
      baseclass_stats_[bclass]->beta_ += posterior_d(m);
      baseclass_stats_[bclass]->K_.AddVecVec(posterior_d(m), inv_var_x,
                                             extended_mean);
      vector< SpMatrix<double> > &G = baseclass_stats_[bclass]->G_;
      for (int32 d = 0; d < dim_; d++)
        G[d].AddSp((posterior_d(m) * pdf.inv_vars()(m, d)), mean_scatter);
    }
    return loglike;
  }
  
  void RegtreeMllrDiagGmmAccs::AccumulateForGaussian(
      const RegressionTree &regtree, const AmDiagGmm &am,
      const VectorBase<BaseFloat> &data, int32 pdf_index, int32 gauss_index,
      BaseFloat weight) {
    const DiagGmm &pdf = am.GetPdf(pdf_index);
    Vector<double> data_d(data);
    Vector<double> inv_var_x(dim_);
    Vector<double> extended_mean(dim_+1);
    double weight_d = static_cast<double>(weight);
  
    unsigned bclass = regtree.Gauss2BaseclassId(pdf_index, gauss_index);
    inv_var_x.CopyFromVec(pdf.inv_vars().Row(gauss_index));
    inv_var_x.MulElements(data_d);
  
    // Using SubVector to stop compiler warning
    SubVector<double> tmp_mean(extended_mean, 0, dim_);
    pdf.GetComponentMean(gauss_index, &tmp_mean);  // modifies extended_mean
    extended_mean(dim_) = 1.0;
    SpMatrix<double> mean_scatter(dim_+1);
    mean_scatter.AddVec2(1.0, extended_mean);
  
    baseclass_stats_[bclass]->beta_ += weight_d;
    baseclass_stats_[bclass]->K_.AddVecVec(weight_d, inv_var_x, extended_mean);
    vector< SpMatrix<double> > &G = baseclass_stats_[bclass]->G_;
    for (int32 d = 0; d < dim_; d++)
      G[d].AddSp((weight_d * pdf.inv_vars()(gauss_index, d)), mean_scatter);
  }
  
  void RegtreeMllrDiagGmmAccs::Write(std::ostream &out, bool binary) const {
    WriteToken(out, binary, "<MLLRACCS>");
    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, "</MLLRACCS>");
  }
  
  void RegtreeMllrDiagGmmAccs::Read(std::istream &in, bool binary, bool add) {
    ExpectToken(in, binary, "<MLLRACCS>");
    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, "</MLLRACCS>");
  }
  
  static void ComputeMllrMatrix(const Matrix<double> &K,
                                const vector< SpMatrix<double> > &G,
                                Matrix<BaseFloat> *out) {
    int32 dim = G.size();
    Matrix<double> tmp_out(dim, dim+1);
    for (int32 d = 0; d < dim; d++) {
      if (G[d].Cond() > 1.0e+9) {
        KALDI_WARN << "Dim " << d << ": Badly conditioned stats. Setting MLLR "
                   << "transform to unit.";
        tmp_out.SetUnit();
        break;
      }
      SpMatrix<double> inv_g(G[d]);
  //    KALDI_LOG << "Dim " << d << ": G: max = " << inv_g.Max() << ", min = "
  //              << inv_g.Min() << ", log det = " << inv_g.LogDet(NULL)
  //              << ", cond = " << inv_g.Cond();
      inv_g.Invert();
  //    KALDI_LOG << "Inv G: max = " << inv_g.Max() << ", min = " << inv_g.Min()
  //              << ", log det = " << inv_g.LogDet(NULL) << ", cond = "
  //              << inv_g.Cond();
      tmp_out.Row(d).AddSpVec(1.0, inv_g, K.Row(d), 0.0);
    }
    out->CopyFromMat(tmp_out, kNoTrans);
  }
  
  static BaseFloat MllrAuxFunction(const Matrix<BaseFloat> &xform,
                                   const AffineXformStats &stats) {
    int32 dim = stats.G_.size();
    Matrix<double> xform_d(xform);
    Vector<double> xform_row_g(dim + 1);
    SubMatrix<double> A(xform_d, 0, dim, 0, dim);
    double obj = TraceMatMat(xform_d, stats.K_, kTrans);
    for (int32 d = 0; d < dim; d++) {
      xform_row_g.AddSpVec(1.0, stats.G_[d], xform_d.Row(d), 0.0);
      obj -= 0.5 * VecVec(xform_row_g, xform_d.Row(d));
    }
    return obj;
  }
  
  void RegtreeMllrDiagGmmAccs::Update(const RegressionTree &regtree,
                                      const RegtreeMllrOptions &opts,
                                      RegtreeMllrDiagGmm *out_mllr,
                                      BaseFloat *auxf_impr,
                                      BaseFloat *t) const {
    BaseFloat tot_auxf_impr = 0, tot_t = 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_mllr->set_bclass2xforms(base2regclass);
      // If update_xforms == true, none should be negative, else all should be -1
      if (update_xforms) {
        out_mllr->Init(regclass_stats.size(), dim_);
        for (int32 rclass_index = 0, num_rclass = regclass_stats.size();
             rclass_index < num_rclass; ++rclass_index) {
          KALDI_ASSERT(regclass_stats[rclass_index]->beta_ >= opts.min_count);
          xform_mat.SetUnit();
          BaseFloat obj_old = MllrAuxFunction(xform_mat,
                                              *(regclass_stats[rclass_index]));
          ComputeMllrMatrix(regclass_stats[rclass_index]->K_,
                            regclass_stats[rclass_index]->G_, &xform_mat);
          out_mllr->SetParameters(xform_mat, rclass_index);
          BaseFloat obj_new = MllrAuxFunction(xform_mat,
                                              *(regclass_stats[rclass_index]));
          KALDI_LOG << "MLLR: regclass " << (rclass_index)
                    << ": Objective function impr per frame is "
                    << ((obj_new - obj_old)/regclass_stats[rclass_index]->beta_)
                    << " over " << regclass_stats[rclass_index]->beta_
                    << " frames.";
          KALDI_ASSERT(obj_new >= obj_old - (std::abs(obj_new)+std::abs(obj_old))*1.0e-05);
          tot_t += regclass_stats[rclass_index]->beta_;
          tot_auxf_impr += obj_new - obj_old;
        }
      } else {
        out_mllr->Init(1, dim_);  // Use a unit transform at the root.
      }
      DeletePointers(&regclass_stats);
      // end of estimation using regression tree
    } else {  // estimate 1 transform per baseclass (if enough count)
      out_mllr->Init(num_baseclasses_, dim_);
      vector<int32> base2xforms(num_baseclasses_, -1);
      for (int32 bclass_index = 0; bclass_index < num_baseclasses_;
           ++bclass_index) {
        if (baseclass_stats_[bclass_index]->beta_ > opts.min_count) {
          base2xforms[bclass_index] = bclass_index;
          xform_mat.SetUnit();
          BaseFloat obj_old = MllrAuxFunction(xform_mat,
                                              *(baseclass_stats_[bclass_index]));
          ComputeMllrMatrix(baseclass_stats_[bclass_index]->K_,
                            baseclass_stats_[bclass_index]->G_, &xform_mat);
          out_mllr->SetParameters(xform_mat, bclass_index);
          BaseFloat obj_new = MllrAuxFunction(xform_mat,
                                              *(baseclass_stats_[bclass_index]));
          KALDI_LOG << "MLLR: base-class " << (bclass_index)
                    << ": Auxiliary function impr per frame is "
                    << ((obj_new-obj_old)/baseclass_stats_[bclass_index]->beta_);
          KALDI_ASSERT(obj_new >= obj_old - (std::abs(obj_new)+std::abs(obj_old))*1.0e-05);
          tot_t += baseclass_stats_[bclass_index]->beta_;
          tot_auxf_impr += obj_new - obj_old;
        } else {
          KALDI_WARN << "For baseclass "  << (bclass_index) << " count = "
                     << (baseclass_stats_[bclass_index]->beta_) << " < "
                     << opts.min_count << ": not updating MLLR";
          tot_t += baseclass_stats_[bclass_index]->beta_;
        }
      }  // end looping over all baseclasses
      out_mllr->set_bclass2xforms(base2xforms);
    }  // end of estimating one transform per baseclass
    if (auxf_impr != NULL) *auxf_impr = tot_auxf_impr;
    if (t != NULL) *t = tot_t;
  }
  
  }  // namespace kaldi