// transform/regtree-mllr-diag-gmm-test.cc

// Copyright 2009-2011   Saarland University
// Author:  Arnab Ghoshal

// 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-mllr-diag-gmm.h"

using kaldi::int32;
using kaldi::BaseFloat;
using kaldi::RegtreeMllrDiagGmmAccs;
namespace ut = kaldi::unittest;

void TestMllrAccsIO(const kaldi::AmDiagGmm &am_gmm,
                    const kaldi::RegressionTree &regtree,
                    const RegtreeMllrDiagGmmAccs &accs,
                    const kaldi::Matrix<BaseFloat> adapt_data) {
  // First, non-binary write
  accs.Write(kaldi::Output("tmpf", false).Stream(), false);

  kaldi::RegtreeMllrDiagGmm mllr;
  kaldi::RegtreeMllrOptions opts;
  opts.min_count = 100;
  opts.use_regtree = false;
  accs.Update(regtree, opts, &mllr, NULL, NULL);
  kaldi::AmDiagGmm am1;
  am1.CopyFromAmDiagGmm(am_gmm);
  mllr.TransformModel(regtree, &am1);

  BaseFloat loglike = 0;
  int32 npoints = adapt_data.NumRows();
  for (int32 j = 0; j < npoints; j++) {
    loglike += am1.LogLikelihood(0, adapt_data.Row(j));
  }
  KALDI_LOG << "Per-frame loglike after adaptation = " << (loglike/npoints)
            << " over " << npoints << " frames.";

  size_t num_comp2 = 1 + kaldi::RandInt(0, 9);  // random number of mixtures
  int32 dim = am_gmm.Dim();
  kaldi::DiagGmm gmm2;
  ut::InitRandDiagGmm(dim, num_comp2, &gmm2);
  kaldi::Vector<BaseFloat> data(dim);
  gmm2.Generate(&data);
  BaseFloat loglike1 = am1.LogLikelihood(0, data);
//  KALDI_LOG << "LL0 = " << loglike0 << "; LL1 = " << loglike1;

  KALDI_LOG << "Test ASCII IO.";
  bool binary_in;
  kaldi::RegtreeMllrDiagGmm mllr1;
  RegtreeMllrDiagGmmAccs *accs1 = new RegtreeMllrDiagGmmAccs();
  // Non-binary read
  kaldi::Input ki1("tmpf", &binary_in);
  accs1->Read(ki1.Stream(), binary_in, false);
  accs1->Update(regtree, opts, &mllr1, NULL, NULL);
  delete accs1;
  kaldi::AmDiagGmm am2;
  am2.CopyFromAmDiagGmm(am_gmm);
  mllr.TransformModel(regtree, &am2);
  BaseFloat loglike2 = am2.LogLikelihood(0, data);
//  KALDI_LOG << "LL1 = " << loglike1 << "; LL2 = " << loglike2;
  kaldi::AssertEqual(loglike1, loglike2, 1e-6);

  kaldi::RegtreeMllrDiagGmm mllr2;
  // Next, binary write
  KALDI_LOG << "Test Binary IO.";
  accs.Write(kaldi::Output("tmpfb", true).Stream(), true);
  RegtreeMllrDiagGmmAccs *accs2 = new RegtreeMllrDiagGmmAccs();
  // Binary read
  kaldi::Input ki2("tmpfb", &binary_in);
  accs2->Read(ki2.Stream(), binary_in, false);
  accs2->Update(regtree, opts, &mllr2, NULL, NULL);
  delete accs2;
  kaldi::AmDiagGmm am3;
  am3.CopyFromAmDiagGmm(am_gmm);
  mllr.TransformModel(regtree, &am3);
  BaseFloat loglike3 = am3.LogLikelihood(0, data);
//  KALDI_LOG << "LL1 = " << loglike1 << "; LL3 = " << loglike3;
  kaldi::AssertEqual(loglike1, loglike3, 1e-6);
  
  unlink("tmpf");
  unlink("tmpfb");
}

void TestXformMean(const kaldi::AmDiagGmm &am_gmm,
                   const kaldi::RegressionTree &regtree,
                   const RegtreeMllrDiagGmmAccs &accs,
                   const kaldi::Matrix<BaseFloat> adapt_data) {
  kaldi::RegtreeMllrDiagGmm mllr;
  kaldi::RegtreeMllrOptions opts;
  opts.min_count = 100;
  opts.use_regtree = false;
  accs.Update(regtree, opts, &mllr, NULL, NULL);

  kaldi::AmDiagGmm am1;
  am1.CopyFromAmDiagGmm(am_gmm);
  mllr.TransformModel(regtree, &am1);

  kaldi::DiagGmm tmp_pdf;
  tmp_pdf.CopyFromDiagGmm(am_gmm.GetPdf(0));
  kaldi::Matrix<BaseFloat> tmp_means(am_gmm.GetPdf(0).NumGauss(), am_gmm.Dim());
  mllr.GetTransformedMeans(regtree, am_gmm, 0, &tmp_means);
  tmp_pdf.SetInvVarsAndMeans(tmp_pdf.inv_vars(), tmp_means);
  tmp_pdf.ComputeGconsts();

  BaseFloat loglike0 = 0, loglike = 0;
  int32 npoints = adapt_data.NumRows();
  for (int32 j = 0; j < npoints; j++) {
    loglike0 += am1.LogLikelihood(0, adapt_data.Row(j));
    loglike += tmp_pdf.LogLikelihood(adapt_data.Row(j));
  }
  KALDI_LOG << "Per-frame loglike after adaptation = " << (loglike0/npoints)
            << " over " << npoints << " frames.";
//  KALDI_LOG << "LL0 = " << loglike0 << "; LL = " << loglike;
  kaldi::AssertEqual(loglike0, loglike, 1e-6);

  kaldi::Matrix<BaseFloat> tmp_means2(am_gmm.GetPdf(0).NumGauss(), am_gmm.Dim());
  mllr.GetTransformedMeans(regtree, am_gmm, 0, &tmp_means2);
  tmp_pdf.SetInvVarsAndMeans(tmp_pdf.inv_vars(), tmp_means2);
  tmp_pdf.ComputeGconsts();

  BaseFloat loglike1 = 0;
  for (int32 j = 0; j < npoints; j++) {
    loglike1 += tmp_pdf.LogLikelihood(adapt_data.Row(j));
  }
//  KALDI_LOG << "LL = " << loglike << "; LL1 = " << loglike1;
  kaldi::AssertEqual(loglike, loglike1, 1e-6);
}


void UnitTestRegtreeMllrDiagGmm() {
  size_t dim = 1 + kaldi::RandInt(1, 9);  // random dimension of the gmm
  size_t num_comp = 1 + kaldi::RandInt(0, 5);  // random number of mixtures
  kaldi::DiagGmm gmm;
  ut::InitRandDiagGmm(dim, num_comp, &gmm);
  kaldi::AmDiagGmm am_gmm;
  am_gmm.Init(gmm, 1);

  size_t num_comp2 = 1 + kaldi::RandInt(0, 5);  // random number of mixtures
  kaldi::DiagGmm gmm2;
  ut::InitRandDiagGmm(dim, num_comp2, &gmm2);
  int32 npoints = dim*(dim+1)*10 + 500;
  kaldi::Matrix<BaseFloat> adapt_data(npoints, dim);
  for (int32 j = 0; j < npoints; j++) {
    kaldi::SubVector<BaseFloat> row(adapt_data, j);
    gmm2.Generate(&row);
  }

  kaldi::RegressionTree regtree;
  std::vector<int32> sil_indices;
  kaldi::Vector<BaseFloat> state_occs(1);
  state_occs(0) = npoints;
  regtree.BuildTree(state_occs, sil_indices, am_gmm, 1);
  int32 num_bclass = regtree.NumBaseclasses();

  kaldi::RegtreeMllrDiagGmmAccs accs;
  BaseFloat loglike = 0;
  accs.Init(num_bclass, dim);
  for (int32 j = 0; j < npoints; j++) {
    loglike += accs.AccumulateForGmm(regtree, am_gmm, adapt_data.Row(j),
                                     0, 1.0);
  }
  KALDI_LOG << "Per-frame loglike during accumulations = " << (loglike/npoints)
            << " over " << npoints << " frames.";

  TestMllrAccsIO(am_gmm, regtree, accs, adapt_data);
  TestXformMean(am_gmm, regtree, accs, adapt_data);
}

int main() {
  kaldi::g_kaldi_verbose_level = 5;
  for (int i = 0; i <= 10; i++)
    UnitTestRegtreeMllrDiagGmm();
  std::cout << "Test OK.\n";
}