// gmm/full-gmm-test.cc
  
  // Copyright 2009-2011  Jan Silovsky;  Saarland University;
  //                      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 "gmm/full-gmm.h"
  #include "gmm/diag-gmm.h"
  #include "gmm/model-test-common.h"
  #include "util/stl-utils.h"
  #include "util/kaldi-io.h"
  #include "gmm/full-gmm-normal.h"
  #include "gmm/mle-full-gmm.h"
  
  using namespace kaldi;
  
  void RandPosdefSpMatrix(size_t dim, SpMatrix<BaseFloat> *matrix,
                            TpMatrix<BaseFloat> *matrix_sqrt = NULL,
                            BaseFloat *logdet = NULL) {
    // generate random (non-singular) matrix
    Matrix<BaseFloat> tmp(dim, dim);
    while (1) {
      tmp.SetRandn();
      if (tmp.Cond() < 100) break;
      std::cout << "Condition number of random matrix large "
        << static_cast<float>(tmp.Cond()) << ", trying again (this is normal)"
        << '
  ';
    }
    // tmp * tmp^T will give positive definite matrix
    matrix->AddMat2(1.0, tmp, kNoTrans, 0.0);
  
    if (matrix_sqrt != NULL) matrix_sqrt->Cholesky(*matrix);
    if (logdet != NULL) *logdet = matrix->LogPosDefDet();
    if ((matrix_sqrt == NULL) && (logdet == NULL)) {
      TpMatrix<BaseFloat> sqrt(dim);
      sqrt.Cholesky(*matrix);
    }
  }
  
  void init_rand_diag_gmm(DiagGmm *gmm) {
    size_t num_comp = gmm->NumGauss(), dim = gmm->Dim();
    Vector<BaseFloat> weights(num_comp);
    Matrix<BaseFloat> means(num_comp, dim), vars(num_comp, dim);
  
    BaseFloat tot_weight = 0.0;
    for (size_t m = 0; m < num_comp; m++) {
      weights(m) = kaldi::RandUniform();
      for (size_t d= 0; d < dim; d++) {
        means(m, d) = kaldi::RandGauss();
        vars(m, d) = Exp(kaldi::RandGauss()) + 1e-5;
      }
      tot_weight += weights(m);
    }
  
    // normalize weights
    for (size_t m = 0; m < num_comp; m++) {
      weights(m) /= tot_weight;
    }
  
    vars.InvertElements();
    gmm->SetWeights(weights);
    gmm->SetInvVarsAndMeans(vars, means);
    gmm->Perturb(0.5 * RandUniform());
    gmm->ComputeGconsts();  // this is unnecassary; computed in Perturb
  }
  
  void UnitTestFullGmmEst() {
    FullGmm fgmm;
    int32 dim = 10 + Rand() % 10, num_comp = 1 + Rand() % 10;
    unittest::InitRandFullGmm(dim, num_comp, &fgmm);
    int32 num_frames = 5000;
    Matrix<BaseFloat> feats(num_frames, dim);
    FullGmmNormal fgmm_normal(fgmm);
    fgmm_normal.Rand(&feats);
  
    AccumFullGmm acc(fgmm, kGmmAll);
    for (int32 t = 0; t < num_frames; t++)
      acc.AccumulateFromFull(fgmm, feats.Row(t), 1.0);
    BaseFloat objf_change, count;
  
    MleFullGmmOptions opts;
  
    MleFullGmmUpdate(opts, acc, kGmmAll, &fgmm, &objf_change, &count);
    BaseFloat change = objf_change / count,
        num_params = (num_comp * (dim + 1 + (dim*(dim+1)/2))),
        predicted_change = 0.5 * num_params / num_frames; // Was there
    KALDI_LOG << "Objf change per frame was " << change << " vs. predicted "
              << predicted_change;
    KALDI_ASSERT(change < 2.0 * predicted_change && change > 0.0);
  }
  
  
  void
  UnitTestFullGmm() {
    // random dimension of the gmm
    size_t dim = 1 + kaldi::RandInt(0, 9);
    // random number of mixtures
    size_t nMix = 1 + kaldi::RandInt(0, 9);
  
    std::cout << "Testing NumGauss: " << nMix << ", " << "Dim: " << dim
      << '
  ';
  
    // generate random feature vector and
    // random mean vectors and covariance matrices
    Vector<BaseFloat> feat(dim);
    Vector<BaseFloat> weights(nMix);
    Vector<BaseFloat> loglikes(nMix);
    Matrix<BaseFloat> means(nMix, dim);
    std::vector<SpMatrix<BaseFloat> > invcovars(nMix);
    for (size_t mix = 0; mix < nMix; mix++) {
      invcovars[mix].Resize(dim);
    }
    Vector<BaseFloat> covars_logdet(nMix);
  
    for (size_t d = 0; d < dim; d++) {
      feat(d) = kaldi::RandGauss();
    }
  
    float tot_weight = 0.0;
    for (size_t m = 0; m < nMix; m++) {
      weights(m) = kaldi::RandUniform();
      for (size_t d = 0; d < dim; d++) {
        means(m, d) = kaldi::RandGauss();
      }
      SpMatrix<BaseFloat> covar(dim);
      RandPosdefSpMatrix(dim, &covar, NULL, &covars_logdet(m));
      invcovars[m].CopyFromSp(covar);
      invcovars[m].InvertDouble();
      tot_weight += weights(m);
    }
  
    // normalize weights and compute loglike for feature vector
    for (size_t m = 0; m < nMix; m++) {
      weights(m) /= tot_weight;
    }
  
    // compute loglike for feature vector
    float loglike = 0.0;
    for (size_t m = 0; m < nMix; m++) {
      loglikes(m) += -0.5 * (M_LOG_2PI * dim
        + covars_logdet(m)
        + VecSpVec(means.Row(m), invcovars[m], means.Row(m))
        + VecSpVec(feat, invcovars[m], feat))
        + VecSpVec(means.Row(m), invcovars[m], feat);
      loglikes(m) += Log(weights(m));
    }
  
    loglike = loglikes.LogSumExp();
  
  
    // new GMM
    FullGmm *gmm = new FullGmm();
    gmm->Resize(nMix, dim);
    gmm->SetWeights(weights);
    gmm->SetInvCovarsAndMeans(invcovars, means);
    gmm->ComputeGconsts();
  
    Vector<BaseFloat> posterior1(nMix);
    float loglike1 = gmm->ComponentPosteriors(feat, &posterior1);
  
    // std::cout << "LogLike: " << loglike << '
  ';
    // std::cout << "LogLike1: " << loglike1 << '
  ';
  
    AssertEqual(loglike, loglike1, 0.01);
  
    KALDI_ASSERT(fabs(1.0 - posterior1.Sum()) < 0.001);
  
    {  // Test various accessors / mutators
      Vector<BaseFloat> weights_bak(nMix);
      Matrix<BaseFloat> means_bak(nMix, dim);
      std::vector<SpMatrix<BaseFloat> > invcovars_bak(nMix);
      for (size_t i = 0; i < nMix; i++) {
        invcovars_bak[i].Resize(dim);
      }
  
      weights_bak.CopyFromVec(gmm->weights());
      gmm->GetMeans(&means_bak);
      gmm->GetCovars(&invcovars_bak);
      for (size_t i = 0; i < nMix; i++) {
        invcovars_bak[i].InvertDouble();
      }
  
      // set all params one-by-one to new model
      FullGmm gmm2;
      gmm2.Resize(gmm->NumGauss(), gmm->Dim());
      gmm2.SetWeights(weights_bak);
      gmm2.SetMeans(means_bak);
      gmm2.SetInvCovars(invcovars_bak);
      gmm2.ComputeGconsts();
      BaseFloat loglike_gmm2 = gmm2.LogLikelihood(feat);
      AssertEqual(loglike1, loglike_gmm2);
      {
        Vector<BaseFloat> loglikes;
        gmm2.LogLikelihoods(feat, &loglikes);
        AssertEqual(loglikes.LogSumExp(), loglike_gmm2);
      }
      {
        std::vector<int32> indices;
        for (int32 i = 0; i < gmm2.NumGauss(); i++)
          indices.push_back(i);
        Vector<BaseFloat> loglikes;
        gmm2.LogLikelihoodsPreselect(feat, indices, &loglikes);
        AssertEqual(loglikes.LogSumExp(), loglike_gmm2);
      }
  
  
      // single component mean accessor + mutator
      FullGmm gmm3;
      gmm3.Resize(gmm->NumGauss(), gmm->Dim());
      gmm3.SetWeights(weights_bak);
      means_bak.SetZero();
      for (size_t i = 0; i < nMix; i++) {
        SubVector<BaseFloat> tmp = means_bak.Row(i);
        gmm->GetComponentMean(i, &tmp);
      }
      gmm3.SetMeans(means_bak);
      gmm3.SetInvCovars(invcovars_bak);
      gmm3.ComputeGconsts();
      float loglike_gmm3 = gmm3.LogLikelihood(feat);
      AssertEqual(loglike1, loglike_gmm3, 0.01);
  
      // set all params one-by-one to new model
      FullGmm gmm4;
      gmm4.Resize(gmm->NumGauss(), gmm->Dim());
      gmm4.SetWeights(weights_bak);
      gmm->GetCovarsAndMeans(&invcovars_bak, &means_bak);
      for (size_t i = 0; i < nMix; i++) {
        invcovars_bak[i].InvertDouble();
      }
      gmm4.SetInvCovarsAndMeans(invcovars_bak, means_bak);
      gmm4.ComputeGconsts();
      BaseFloat loglike_gmm4 = gmm4.LogLikelihood(feat);
      AssertEqual(loglike1, loglike_gmm4, 0.001);
  
    }  // Test various accessors / mutators end
  
     // First, non-binary write
    gmm->Write(Output("tmpf", false).Stream(), false);
  
    {  // I/O tests
      bool binary_in;
      FullGmm *gmm2 = new FullGmm();
      Input ki("tmpf", &binary_in);
      gmm2->Read(ki.Stream(), binary_in);
  
      float loglike3 = gmm2->ComponentPosteriors(feat, &posterior1);
      AssertEqual(loglike, loglike3, 0.01);
  
      // binary write
      gmm2->Write(Output("tmpfb", true).Stream(), true);
      delete gmm2;
  
      // binary read
      FullGmm *gmm3;
      gmm3 = new FullGmm();
  
      Input ki2("tmpfb", &binary_in);
      gmm3->Read(ki2.Stream(), binary_in);
  
      AssertEqual(loglike, loglike3, 0.01);
  
      delete gmm3;
    }
  
    {  // CopyFromFullGmm
      FullGmm gmm4;
      gmm4.CopyFromFullGmm(*gmm);
      float loglike5 = gmm4.ComponentPosteriors(feat, &posterior1);
      AssertEqual(loglike, loglike5, 0.01);
    }
  
    {  // test copy from DiagGmm and back to DiagGmm
      DiagGmm gmm_diag;
      gmm_diag.Resize(nMix, dim);
      init_rand_diag_gmm(&gmm_diag);
      float loglike_diag = gmm_diag.LogLikelihood(feat);
  
      FullGmm gmm_full;
      gmm_full.CopyFromDiagGmm(gmm_diag);
      float loglike_full = gmm_full.LogLikelihood(feat);
  
      DiagGmm gmm_diag2;
      gmm_diag2.CopyFromFullGmm(gmm_full);
      float loglike_diag2 = gmm_diag2.LogLikelihood(feat);
  
      AssertEqual(loglike_diag, loglike_full, 0.01);
      AssertEqual(loglike_diag, loglike_diag2, 0.01);
    }
  
    {  // split and merge test for 1 component GMM (doesn't test the merge crit.)
      FullGmm gmm1;
      Vector<BaseFloat> weights1(1);
      Matrix<BaseFloat> means1(1, dim);
      std::vector<SpMatrix<BaseFloat> > invcovars1(1);
      weights1(0) = 1.0;
      means1.CopyFromMat(means.Range(0, 1, 0, dim));
      invcovars1[0].Resize(dim);
      invcovars1[0].CopyFromSp(invcovars[0]);
      gmm1.Resize(1, dim);
      gmm1.SetWeights(weights1);
      gmm1.SetInvCovarsAndMeans(invcovars1, means1);
      gmm1.ComputeGconsts();
      FullGmm gmm2;
      gmm2.CopyFromFullGmm(gmm1);
      gmm2.Split(2, 0.001);
      gmm2.Merge(1);
      float loglike1 = gmm1.LogLikelihood(feat);
      float loglike2 = gmm2.LogLikelihood(feat);
      AssertEqual(loglike1, loglike2, 0.01);
    }
  
    delete gmm;
  
    unlink("tmpf");
    unlink("tmpfb");
  }
  
  int
  main() {
    // repeat the test ten times
    for (int i = 0; i < 2; i++) {
      UnitTestFullGmm();
      UnitTestFullGmmEst();
    }
    std::cout << "Test OK.
  ";
  }