// transform/decodable-am-diag-gmm-regtree.cc

// Copyright 2009-2011  Saarland University;  Lukas Burget
//                2013  Johns Hopkins Universith (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 <vector>
using std::vector;

#include "transform/decodable-am-diag-gmm-regtree.h"

namespace kaldi {


BaseFloat DecodableAmDiagGmmRegtreeFmllr::LogLikelihoodZeroBased(int32 frame,
                                                          int32 state) {
  KALDI_ASSERT(frame < NumFramesReady() && frame >= 0);
  KALDI_ASSERT(state < NumIndices() && state >= 0);

  if (!valid_logdets_) {
    logdets_.Resize(fmllr_xform_.NumRegClasses());
    fmllr_xform_.GetLogDets(&logdets_);
    valid_logdets_ = true;
  }

  if (log_like_cache_[state].hit_time == frame) {
    return log_like_cache_[state].log_like;  // return cached value, if found
  }

  const DiagGmm &pdf = acoustic_model_.GetPdf(state);
  const VectorBase<BaseFloat> &data = feature_matrix_.Row(frame);

  // check if everything is in order
  if (pdf.Dim() != data.Dim()) {
    KALDI_ERR << "Dim mismatch: data dim = "  << data.Dim()
        << " vs. model dim = " << pdf.Dim();
  }
  if (!pdf.valid_gconsts()) {
    KALDI_ERR << "State "  << (state)  << ": Must call ComputeGconsts() "
        "before computing likelihood.";
  }

  if (frame != previous_frame_) {  // cache the transformed & squared stats.
    fmllr_xform_.TransformFeature(data, &xformed_data_);
    xformed_data_squared_ = xformed_data_;
    vector< Vector <BaseFloat> >::iterator it = xformed_data_squared_.begin(),
        end = xformed_data_squared_.end();
    for (; it != end; ++it) { it->ApplyPow(2.0); }
    previous_frame_ = frame;
  }

  Vector<BaseFloat> loglikes(pdf.gconsts());  // need to recreate for each pdf
  int32 baseclass, regclass;
  for (int32 comp_id = 0, num_comp = pdf.NumGauss(); comp_id < num_comp;
      ++comp_id) {
    baseclass = regtree_.Gauss2BaseclassId(state, comp_id);
    regclass = fmllr_xform_.Base2RegClass(baseclass);
    // loglikes +=  means * inv(vars) * data.
    loglikes(comp_id) += VecVec(pdf.means_invvars().Row(comp_id),
                                xformed_data_[regclass]);
    // loglikes += -0.5 * inv(vars) * data_sq.
    loglikes(comp_id) -= 0.5 * VecVec(pdf.inv_vars().Row(comp_id),
                                      xformed_data_squared_[regclass]);
    loglikes(comp_id) += logdets_(regclass);
  }

  BaseFloat log_sum = loglikes.LogSumExp(log_sum_exp_prune_);
  if (KALDI_ISNAN(log_sum) || KALDI_ISINF(log_sum))
    KALDI_ERR << "Invalid answer (overflow or invalid variances/features?)";

  log_like_cache_[state].log_like = log_sum;
  log_like_cache_[state].hit_time = frame;

  return log_sum;
}

DecodableAmDiagGmmRegtreeMllr::~DecodableAmDiagGmmRegtreeMllr() {
  DeletePointers(&xformed_mean_invvars_);
  DeletePointers(&xformed_gconsts_);
}


void DecodableAmDiagGmmRegtreeMllr::InitCache() {
  if (xformed_mean_invvars_.size() != 0)
    DeletePointers(&xformed_mean_invvars_);
  if (xformed_gconsts_.size() != 0)
    DeletePointers(&xformed_gconsts_);
  int32 num_pdfs = acoustic_model_.NumPdfs();
  xformed_mean_invvars_.resize(num_pdfs);
  xformed_gconsts_.resize(num_pdfs);
  is_cached_.resize(num_pdfs, false);
  ResetLogLikeCache();
}


// This is almost the same code as DiagGmm::ComputeGconsts, except that
// means are used instead of means * inv(vars). This saves some computation.
static void ComputeGconsts(const VectorBase<BaseFloat> &weights,
                           const MatrixBase<BaseFloat> &means,
                           const MatrixBase<BaseFloat> &inv_vars,
                           VectorBase<BaseFloat> *gconsts_out) {
  int32 num_gauss = weights.Dim();
  int32 dim = means.NumCols();
  KALDI_ASSERT(means.NumRows() == num_gauss
      && inv_vars.NumRows() == num_gauss && inv_vars.NumCols() == dim);
  KALDI_ASSERT(gconsts_out->Dim() == num_gauss);

  BaseFloat offset = -0.5 * M_LOG_2PI * dim;  // constant term in gconst.
  int32 num_bad = 0;

  for (int32 gauss = 0; gauss < num_gauss; gauss++) {
    KALDI_ASSERT(weights(gauss) >= 0);  // Cannot have negative weights.
    BaseFloat gc = Log(weights(gauss)) + offset;  // May be -inf if weights == 0
    for (int32 d = 0; d < dim; d++) {
      gc += 0.5 * Log(inv_vars(gauss, d)) - 0.5 * means(gauss, d)
        * means(gauss, d) * inv_vars(gauss, d);  // diff from DiagGmm version.
    }

    if (KALDI_ISNAN(gc)) {  // negative infinity is OK but NaN is not acceptable
      KALDI_ERR << "At component "  << gauss
                << ", 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_out)(gauss) = gc;
  }
  if (num_bad > 0)
    KALDI_WARN << num_bad << " unusable components found while computing "
               << "gconsts.";
}


const Matrix<BaseFloat>& DecodableAmDiagGmmRegtreeMllr::GetXformedMeanInvVars(
    int32 state) {
  if (is_cached_[state]) {  // found in cache
    KALDI_ASSERT(xformed_mean_invvars_[state] != NULL);
    KALDI_VLOG(3) << "For PDF index " << state << ": transformed means "
                  << "found in cache.";
    return *xformed_mean_invvars_[state];
  } else {  // transform the means and cache them
    KALDI_ASSERT(xformed_mean_invvars_[state] == NULL);
    KALDI_VLOG(3) << "For PDF index " << state << ": transforming means.";
    int32 num_gauss = acoustic_model_.GetPdf(state).NumGauss(),
        dim = acoustic_model_.Dim();
    const Vector<BaseFloat> &weights = acoustic_model_.GetPdf(state).weights();
    const Matrix<BaseFloat> &invvars = acoustic_model_.GetPdf(state).inv_vars();
    xformed_mean_invvars_[state] = new Matrix<BaseFloat>(num_gauss, dim);
    mllr_xform_.GetTransformedMeans(regtree_, acoustic_model_, state,
                                    xformed_mean_invvars_[state]);
    xformed_gconsts_[state] = new Vector<BaseFloat>(num_gauss);
    // At this point, the transformed means haven't been multiplied with
    // the inv vars, and they are used to compute gconsts first.
    ComputeGconsts(weights, *xformed_mean_invvars_[state], invvars,
                   xformed_gconsts_[state]);
    // Finally, multiply the transformed means with the inv vars.
    xformed_mean_invvars_[state]->MulElements(invvars);
    is_cached_[state] = true;
    return *xformed_mean_invvars_[state];
  }
}

const Vector<BaseFloat>& DecodableAmDiagGmmRegtreeMllr::GetXformedGconsts(
    int32 state) {
  if (!is_cached_[state]) {
    KALDI_ERR << "GConsts not cached for state: " << state << ". Must call "
              << "GetXformedMeanInvVars() first.";
  }
  KALDI_ASSERT(xformed_gconsts_[state] != NULL);
  return *xformed_gconsts_[state];
}

BaseFloat DecodableAmDiagGmmRegtreeMllr::LogLikelihoodZeroBased(int32 frame,
                                                                int32 state) {
//  KALDI_ERR << "Function not completely implemented yet.";
  KALDI_ASSERT(frame < NumFramesReady() && frame >= 0);
  KALDI_ASSERT(state < NumIndices() && state >= 0);

  if (log_like_cache_[state].hit_time == frame) {
    return log_like_cache_[state].log_like;  // return cached value, if found
  }

  const DiagGmm &pdf = acoustic_model_.GetPdf(state);
  const VectorBase<BaseFloat> &data = feature_matrix_.Row(frame);

  // check if everything is in order
  if (pdf.Dim() != data.Dim()) {
    KALDI_ERR << "Dim mismatch: data dim = "  << data.Dim()
        << " vs. model dim = " << pdf.Dim();
  }

  if (frame != previous_frame_) {  // cache the squared stats.
    data_squared_.CopyFromVec(feature_matrix_.Row(frame));
    data_squared_.ApplyPow(2.0);
    previous_frame_ = frame;
  }

  const Matrix<BaseFloat> &means_invvars = GetXformedMeanInvVars(state);
  const Vector<BaseFloat> &gconsts = GetXformedGconsts(state);

  Vector<BaseFloat> loglikes(gconsts);  // need to recreate for each pdf
  // 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, pdf.inv_vars(), kNoTrans, data_squared_, 1.0);

  BaseFloat log_sum = loglikes.LogSumExp(log_sum_exp_prune_);
  if (KALDI_ISNAN(log_sum) || KALDI_ISINF(log_sum))
    KALDI_ERR << "Invalid answer (overflow or invalid variances/features?)";

  log_like_cache_[state].log_like = log_sum;
  log_like_cache_[state].hit_time = frame;

  return log_sum;
}

}  // namespace kaldi