// transform/regtree-fmllr-diag-gmm-test.cc
  
  // Copyright 2009-2011  Georg Stemmer;  Saarland University
  
  // 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 "util/common-utils.h"
  #include "gmm/diag-gmm.h"
  #include "gmm/mle-diag-gmm.h"
  #include "gmm/mle-am-diag-gmm.h"
  #include "gmm/model-test-common.h"
  #include "transform/regtree-fmllr-diag-gmm.h"
  
  namespace kaldi {
  
  static void
  RandFullCova(Matrix<BaseFloat> *matrix) {
    size_t dim = matrix->NumCols();
    KALDI_ASSERT(matrix->NumCols() == matrix->NumRows());
  
    size_t iter = 0;
    size_t max_iter = 10000;
    // generate random (non-singular) matrix
    // until condition
    Matrix<BaseFloat> tmp(dim, dim);
    SpMatrix<BaseFloat> tmp2(dim);
    while (iter < max_iter) {
      tmp.SetRandn();
      if (tmp.Cond() < 100) break;
      iter++;
    }
    if (iter >= max_iter) {
      KALDI_ERR << "Internal error: found no random covariance matrix.";
    }
    // tmp * tmp^T will give positive definite matrix
    tmp2.AddMat2(1.0, tmp, kNoTrans, 0.0);
    matrix->CopyFromSp(tmp2);
  }
  
  
  /// Generate features for a certain covariance type
  /// covariance_type == 0: full covariance
  /// covariance_type == 1: diagonal covariance
  
  enum cova_type {
    full,
    diag
  };
  
  static void
  generate_features(cova_type covariance_type,
                    size_t n_gaussians,
                    size_t dim,
                    Matrix<BaseFloat> &trans_mat,
                    size_t frames_per_gaussian,
                    std::vector<Vector<BaseFloat>*> & train_feats,
                    std::vector<Vector<BaseFloat>*> & adapt_feats
                    ) {
    // compute inverse of the transformation matrix
    Matrix<BaseFloat> inv_trans_mat(dim, dim);
    inv_trans_mat.CopyFromMat(trans_mat, kNoTrans);
    inv_trans_mat.Invert();
    // the untransformed means are random
    Matrix<BaseFloat> untransformed_means(dim, n_gaussians);
    untransformed_means.SetRandn();
    untransformed_means.Scale(10);
  
    // the actual means result from
    // transformation with inv_trans_mat
    Matrix<BaseFloat> actual_means(dim, n_gaussians);
  
    // actual_means = inv_trans_mat * untransformed_means
    actual_means.AddMatMat(1.0, inv_trans_mat, kNoTrans,
                           untransformed_means, kNoTrans, 0.0);
  
    size_t train_counter = 0;
  
    // temporary variables
    Vector<BaseFloat> randomvec(dim);
    Matrix<BaseFloat> Sj(dim, dim);
  
    // loop over all gaussians
    for (size_t j = 0; j < n_gaussians; j++) {
      if (covariance_type == diag) {
        // random diagonal covariance for gaussian j
        Sj.SetZero();
        for (size_t d = 0; d < dim; d++) {
          Sj(d, d) = 2*Exp(RandGauss());
        }
      }
      if (covariance_type == full) {
        // random full covariance for gaussian j
        RandFullCova(&Sj);
      }
      // compute inv_trans_mat * Sj
      Matrix<BaseFloat> tmp_matrix(dim, dim);
      tmp_matrix.AddMatMat(1.0, inv_trans_mat, kNoTrans, Sj, kNoTrans, 0.0);
  
      // compute features
      for (size_t i = 0; i < frames_per_gaussian; i++) {
        train_feats[train_counter] = new Vector<BaseFloat>(dim);
        adapt_feats[train_counter] = new Vector<BaseFloat>(dim);
  
        // initalize feature vector with mean of class j
        train_feats[train_counter]->CopyColFromMat(untransformed_means, j);
        adapt_feats[train_counter]->CopyColFromMat(actual_means, j);
  
        // determine random vector and
        // multiply the random vector with SJ
        // and add it to train_feats:
        // train_feats = train_feats + SJ * random
        // for adapt_feats we include the invtrans_mat:
        // adapt_feats = adapt_feats + invtrans_mat * SJ * random
        for (size_t d = 0; d < dim; d++) {
          randomvec(d) = RandGauss();
        }
        train_feats[train_counter]->AddMatVec(1.0, Sj, kNoTrans,
                                              randomvec, 1.0);
        adapt_feats[train_counter]->AddMatVec(1.0, tmp_matrix, kNoTrans,
                                              randomvec, 1.0);
        train_counter++;
      }
    }
    return;
  }
  
  void UnitTestRegtreeFmllrDiagGmm(cova_type feature_type, size_t max_bclass) {
    // dimension of the feature space
    size_t dim = 5 + Rand() % 3;
  
    // number of components in the data
    size_t n_gaussians = 8;
  
    // number of data points to generate for every gaussian
    size_t frames_per_gaussian = 100;
  
    // generate random transformation matrix trans_mat
    Matrix<BaseFloat> trans_mat(dim, dim);
    int i = 0;
    while (i < 10000) {
      trans_mat.SetRandn();
      if (trans_mat.Cond() < 100) break;
      i++;
    }
    std::cout << "Condition of original Trans_Mat: " << trans_mat.Cond() << '
  ';
  
    // generate many feature vectors for each of the mixture components
    std::vector<Vector<BaseFloat>*>
        train_feats(n_gaussians * frames_per_gaussian);
    std::vector<Vector<BaseFloat>*>
        adapt_feats(n_gaussians * frames_per_gaussian);
  
    generate_features(feature_type,
                      n_gaussians,
                      dim,
                      trans_mat,
                      frames_per_gaussian,
                      train_feats,
                      adapt_feats);
  
    // initial values for a GMM
    Vector<BaseFloat> weights(1);
    Matrix<BaseFloat> means(1, dim), vars(1, dim), invvars(1, dim);
    for (size_t d= 0; d < dim; d++) {
      means(0, d) = 0.0F;
      vars(0, d) = 1.0F;
    }
    weights(0) = 1.0F;
    invvars.CopyFromMat(vars);
    invvars.InvertElements();
  
    // new HMM with 1 state
    DiagGmm *gmm = new DiagGmm();
    gmm->Resize(1, dim);
    gmm->SetWeights(weights);
    gmm->SetInvVarsAndMeans(invvars, means);
    gmm->ComputeGconsts();
    GmmFlagsType flags = kGmmAll;
    MleDiagGmmOptions opts;
  
    AmDiagGmm *am = new AmDiagGmm();
    am->AddPdf(*gmm);
    AccumAmDiagGmm *est_am = new AccumAmDiagGmm();
  
    // train HMM
    size_t iteration = 0;
    size_t maxiterations = 10;
    int32 maxcomponents = n_gaussians;
    BaseFloat loglike = 0;
    while (iteration < maxiterations) {
      est_am->Init(*am, flags);
  
      loglike = 0;
      for (size_t j = 0; j < train_feats.size(); j++) {
        loglike += est_am->AccumulateForGmm(*am, *train_feats[j], 0, 1.0);
      }
      MleAmDiagGmmUpdate(opts, *est_am, flags, am, NULL, NULL);
  
      std::cout << "Loglikelihood before iteration " << iteration << " : "
                << std::scientific << loglike << " number of components: "
                << am->NumGaussInPdf(0) << '
  ';
  
      if ((iteration % 3 == 1) &&
          (am->NumGaussInPdf(0) * 2 <= maxcomponents)) {
        size_t n = am->NumGaussInPdf(0)*2;
        am->SplitPdf(0, n, 0.001);
      }
      iteration++;
    }
  
    // adapt HMM to the transformed feature vectors
    iteration = 0;
    RegtreeFmllrDiagGmmAccs * fmllr_accs = new RegtreeFmllrDiagGmmAccs();
    RegressionTree regtree;
  
    RegtreeFmllrOptions xform_opts;
    xform_opts.min_count = 100 * (1 + Rand() % 10);
    xform_opts.use_regtree = (RandUniform() < 0.5)? false : true;
  
    size_t num_pdfs = 1;
    Vector<BaseFloat> occs(num_pdfs);
    for (int32 i = 0; i < static_cast<int32>(num_pdfs); i++) {
      occs(i) = 1.0/static_cast<BaseFloat>(num_pdfs);
    }
    std::vector<int32> silphones;
    regtree.BuildTree(occs, silphones, *am, max_bclass);
    maxiterations = 10;
    std::vector<Vector<BaseFloat>*> logdet(adapt_feats.size());
    for (size_t j = 0; j < adapt_feats.size(); j++) {
      logdet[j] = new Vector<BaseFloat>(1);
      logdet[j]->operator()(0) = 0.0;
    }
    while (iteration < maxiterations) {
      fmllr_accs->Init(regtree.NumBaseclasses(), dim);
      fmllr_accs->SetZero();
      RegtreeFmllrDiagGmm *new_fmllr = new RegtreeFmllrDiagGmm();
      loglike = 0;
      for (size_t j = 0; j < adapt_feats.size(); j++) {
        loglike += fmllr_accs->AccumulateForGmm(regtree, *am, *adapt_feats[j], 0, 1.0);
        loglike += logdet[j]->operator()(0);
      }
      std::cout << "FMLLR: Loglikelihood before iteration " << iteration << " : "
                << std::scientific << loglike << '
  ';
  
      fmllr_accs->Update(regtree, xform_opts, new_fmllr, NULL, NULL);
      std::cout << "Got " << new_fmllr->NumBaseClasses() << " baseclasses
  ";
      bool binary = (RandUniform() < 0.5)? true : false;
      std::cout << "Writing the transform to disk.
  ";
      new_fmllr->Write(Output("tmpf", binary).Stream(), binary);
      RegtreeFmllrDiagGmm *fmllr_read = new RegtreeFmllrDiagGmm();
      bool binary_in;
      Input ki("tmpf", &binary_in);
      std::cout << "Reading the transform from disk.
  ";
      fmllr_read->Read(ki.Stream(), binary_in);
      fmllr_read->Validate();
  
      // transform features
      std::vector<Vector<BaseFloat> > trans_feats(1);
      Vector<BaseFloat> trans_logdet;
  //    new_fmllr->ComputeLogDets();
      trans_logdet.Resize(fmllr_read->NumRegClasses());
      fmllr_read->GetLogDets(&trans_logdet);
      for (size_t j = 0; j < adapt_feats.size(); j++) {
        fmllr_read->TransformFeature(*adapt_feats[j], &trans_feats);
        logdet[j]->operator()(0) += trans_logdet(0);
        adapt_feats[j]->CopyFromVec(trans_feats[0]);
      }
      iteration++;
      delete new_fmllr;
      delete fmllr_read;
      
      unlink("tmpf");
    }
  
  //  // transform features with empty transform
  //  std::vector<Vector<BaseFloat> > trans_feats(1);
  //  RegtreeFmllrDiagGmm *empty_fmllr = new RegtreeFmllrDiagGmm();
  //  empty_fmllr->Init(0, 0);
  //  for (size_t j = 0; j < adapt_feats.size(); j++) {
  //    empty_fmllr->TransformFeature(*adapt_feats[j], &trans_feats);
  //  }
  //  delete empty_fmllr;
  
    // clean up
    delete fmllr_accs;
    delete est_am;
    delete am;
    delete gmm;
    DeletePointers(&logdet);
    DeletePointers(&train_feats);
    DeletePointers(&adapt_feats);
  }
  }  // namespace kaldi ends here
  
  int main() {
    for (int i = 0; i <= 8; i+=2) {  // test is too slow so can't do too many
      std::cout << "--------------------------------------" << '
  ';
      std::cout << "Test number " << i << '
  ';
      std::cout << "--
  features = full
  ";
      kaldi::UnitTestRegtreeFmllrDiagGmm(kaldi::full, (i%10+1));
      std::cout << "--
  features = diag
  ";
      kaldi::UnitTestRegtreeFmllrDiagGmm(kaldi::diag, (i%10+1));
      std::cout << "--------------------------------------" << '
  ';
    }
    std::cout << "Test OK.
  ";
  }