// gmm/diag-gmm.cc

// Copyright 2009-2011  Microsoft Corporation;
//                      Saarland University (Author: Arnab Ghoshal);
//                      Georg Stemmer;  Jan Silovsky
//           2012       Arnab Ghoshal
//           2013-2014  Johns Hopkins University (author: Daniel Povey)

// 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 <algorithm>
#include <functional>
#include <limits>
#include <string>
#include <vector>

#include "gmm/diag-gmm.h"
#include "gmm/diag-gmm-normal.h"
#include "gmm/full-gmm.h"
#include "gmm/full-gmm-normal.h"
#include "tree/clusterable-classes.h"

namespace kaldi {

// Constructor that allows us to merge GMMs.
DiagGmm::DiagGmm(const std::vector<std::pair<BaseFloat, const DiagGmm*> > &gmms)
  : valid_gconsts_(false) {
  if (gmms.empty()) {
    return;  // GMM will be empty.
  } else {
    int32 num_gauss = 0, dim = gmms[0].second->Dim();
    for (size_t i = 0; i < gmms.size(); i++)
      num_gauss += gmms[i].second->NumGauss();
    Resize(num_gauss, dim);
    int32 cur_gauss = 0;
    for (size_t i = 0; i < gmms.size(); i++) {
      BaseFloat weight = gmms[i].first;
      KALDI_ASSERT(weight > 0.0);
      const DiagGmm &gmm = *(gmms[i].second);
      for (int32 g = 0; g < gmm.NumGauss(); g++, cur_gauss++) {
        means_invvars_.Row(cur_gauss).CopyFromVec(gmm.means_invvars().Row(g));
        inv_vars_.Row(cur_gauss).CopyFromVec(gmm.inv_vars().Row(g));
        weights_(cur_gauss) = weight * gmm.weights()(g);
      }
    }
    KALDI_ASSERT(cur_gauss == NumGauss());
    ComputeGconsts();
  }
}



void DiagGmm::Resize(int32 nmix, int32 dim) {
  KALDI_ASSERT(nmix > 0 && dim > 0);
  if (gconsts_.Dim() != nmix) gconsts_.Resize(nmix);
  if (weights_.Dim() != nmix) weights_.Resize(nmix);
  if (inv_vars_.NumRows() != nmix ||
      inv_vars_.NumCols() != dim) {
    inv_vars_.Resize(nmix, dim);
    inv_vars_.Set(1.0);
    // must be initialized to unit for case of calling SetMeans while having
    // covars/invcovars that are not set yet (i.e. zero)
  }
  if (means_invvars_.NumRows() != nmix ||
      means_invvars_.NumCols() != dim)
    means_invvars_.Resize(nmix, dim);
  valid_gconsts_ = false;
}

void DiagGmm::CopyFromDiagGmm(const DiagGmm &diaggmm) {
  Resize(diaggmm.weights_.Dim(), diaggmm.means_invvars_.NumCols());
  gconsts_.CopyFromVec(diaggmm.gconsts_);
  weights_.CopyFromVec(diaggmm.weights_);
  inv_vars_.CopyFromMat(diaggmm.inv_vars_);
  means_invvars_.CopyFromMat(diaggmm.means_invvars_);
  valid_gconsts_ = diaggmm.valid_gconsts_;
}

void DiagGmm::CopyFromFullGmm(const FullGmm &fullgmm) {
  int32 num_comp = fullgmm.NumGauss(), dim = fullgmm.Dim();
  Resize(num_comp, dim);
  gconsts_.CopyFromVec(fullgmm.gconsts());
  weights_.CopyFromVec(fullgmm.weights());
  Matrix<BaseFloat> means(num_comp, dim);
  fullgmm.GetMeans(&means);
  int32 ncomp = NumGauss();
  for (int32 mix = 0; mix < ncomp; mix++) {
    SpMatrix<double> covar(dim);
    covar.CopyFromSp(fullgmm.inv_covars()[mix]);
    covar.Invert();
    Vector<double> diag(dim);
    diag.CopyDiagFromPacked(covar);
    diag.InvertElements();
    inv_vars_.Row(mix).CopyFromVec(diag);
  }
  means_invvars_.CopyFromMat(means);
  means_invvars_.MulElements(inv_vars_);
  ComputeGconsts();
}

int32 DiagGmm::ComputeGconsts() {
  int32 num_mix = NumGauss();
  int32 dim = Dim();
  BaseFloat offset = -0.5 * M_LOG_2PI * dim;  // constant term in gconst.
  int32 num_bad = 0;

  // Resize if Gaussians have been removed during Update()
  if (num_mix != static_cast<int32>(gconsts_.Dim()))
    gconsts_.Resize(num_mix);

  for (int32 mix = 0; mix < num_mix; mix++) {
    KALDI_ASSERT(weights_(mix) >= 0);  // Cannot have negative weights.
    BaseFloat gc = Log(weights_(mix)) + offset;  // May be -inf if weights == 0
    for (int32 d = 0; d < dim; d++) {
      gc += 0.5 * Log(inv_vars_(mix, d)) - 0.5 * means_invvars_(mix, d)
        * means_invvars_(mix, d) / inv_vars_(mix, d);
    }
    // Change sign for logdet because var is inverted. Also, note that
    // mean_invvars(mix, d)*mean_invvars(mix, d)/inv_vars(mix, d) is the
    // mean-squared times inverse variance, since mean_invvars(mix, d) contains
    // the mean times inverse variance.
    // So gc is the likelihood at zero feature value.

    if (KALDI_ISNAN(gc)) {  // negative infinity is OK but NaN is not acceptable
      KALDI_ERR << "At component "  << mix
                << ", not a number in gconst computation";
    }
    if (KALDI_ISINF(gc)) {
      num_bad++;
      // If positive infinity, make it negative infinity.
      // Want to make sure the answer becomes -inf in the end, not NaN.
      if (gc > 0) gc = -gc;
    }
    gconsts_(mix) = gc;
  }

  valid_gconsts_ = true;
  return num_bad;
}

void DiagGmm::Split(int32 target_components, float perturb_factor,
                    std::vector<int32> *history) {
  if (target_components < NumGauss() || NumGauss() == 0) {
    KALDI_ERR << "Cannot split from "  << NumGauss() << " to "
              << target_components  << " components";
  }
  if (target_components == NumGauss()) {
    KALDI_WARN << "Already have the target # of Gaussians. Doing nothing.";
    return;
  }

  int32 current_components = NumGauss(), dim = Dim();
  DiagGmm *tmp = new DiagGmm;
  tmp->CopyFromDiagGmm(*this);  // so we have copies of matrices
  // First do the resize:
  weights_.Resize(target_components);
  weights_.Range(0, current_components).CopyFromVec(tmp->weights_);
  means_invvars_.Resize(target_components, dim);
  means_invvars_.Range(0, current_components, 0, dim).CopyFromMat(
      tmp->means_invvars_);
  inv_vars_.Resize(target_components, dim);
  inv_vars_.Range(0, current_components, 0, dim).CopyFromMat(tmp->inv_vars_);
  gconsts_.Resize(target_components);

  delete tmp;

  // future work(arnab): Use a priority queue instead?
  while (current_components < target_components) {
    BaseFloat max_weight = weights_(0);
    int32 max_idx = 0;
    for (int32 i = 1; i < current_components; i++) {
      if (weights_(i) > max_weight) {
        max_weight = weights_(i);
        max_idx = i;
      }
    }

    // remember what component was split
    if (history != NULL)
      history->push_back(max_idx);

    weights_(max_idx) /= 2;
    weights_(current_components) = weights_(max_idx);
    Vector<BaseFloat> rand_vec(dim);
    for (int32 i = 0; i < dim; i++) {
      rand_vec(i) = RandGauss() * std::sqrt(inv_vars_(max_idx, i));
      // note, this looks wrong but is really right because it's the
      // means_invvars we're multiplying and they have the dimension
      // of an inverse standard variance. [dan]
    }
    inv_vars_.Row(current_components).CopyFromVec(inv_vars_.Row(max_idx));
    means_invvars_.Row(current_components).CopyFromVec(means_invvars_.Row(
        max_idx));
    means_invvars_.Row(current_components).AddVec(perturb_factor, rand_vec);
    means_invvars_.Row(max_idx).AddVec(-perturb_factor, rand_vec);
    current_components++;
  }
  ComputeGconsts();
}


void DiagGmm::Perturb(float perturb_factor) {
  int32 num_comps = NumGauss(),
      dim = Dim();
  Matrix<BaseFloat> rand_mat(num_comps, dim);
  for (int32 i = 0; i < num_comps; i++) {
    for (int32 d = 0; d < dim; d++) {
      rand_mat(i, d) = RandGauss() * std::sqrt(inv_vars_(i, d));
      // as in DiagGmm::Split, we perturb the means_invvars using a random
      // fraction of inv_vars_
    }
  }
  means_invvars_.AddMat(perturb_factor, rand_mat, kNoTrans);
  ComputeGconsts();
}


void DiagGmm::MergeKmeans(int32 target_components,
                          ClusterKMeansOptions cfg) {
  if (target_components <= 0 || NumGauss() < target_components) {
    KALDI_ERR << "Invalid argument for target number of Gaussians (="
              << target_components << "), #Gauss = " << NumGauss();
  }
  if (NumGauss() == target_components) {
    KALDI_VLOG(2) << "No components merged, as target (" << target_components
                  << ") = total.";
    return;  // Nothing to do.
  }
  double min_var = 1.0e-10;
  std::vector<Clusterable*> clusterable_vec;
  for (int32 g = 0; g < NumGauss(); g++) {
    if (weights_(g) == 0) {
      KALDI_WARN << "Not using zero-weight Gaussians in clustering.";
      continue;
    }
    Vector<BaseFloat> x_stats(Dim()),
        x2_stats(Dim());
    BaseFloat count = weights_(g);

    SubVector<BaseFloat> inv_var(inv_vars_, g),
        mean_invvar(means_invvars_, g);
    x_stats.AddVecDivVec(1.0, mean_invvar, inv_var, count);  // x_stats is now mean.
    x2_stats.CopyFromVec(inv_var);
    x2_stats.InvertElements();  // x2_stats is now var.
    x2_stats.AddVec2(1.0, x_stats);  // x2_stats is now var + mean^2
    x_stats.Scale(count);  // x_stats is now scaled by count.
    x2_stats.Scale(count);  // x2_stats is now scaled by count.
    clusterable_vec.push_back(new GaussClusterable(x_stats, x2_stats, min_var,
                                                   count));
  }
  if (clusterable_vec.size() <= target_components) {
    KALDI_WARN << "Not doing clustering phase since lost too many Gaussians "
               << "due to zero weight. Warning: zero-weight Gaussians are "
               << "still there.";
    DeletePointers(&clusterable_vec);
    return;
  } else {
    std::vector<Clusterable*> clusters;
    ClusterKMeans(clusterable_vec,
                  target_components,
                  &clusters, NULL, cfg);
    Resize(clusters.size(), Dim());
    for (int32 g = 0; g < static_cast<int32>(clusters.size()); g++) {
      GaussClusterable *gc = static_cast<GaussClusterable*>(clusters[g]);
      weights_(g) = gc->count();
      SubVector<BaseFloat> inv_var(inv_vars_, g),
          mean_invvar(means_invvars_, g);
      inv_var.CopyFromVec(gc->x2_stats());
      inv_var.Scale(1.0 / gc->count());  // inv_var is now the var + mean^2
      mean_invvar.CopyFromVec(gc->x_stats());
      mean_invvar.Scale(1.0 / gc->count());  // mean_invvar is now the mean.
      inv_var.AddVec2(-1.0, mean_invvar);  // subtract mean^2; inv_var is now the var
      inv_var.InvertElements();  // inv_var is now the inverse var.
      mean_invvar.MulElements(inv_var);  // mean_invvar is now mean * inverse var.
    }
    ComputeGconsts();
    DeletePointers(&clusterable_vec);
    DeletePointers(&clusters);
  }
}

void DiagGmm::Merge(int32 target_components, std::vector<int32> *history) {
  if (target_components <= 0 || NumGauss() < target_components) {
    KALDI_ERR << "Invalid argument for target number of Gaussians (="
              << target_components << "), #Gauss = " << NumGauss();
  }
  if (NumGauss() == target_components) {
    KALDI_VLOG(2) << "No components merged, as target (" << target_components
                  << ") = total.";
    return;  // Nothing to do.
  }

  int32 num_comp = NumGauss(), dim = Dim();

  if (target_components == 1) {  // global mean and variance
    Vector<BaseFloat> weights(weights_);
    // Undo variance inversion and multiplication of mean by inv var.
    Matrix<BaseFloat> vars(inv_vars_);
    Matrix<BaseFloat> means(means_invvars_);
    vars.InvertElements();
    means.MulElements(vars);
    // add means square to variances; get second-order stats
    for (int32 i = 0; i < num_comp; i++) {
      vars.Row(i).AddVec2(1.0, means.Row(i));
    }

    // Slightly more efficient than calling this->Resize(1, dim)
    gconsts_.Resize(1);
    weights_.Resize(1);
    means_invvars_.Resize(1, dim);
    inv_vars_.Resize(1, dim);

    for (int32 i = 0; i < num_comp; i++) {
      weights_(0) += weights(i);
      means_invvars_.Row(0).AddVec(weights(i), means.Row(i));
      inv_vars_.Row(0).AddVec(weights(i), vars.Row(i));
    }
    if (!ApproxEqual(weights_(0), 1.0, 1e-6)) {
      KALDI_WARN << "Weights sum to " << weights_(0) << ": rescaling.";
      means_invvars_.Scale(weights_(0));
      inv_vars_.Scale(weights_(0));
      weights_(0) = 1.0;
    }
    inv_vars_.Row(0).AddVec2(-1.0, means_invvars_.Row(0));
    inv_vars_.InvertElements();
    means_invvars_.MulElements(inv_vars_);
    ComputeGconsts();
    return;
  }

  // If more than 1 merged component is required, use the hierarchical
  // clustering of components that lead to the smallest decrease in likelihood.
  std::vector<bool> discarded_component(num_comp);
  Vector<BaseFloat> logdet(num_comp);   // logdet for each component
  for (int32 i = 0; i < num_comp; i++) {
    discarded_component[i] = false;
    for (int32 d = 0; d < dim; d++) {
      logdet(i) += 0.5 * Log(inv_vars_(i, d));  // +0.5 because var is inverted
    }
  }

  // Undo variance inversion and multiplication of mean by this
  // Makes copy of means and vars for all components - memory inefficient?
  Matrix<BaseFloat> vars(inv_vars_);
  Matrix<BaseFloat> means(means_invvars_);
  vars.InvertElements();
  means.MulElements(vars);

  // add means square to variances; get second-order stats
  // (normalized by zero-order stats)
  for (int32 i = 0; i < num_comp; i++) {
    vars.Row(i).AddVec2(1.0, means.Row(i));
  }

  // compute change of likelihood for all combinations of components
  SpMatrix<BaseFloat> delta_like(num_comp);
  for (int32 i = 0; i < num_comp; i++) {
    for (int32 j = 0; j < i; j++) {
      BaseFloat w1 = weights_(i), w2 = weights_(j), w_sum = w1 + w2;
      BaseFloat merged_logdet = merged_components_logdet(w1, w2,
        means.Row(i), means.Row(j), vars.Row(i), vars.Row(j));
      delta_like(i, j) = w_sum * merged_logdet
        - w1 * logdet(i) - w2 * logdet(j);
    }
  }

  // Merge components with smallest impact on the loglike
  for (int32 removed = 0; removed < num_comp - target_components; removed++) {
    // Search for the least significant change in likelihood
    // (maximum of negative delta_likes)
    BaseFloat max_delta_like = -std::numeric_limits<BaseFloat>::max();
    int32 max_i = -1, max_j = -1;
    for (int32 i = 0; i < NumGauss(); i++) {
      if (discarded_component[i]) continue;
      for (int32 j = 0; j < i; j++) {
        if (discarded_component[j]) continue;
        if (delta_like(i, j) > max_delta_like) {
          max_delta_like = delta_like(i, j);
          max_i = i;
          max_j = j;
        }
      }
    }

    // make sure that different components will be merged
    KALDI_ASSERT(max_i != max_j && max_i != -1 && max_j != -1);

    // remember the merge candidates
    if (history != NULL) {
      history->push_back(max_i);
      history->push_back(max_j);
    }

    // Merge components
    BaseFloat w1 = weights_(max_i), w2 = weights_(max_j);
    BaseFloat w_sum = w1 + w2;
    // merge means
    means.Row(max_i).AddVec(w2/w1, means.Row(max_j));
    means.Row(max_i).Scale(w1/w_sum);
    // merge vars
    vars.Row(max_i).AddVec(w2/w1, vars.Row(max_j));
    vars.Row(max_i).Scale(w1/w_sum);
    // merge weights
    weights_(max_i) = w_sum;

    // Update gmm for merged component
    // copy second-order stats (normalized by zero-order stats)
    inv_vars_.Row(max_i).CopyFromVec(vars.Row(max_i));
    // centralize
    inv_vars_.Row(max_i).AddVec2(-1.0, means.Row(max_i));
    // invert
    inv_vars_.Row(max_i).InvertElements();
    // copy first-order stats (normalized by zero-order stats)
    means_invvars_.Row(max_i).CopyFromVec(means.Row(max_i));
    // multiply by inv_vars
    means_invvars_.Row(max_i).MulElements(inv_vars_.Row(max_i));

    // Update logdet for merged component
    logdet(max_i) = 0.0;
    for (int32 d = 0; d < dim; d++) {
      logdet(max_i) += 0.5 * Log(inv_vars_(max_i, d));
      // +0.5 because var is inverted
    }

    // Label the removed component as discarded
    discarded_component[max_j] = true;

    // Update delta_like for merged component
    for (int32 j = 0; j < num_comp; j++) {
      if ((j == max_i) || (discarded_component[j])) continue;
      BaseFloat w1 = weights_(max_i),
                w2 = weights_(j),
                w_sum = w1 + w2;
      BaseFloat merged_logdet = merged_components_logdet(w1, w2,
          means.Row(max_i), means.Row(j), vars.Row(max_i), vars.Row(j));
      delta_like(max_i, j) = w_sum * merged_logdet - w1 * logdet(max_i)
          - w2 * logdet(j);
      // doesn't respect lower triangular indeces,
      // relies on implicitly performed swap of coordinates if necessary
    }
  }

  // Remove the consumed components
  int32 m = 0;
  for (int32 i = 0; i < num_comp; i++) {
    if (discarded_component[i]) {
      weights_.RemoveElement(m);
      means_invvars_.RemoveRow(m);
      inv_vars_.RemoveRow(m);
    } else {
      ++m;
    }
  }

  ComputeGconsts();
}

BaseFloat DiagGmm::merged_components_logdet(BaseFloat w1, BaseFloat w2,
                                            const VectorBase<BaseFloat> &f1,
                                            const VectorBase<BaseFloat> &f2,
                                            const VectorBase<BaseFloat> &s1,
                                            const VectorBase<BaseFloat> &s2)
                                            const {
  int32 dim = f1.Dim();
  Vector<BaseFloat> tmp_mean(dim);
  Vector<BaseFloat> tmp_var(dim);

  BaseFloat w_sum = w1 + w2;
  tmp_mean.CopyFromVec(f1);
  tmp_mean.AddVec(w2/w1, f2);
  tmp_mean.Scale(w1/w_sum);
  tmp_var.CopyFromVec(s1);
  tmp_var.AddVec(w2/w1, s2);
  tmp_var.Scale(w1/w_sum);
  tmp_var.AddVec2(-1.0, tmp_mean);
  BaseFloat merged_logdet = 0.0;
  for (int32 d = 0; d < dim; d++) {
    merged_logdet -= 0.5 * Log(tmp_var(d));
    // -0.5 because var is not inverted
  }
  return merged_logdet;
}

BaseFloat DiagGmm::ComponentLogLikelihood(const VectorBase<BaseFloat> &data,
                                          int32 comp_id) const {
  if (!valid_gconsts_)
    KALDI_ERR << "Must call ComputeGconsts() before computing likelihood";
  if (static_cast<int32>(data.Dim()) != Dim()) {
    KALDI_ERR << "DiagGmm::ComponentLogLikelihood, dimension "
        << "mismatch " << (data.Dim()) << " vs. "<< (Dim());
  }
  BaseFloat loglike;
  Vector<BaseFloat> data_sq(data);
  data_sq.ApplyPow(2.0);

  // loglike =  means * inv(vars) * data.
  loglike = VecVec(means_invvars_.Row(comp_id), data);
  // loglike += -0.5 * inv(vars) * data_sq.
  loglike -= 0.5 * VecVec(inv_vars_.Row(comp_id), data_sq);
  return loglike + gconsts_(comp_id);
}

// Gets likelihood of data given this.
BaseFloat DiagGmm::LogLikelihood(const VectorBase<BaseFloat> &data) const {
  if (!valid_gconsts_)
    KALDI_ERR << "Must call ComputeGconsts() before computing likelihood";
  Vector<BaseFloat> loglikes;
  LogLikelihoods(data, &loglikes);
  BaseFloat log_sum = loglikes.LogSumExp();
  if (KALDI_ISNAN(log_sum) || KALDI_ISINF(log_sum))
    KALDI_ERR << "Invalid answer (overflow or invalid variances/features?)";
  return log_sum;
}

void DiagGmm::LogLikelihoods(const VectorBase<BaseFloat> &data,
                             Vector<BaseFloat> *loglikes) const {
  loglikes->Resize(gconsts_.Dim(), kUndefined);
  loglikes->CopyFromVec(gconsts_);
  if (data.Dim() != Dim()) {
    KALDI_ERR << "DiagGmm::ComponentLogLikelihood, dimension "
              << "mismatch " << data.Dim() << " vs. "<< Dim();
  }
  Vector<BaseFloat> data_sq(data);
  data_sq.ApplyPow(2.0);

  // loglikes +=  means * inv(vars) * data.
  loglikes->AddMatVec(1.0, means_invvars_, kNoTrans, data, 1.0);
  // loglikes += -0.5 * inv(vars) * data_sq.
  loglikes->AddMatVec(-0.5, inv_vars_, kNoTrans, data_sq, 1.0);
}


void DiagGmm::LogLikelihoods(const MatrixBase<BaseFloat> &data,
                             Matrix<BaseFloat> *loglikes) const {
  KALDI_ASSERT(data.NumRows() != 0);
  loglikes->Resize(data.NumRows(), gconsts_.Dim(), kUndefined);
  loglikes->CopyRowsFromVec(gconsts_);
  if (data.NumCols() != Dim()) {
    KALDI_ERR << "DiagGmm::ComponentLogLikelihood, dimension "
              << "mismatch " << data.NumCols() << " vs. "<< Dim();
  }
  Matrix<BaseFloat> data_sq(data);
  data_sq.ApplyPow(2.0);

  // loglikes +=  means * inv(vars) * data.
  loglikes->AddMatMat(1.0, data, kNoTrans, means_invvars_, kTrans, 1.0);
  // loglikes += -0.5 * inv(vars) * data_sq.
  loglikes->AddMatMat(-0.5, data_sq, kNoTrans, inv_vars_, kTrans, 1.0);
}



void DiagGmm::LogLikelihoodsPreselect(const VectorBase<BaseFloat> &data,
                                      const std::vector<int32> &indices,
                                      Vector<BaseFloat> *loglikes) const {
  KALDI_ASSERT(data.Dim() == Dim());
  Vector<BaseFloat> data_sq(data);
  data_sq.ApplyPow(2.0);

  int32 num_indices = static_cast<int32>(indices.size());
  loglikes->Resize(num_indices, kUndefined);
  if (indices.back() + 1 - indices.front() == num_indices) {
    // A special (but common) case when the indices form a contiguous range.
    int32 start_idx = indices.front();
    loglikes->CopyFromVec(SubVector<BaseFloat>(gconsts_, start_idx, num_indices));
    // loglikes +=  means * inv(vars) * data.
    SubMatrix<BaseFloat> means_invvars_sub(means_invvars_, start_idx, num_indices,
                                           0, Dim());
    loglikes->AddMatVec(1.0, means_invvars_sub, kNoTrans, data, 1.0);
    SubMatrix<BaseFloat> inv_vars_sub(inv_vars_, start_idx, num_indices,
                                      0, Dim());
    // loglikes += -0.5 * inv(vars) * data_sq.
    loglikes->AddMatVec(-0.5, inv_vars_sub, kNoTrans, data_sq, 1.0);
  } else {
    for (int32 i = 0; i < num_indices; i++) {
      int32 idx = indices[i];  // The Gaussian index.
      BaseFloat this_loglike =
          gconsts_(idx) + VecVec(means_invvars_.Row(idx), data)
          - 0.5*VecVec(inv_vars_.Row(idx), data_sq);
      (*loglikes)(i) = this_loglike;
    }
  }
}



// Gets likelihood of data given this. Also provides per-Gaussian posteriors.
BaseFloat DiagGmm::ComponentPosteriors(const VectorBase<BaseFloat> &data,
                                       Vector<BaseFloat> *posterior) const {
  if (!valid_gconsts_)
    KALDI_ERR << "Must call ComputeGconsts() before computing likelihood";
  if (posterior == NULL) KALDI_ERR << "NULL pointer passed as return argument.";
  Vector<BaseFloat> loglikes;
  LogLikelihoods(data, &loglikes);
  BaseFloat log_sum = loglikes.ApplySoftMax();
  if (KALDI_ISNAN(log_sum) || KALDI_ISINF(log_sum))
    KALDI_ERR << "Invalid answer (overflow or invalid variances/features?)";
  if (posterior->Dim() != loglikes.Dim())
    posterior->Resize(loglikes.Dim());
  posterior->CopyFromVec(loglikes);
  return log_sum;
}

void DiagGmm::RemoveComponent(int32 gauss, bool renorm_weights) {
  KALDI_ASSERT(gauss < NumGauss());
  if (NumGauss() == 1)
    KALDI_ERR << "Attempting to remove the only remaining component.";
  weights_.RemoveElement(gauss);
  gconsts_.RemoveElement(gauss);
  means_invvars_.RemoveRow(gauss);
  inv_vars_.RemoveRow(gauss);
  BaseFloat sum_weights = weights_.Sum();
  if (renorm_weights) {
    weights_.Scale(1.0/sum_weights);
    valid_gconsts_ = false;
  }
}

void DiagGmm::RemoveComponents(const std::vector<int32> &gauss_in,
                               bool renorm_weights) {
  std::vector<int32> gauss(gauss_in);
  std::sort(gauss.begin(), gauss.end());
  KALDI_ASSERT(IsSortedAndUniq(gauss));
  // If efficiency is later an issue, will code this specially (unlikely).
  for (size_t i = 0; i < gauss.size(); i++) {
    RemoveComponent(gauss[i], renorm_weights);
    for (size_t j = i + 1; j < gauss.size(); j++)
      gauss[j]--;
  }
}

void DiagGmm::Interpolate(BaseFloat rho, const DiagGmm &source,
                          GmmFlagsType flags) {
  KALDI_ASSERT(NumGauss() == source.NumGauss());
  KALDI_ASSERT(Dim() == source.Dim());

  DiagGmmNormal us(*this);
  DiagGmmNormal them(source);

  if (flags & kGmmWeights) {
    us.weights_.Scale(1.0 - rho);
    us.weights_.AddVec(rho, them.weights_);
    us.weights_.Scale(1.0 / us.weights_.Sum());
  }

  if (flags & kGmmMeans) {
    us.means_.Scale(1.0 - rho);
    us.means_.AddMat(rho, them.means_);
  }

  if (flags & kGmmVariances) {
    us.vars_.Scale(1.0 - rho);
    us.vars_.AddMat(rho, them.vars_);
  }

  us.CopyToDiagGmm(this);
  ComputeGconsts();
}

void DiagGmm::Interpolate(BaseFloat rho, const FullGmm &source,
                          GmmFlagsType flags) {
  KALDI_ASSERT(NumGauss() == source.NumGauss());
  KALDI_ASSERT(Dim() == source.Dim());
  DiagGmmNormal us(*this);
  FullGmmNormal them(source);

  if (flags & kGmmWeights) {
    us.weights_.Scale(1.0 - rho);
    us.weights_.AddVec(rho, them.weights_);
    us.weights_.Scale(1.0 / us.weights_.Sum());
  }

  if (flags & kGmmMeans) {
    us.means_.Scale(1.0 - rho);
    us.means_.AddMat(rho, them.means_);
  }

  if (flags & kGmmVariances) {
    for (int32 i = 0; i < NumGauss(); i++) {
      us.vars_.Scale(1. - rho);
      Vector<double> diag(Dim());
      for (int32 j = 0; j < Dim(); j++)
        diag(j) = them.vars_[i](j, j);
      us.vars_.Row(i).AddVec(rho, diag);
    }
  }

  us.CopyToDiagGmm(this);
  ComputeGconsts();
}

void DiagGmm::Write(std::ostream &out_stream, bool binary) const {
  if (!valid_gconsts_)
    KALDI_ERR << "Must call ComputeGconsts() before writing the model.";
  WriteToken(out_stream, binary, "<DiagGMM>");
  if (!binary) out_stream << "\n";
  WriteToken(out_stream, binary, "<GCONSTS>");
  gconsts_.Write(out_stream, binary);
  WriteToken(out_stream, binary, "<WEIGHTS>");
  weights_.Write(out_stream, binary);
  WriteToken(out_stream, binary, "<MEANS_INVVARS>");
  means_invvars_.Write(out_stream, binary);
  WriteToken(out_stream, binary, "<INV_VARS>");
  inv_vars_.Write(out_stream, binary);
  WriteToken(out_stream, binary, "</DiagGMM>");
  if (!binary) out_stream << "\n";
}

std::ostream & operator <<(std::ostream & os,
                           const kaldi::DiagGmm &gmm) {
  gmm.Write(os, false);
  return os;
}

void DiagGmm::Read(std::istream &is, bool binary) {
//  ExpectToken(is, binary, "<DiagGMMBegin>");
  std::string token;
  ReadToken(is, binary, &token);
  // <DiagGMMBegin> is for compatibility. Will be deleted later
  if (token != "<DiagGMMBegin>" && token != "<DiagGMM>")
    KALDI_ERR << "Expected <DiagGMM>, got " << token;
  ReadToken(is, binary, &token);
  if (token == "<GCONSTS>") {  // The gconsts are optional.
    gconsts_.Read(is, binary);
    ExpectToken(is, binary, "<WEIGHTS>");
  } else {
    if (token != "<WEIGHTS>")
      KALDI_ERR << "DiagGmm::Read, expected <WEIGHTS> or <GCONSTS>, got "
                << token;
  }
  weights_.Read(is, binary);
  ExpectToken(is, binary, "<MEANS_INVVARS>");
  means_invvars_.Read(is, binary);
  ExpectToken(is, binary, "<INV_VARS>");
  inv_vars_.Read(is, binary);
//  ExpectToken(is, binary, "<DiagGMMEnd>");
  ReadToken(is, binary, &token);
  // <DiagGMMEnd> is for compatibility. Will be deleted later
  if (token != "<DiagGMMEnd>" && token != "</DiagGMM>")
    KALDI_ERR << "Expected </DiagGMM>, got " << token;

  ComputeGconsts();  // safer option than trusting the read gconsts
}

std::istream & operator >>(std::istream &is, kaldi::DiagGmm &gmm) {
  gmm.Read(is, false);  // false == non-binary.
  return is;
}


/// Get gaussian selection information for one frame.
BaseFloat DiagGmm::GaussianSelection(const VectorBase<BaseFloat> &data,
                                     int32 num_gselect,
                                     std::vector<int32> *output) const {
  int32 num_gauss = NumGauss();
  Vector<BaseFloat> loglikes(num_gauss, kUndefined);
  output->clear();
  this->LogLikelihoods(data, &loglikes);

  BaseFloat thresh;
  if (num_gselect < num_gauss) {
    Vector<BaseFloat> loglikes_copy(loglikes);
    BaseFloat *ptr = loglikes_copy.Data();
    std::nth_element(ptr, ptr+num_gauss-num_gselect, ptr+num_gauss);
    thresh = ptr[num_gauss-num_gselect];
  } else {
    thresh = -std::numeric_limits<BaseFloat>::infinity();
  }
  BaseFloat tot_loglike = -std::numeric_limits<BaseFloat>::infinity();
  std::vector<std::pair<BaseFloat, int32> > pairs;
  for (int32 p = 0; p < num_gauss; p++) {
    if (loglikes(p) >= thresh) {
      pairs.push_back(std::make_pair(loglikes(p), p));
    }
  }
  std::sort(pairs.begin(), pairs.end(),
            std::greater<std::pair<BaseFloat, int32> >());
  for (int32 j = 0;
       j < num_gselect && j < static_cast<int32>(pairs.size());
       j++) {
    output->push_back(pairs[j].second);
    tot_loglike = LogAdd(tot_loglike, pairs[j].first);
  }
  KALDI_ASSERT(!output->empty());
  return tot_loglike;
}

BaseFloat DiagGmm::GaussianSelection(const MatrixBase<BaseFloat> &data,
                                     int32 num_gselect,
                                     std::vector<std::vector<int32> > *output) const {
  double ans = 0.0;
  int32 num_frames = data.NumRows(), num_gauss = NumGauss();

  int32 max_mem = 10000000; // Don't devote more than 10Mb to loglikes_mat;
                            // break up the utterance if needed.
  int32 mem_needed = num_frames * num_gauss * sizeof(BaseFloat);
  if (mem_needed > max_mem) {
    // Break into parts and recurse, we don't want to consume too
    // much memory.
    int32 num_parts = (mem_needed + max_mem - 1) / max_mem;
    int32 part_frames = (data.NumRows() + num_parts - 1) / num_parts;
    double tot_ans = 0.0;
    std::vector<std::vector<int32> > part_output;
    output->clear();
    output->resize(num_frames);
    for (int32 p = 0; p < num_parts; p++) {
      int32 start_frame = p * part_frames,
          this_num_frames = std::min(num_frames - start_frame, part_frames);
      SubMatrix<BaseFloat> data_part(data, start_frame, this_num_frames,
                                     0, data.NumCols());
      tot_ans += GaussianSelection(data_part, num_gselect, &part_output);
      for (int32 t = 0; t < this_num_frames; t++)
        (*output)[start_frame + t].swap(part_output[t]);
    }
    KALDI_ASSERT(!output->back().empty());
    return tot_ans;
  }
  
  KALDI_ASSERT(num_frames != 0);
  Matrix<BaseFloat> loglikes_mat(num_frames, num_gauss, kUndefined);
  this->LogLikelihoods(data, &loglikes_mat);
  
  output->clear();
  output->resize(num_frames);

  for (int32 i = 0; i < num_frames; i++) {
    SubVector<BaseFloat> loglikes(loglikes_mat, i);

    BaseFloat thresh;
    if (num_gselect < num_gauss) {
      Vector<BaseFloat> loglikes_copy(loglikes);
      BaseFloat *ptr = loglikes_copy.Data();
      std::nth_element(ptr, ptr+num_gauss-num_gselect, ptr+num_gauss);
      thresh = ptr[num_gauss-num_gselect];
    } else {
      thresh = -std::numeric_limits<BaseFloat>::infinity();
    }
    BaseFloat tot_loglike = -std::numeric_limits<BaseFloat>::infinity();
    std::vector<std::pair<BaseFloat, int32> > pairs;
    for (int32 p = 0; p < num_gauss; p++) {
      if (loglikes(p) >= thresh) {
        pairs.push_back(std::make_pair(loglikes(p), p));
      }
    }
    std::sort(pairs.begin(), pairs.end(),
              std::greater<std::pair<BaseFloat, int32> >());
    std::vector<int32> &this_output = (*output)[i];
    for (int32 j = 0;
         j < num_gselect && j < static_cast<int32>(pairs.size());
         j++) {
      this_output.push_back(pairs[j].second);
      tot_loglike = LogAdd(tot_loglike, pairs[j].first);
    }
    KALDI_ASSERT(!this_output.empty());
    ans += tot_loglike;
  }
  return ans;
}



BaseFloat DiagGmm::GaussianSelectionPreselect(
    const VectorBase<BaseFloat> &data,
    const std::vector<int32> &preselect,
    int32 num_gselect,
    std::vector<int32> *output) const {
  static bool warned_size = false;
  int32 preselect_sz = preselect.size();
  int32 this_num_gselect = std::min(num_gselect, preselect_sz);
  if (preselect_sz <= num_gselect && !warned_size) {
    warned_size = true;
    KALDI_WARN << "Preselect size is less or equal to than final size, "
               << "doing nothing: " << preselect_sz << " < " <<  num_gselect
               << " [won't warn again]";
  }
  Vector<BaseFloat> loglikes(preselect_sz);
  LogLikelihoodsPreselect(data, preselect, &loglikes);
  
  Vector<BaseFloat> loglikes_copy(loglikes);
  BaseFloat *ptr = loglikes_copy.Data();
  std::nth_element(ptr, ptr+preselect_sz-this_num_gselect,
                   ptr+preselect_sz);
  BaseFloat thresh = ptr[preselect_sz-this_num_gselect];

  BaseFloat tot_loglike = -std::numeric_limits<BaseFloat>::infinity();
  // we want the output sorted from best likelihood to worse
  // (so we can prune further without the model)...
  std::vector<std::pair<BaseFloat, int32> > pairs;
  for (int32 p = 0; p < preselect_sz; p++)
    if (loglikes(p) >= thresh)
      pairs.push_back(std::make_pair(loglikes(p), preselect[p]));
  std::sort(pairs.begin(), pairs.end(),
            std::greater<std::pair<BaseFloat, int32> >());
  output->clear();
  for (int32 j = 0;
       j < this_num_gselect && j < static_cast<int32>(pairs.size());
       j++) {
    output->push_back(pairs[j].second);
    tot_loglike = LogAdd(tot_loglike, pairs[j].first);
  }
  KALDI_ASSERT(!output->empty());
  return tot_loglike;
}

void DiagGmm::CopyFromNormal(const DiagGmmNormal &diag_gmm_normal) {
  diag_gmm_normal.CopyToDiagGmm(this);
}

void DiagGmm::Generate(VectorBase<BaseFloat> *output) {
  KALDI_ASSERT(static_cast<int32>(output->Dim()) == Dim());
  BaseFloat tot = weights_.Sum();
  KALDI_ASSERT(tot > 0.0);
  double r = tot * RandUniform() * 0.99999;
  int32 i = 0;
  double sum = 0.0;
  while (sum + weights_(i) < r) {
    sum += weights_(i);
    i++;
    KALDI_ASSERT(i < static_cast<int32>(weights_.Dim()));
  }
  // now i is the index of the Gaussian we chose.
  SubVector<BaseFloat> inv_var(inv_vars_, i),
      mean_invvar(means_invvars_, i);
  for (int32 d = 0; d < inv_var.Dim(); d++) {
    BaseFloat stddev = 1.0 / sqrt(inv_var(d)),
        mean = mean_invvar(d) / inv_var(d);
    (*output)(d) = mean + RandGauss() * stddev;
  }
}

DiagGmm::DiagGmm(const GaussClusterable &gc,
                 BaseFloat var_floor): valid_gconsts_(false) {
  Vector<BaseFloat> x (gc.x_stats());
  Vector<BaseFloat> x2 (gc.x2_stats());
  BaseFloat count =  gc.count();
  KALDI_ASSERT(count > 0.0);
  this->Resize(1, x.Dim());
  x.Scale(1.0/count);
  x2.Scale(1.0/count);
  x2.AddVec2(-1.0, x);  // subtract mean^2.
  x2.ApplyFloor(var_floor);
  x2.InvertElements();  // get inv-var.
  KALDI_ASSERT(x2.Min() > 0);
  Matrix<BaseFloat> mean(1, x.Dim());
  mean.Row(0).CopyFromVec(x);
  Matrix<BaseFloat> inv_var(1, x.Dim());
  inv_var.Row(0).CopyFromVec(x2);
  this->SetInvVarsAndMeans(inv_var, mean);
  Vector<BaseFloat> weights(1);
  weights(0) = 1.0;
  this->SetWeights(weights);
  this->ComputeGconsts();
}

}  // End namespace kaldi