// gmm/am-diag-gmm.cc
  
  // Copyright 2012   Arnab Ghoshal  Johns Hopkins University (Author: Daniel Povey)  Karel Vesely
  // Copyright 2009-2011  Saarland University;  Microsoft Corporation;
  //                      Georg Stemmer
  
  // 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 <queue>
  #include <string>
  #include <vector>
  using std::vector;
  
  #include "gmm/am-diag-gmm.h"
  #include "util/stl-utils.h"
  #include "tree/clusterable-classes.h"
  #include "tree/cluster-utils.h"
  
  namespace kaldi {
  
  AmDiagGmm::~AmDiagGmm() {
    DeletePointers(&densities_);
  }
  
  void AmDiagGmm::Init(const DiagGmm &proto, int32 num_pdfs) {
    if (densities_.size() != 0) {
      KALDI_WARN << "Init() called on a non-empty object. Contents will be "
          "overwritten";
      DeletePointers(&densities_);
    }
    if (num_pdfs == 0) {
      KALDI_WARN << "Init() called with number of pdfs = 0. Will do nothing.";
      return;
    }
  
    densities_.resize(num_pdfs, NULL);
    for (vector<DiagGmm*>::iterator itr = densities_.begin(),
        end = densities_.end(); itr != end; ++itr) {
      *itr = new DiagGmm();
      (*itr)->CopyFromDiagGmm(proto);
    }
  }
  
  void AmDiagGmm::AddPdf(const DiagGmm &gmm) {
    if (densities_.size() != 0)  // not the first gmm
      KALDI_ASSERT(gmm.Dim() == this->Dim());
  
    DiagGmm *gmm_ptr = new DiagGmm();
    gmm_ptr->CopyFromDiagGmm(gmm);
    densities_.push_back(gmm_ptr);
  }
  
  void AmDiagGmm::RemovePdf(int32 pdf_index) {
    KALDI_ASSERT(static_cast<size_t>(pdf_index) < densities_.size());
    delete densities_[pdf_index];
    densities_.erase(densities_.begin() + pdf_index);
  }
  
  int32 AmDiagGmm::NumGauss() const {
    int32 ans = 0;
    for (size_t i = 0; i < densities_.size(); i++)
      ans += densities_[i]->NumGauss();
    return ans;
  }
  
  void AmDiagGmm::CopyFromAmDiagGmm(const AmDiagGmm &other) {
    if (densities_.size() != 0) {
      DeletePointers(&densities_);
    }
    densities_.resize(other.NumPdfs(), NULL);
    for (int32 i = 0, end = densities_.size(); i < end; i++) {
      densities_[i] = new DiagGmm();
      densities_[i]->CopyFromDiagGmm(*other.densities_[i]);
    }
  }
  
  int32 AmDiagGmm::ComputeGconsts() {
    int32 num_bad = 0;
    for (std::vector<DiagGmm*>::iterator itr = densities_.begin(),
        end = densities_.end(); itr != end; ++itr) {
      num_bad += (*itr)->ComputeGconsts();
    }
    if (num_bad > 0)
      KALDI_WARN << "Found " << num_bad << " Gaussian components.";
    return num_bad;
  }
  
  
  void AmDiagGmm::SplitByCount(const Vector<BaseFloat> &state_occs,
                               int32 target_components,
                               float perturb_factor, BaseFloat power,
                               BaseFloat min_count) {
    int32 gauss_at_start = NumGauss();
    std::vector<int32> targets;
    GetSplitTargets(state_occs, target_components, power,
                    min_count, &targets);
  
    for (int32 i = 0; i < NumPdfs(); i++) {
      if (densities_[i]->NumGauss() < targets[i])
        densities_[i]->Split(targets[i], perturb_factor);
    }
  
    KALDI_LOG << "Split " << NumPdfs() << " states with target = "
              << target_components << ", power = " << power
              << ", perturb_factor = " << perturb_factor
              << " and min_count = " << min_count
              << ", split #Gauss from " << gauss_at_start << " to "
              << NumGauss();
  }
  
  
  void AmDiagGmm::MergeByCount(const Vector<BaseFloat> &state_occs,
                               int32 target_components,
                               BaseFloat power,
                               BaseFloat min_count) {
    int32 gauss_at_start = NumGauss();
    std::vector<int32> targets;
    GetSplitTargets(state_occs, target_components,
                    power, min_count, &targets);
  
    for (int32 i = 0; i < NumPdfs(); i++) {
      if (targets[i] == 0) targets[i] = 1;  // can't merge below 1.
      if (densities_[i]->NumGauss() > targets[i])
        densities_[i]->Merge(targets[i]);
    }
  
    KALDI_LOG << "Merged " << NumPdfs() << " states with target = "
              << target_components << ", power = " << power
              << " and min_count = " << min_count
              << ", merged from " << gauss_at_start << " to "
              << NumGauss();
  }
  
  void AmDiagGmm::Read(std::istream &in_stream, bool binary) {
    int32 num_pdfs, dim;
  
    ExpectToken(in_stream, binary, "<DIMENSION>");
    ReadBasicType(in_stream, binary, &dim);
    ExpectToken(in_stream, binary, "<NUMPDFS>");
    ReadBasicType(in_stream, binary, &num_pdfs);
    KALDI_ASSERT(num_pdfs > 0);
    densities_.reserve(num_pdfs);
    for (int32 i = 0; i < num_pdfs; i++) {
      densities_.push_back(new DiagGmm());
      densities_.back()->Read(in_stream, binary);
      KALDI_ASSERT(densities_.back()->Dim() == dim);
    }
  }
  
  void AmDiagGmm::Write(std::ostream &out_stream, bool binary) const {
    int32 dim = this->Dim();
    if (dim == 0) {
      KALDI_WARN << "Trying to write empty AmDiagGmm object.";
    }
    WriteToken(out_stream, binary, "<DIMENSION>");
    WriteBasicType(out_stream, binary, dim);
    WriteToken(out_stream, binary, "<NUMPDFS>");
    WriteBasicType(out_stream, binary, static_cast<int32>(densities_.size()));
    for (std::vector<DiagGmm*>::const_iterator it = densities_.begin(),
        end = densities_.end(); it != end; ++it) {
      (*it)->Write(out_stream, binary);
    }
  }
  
  void UbmClusteringOptions::Check() {
    if (ubm_num_gauss > intermediate_num_gauss)
      KALDI_ERR << "Invalid parameters: --ubm-num_gauss=" << ubm_num_gauss
                << " > --intermediate-num_gauss=" << intermediate_num_gauss;
    if (ubm_num_gauss > max_am_gauss)
      KALDI_ERR << "Invalid parameters: --ubm-num_gauss=" << ubm_num_gauss
                << " > --max-am-gauss=" << max_am_gauss;
    if (ubm_num_gauss <= 0)
      KALDI_ERR << "Invalid parameters: --ubm-num_gauss=" << ubm_num_gauss;
    if (cluster_varfloor <= 0)
      KALDI_ERR << "Invalid parameters: --cluster-varfloor="
                << cluster_varfloor;
    if (reduce_state_factor <= 0 || reduce_state_factor > 1)
      KALDI_ERR << "Invalid parameters: --reduce-state-factor="
                << reduce_state_factor;
  }
  
  void ClusterGaussiansToUbm(const AmDiagGmm &am,
                             const Vector<BaseFloat> &state_occs,
                             UbmClusteringOptions opts,
                             DiagGmm *ubm_out) {
    opts.Check();  // Make sure the various # of Gaussians make sense.
    if (am.NumGauss() > opts.max_am_gauss) {
      KALDI_LOG << "ClusterGaussiansToUbm: first reducing num-gauss from " << am.NumGauss()
                << " to " << opts.max_am_gauss;
      AmDiagGmm tmp_am;
      tmp_am.CopyFromAmDiagGmm(am);
      BaseFloat power = 1.0, min_count = 1.0; // Make the power 1, which I feel
      // is appropriate to the way we're doing the overall clustering procedure.
      tmp_am.MergeByCount(state_occs, opts.max_am_gauss, power, min_count);
  
      if (tmp_am.NumGauss() > opts.max_am_gauss) {
        KALDI_LOG << "Clustered down to " << tmp_am.NumGauss()
                  << "; will not cluster further";
        opts.max_am_gauss = tmp_am.NumGauss();
      }
      ClusterGaussiansToUbm(tmp_am, state_occs, opts, ubm_out);
      return;
    }
  
    int32 num_pdfs = static_cast<int32>(am.NumPdfs()),
        dim = am.Dim(),
        num_clust_states = static_cast<int32>(opts.reduce_state_factor*num_pdfs);
  
    Vector<BaseFloat> tmp_mean(dim);
    Vector<BaseFloat> tmp_var(dim);
    DiagGmm tmp_gmm;
    vector<Clusterable*> states;
    states.reserve(num_pdfs);  // NOT resize(); uses push_back.
  
    // Replace the GMM for each state with a single Gaussian.
    KALDI_VLOG(1) << "Merging densities to 1 Gaussian per state.";
    for (int32 pdf_index = 0; pdf_index < num_pdfs; pdf_index++) {
      KALDI_VLOG(3) << "Merging Gausians for state : " << pdf_index;
      tmp_gmm.CopyFromDiagGmm(am.GetPdf(pdf_index));
      tmp_gmm.Merge(1);
      tmp_gmm.GetComponentMean(0, &tmp_mean);
      tmp_gmm.GetComponentVariance(0, &tmp_var);
      tmp_var.AddVec2(1.0, tmp_mean);  // make it x^2 stats.
      // It may cause problems downstream if we add states with zero weights (see
      // KALDI_ASSERT(weight > 0) below), so we put in a very small floor.
      // These states with tiny weights will later get merged into other states.
      BaseFloat this_weight = 1.0e-10 + state_occs(pdf_index);
      tmp_mean.Scale(this_weight);
      tmp_var.Scale(this_weight);
      states.push_back(new GaussClusterable(tmp_mean, tmp_var,
                            opts.cluster_varfloor, this_weight));
    }
  
    // Bottom-up clustering of the Gaussians corresponding to each state, which
    // gives a partial clustering of states in the 'state_clusters' vector.
    vector<int32> state_clusters;
    KALDI_VLOG(1) << "Creating " << num_clust_states << " clusters of states.";
    ClusterBottomUp(states, std::numeric_limits<BaseFloat>::max(), num_clust_states,
                    NULL /*actual clusters not needed*/,
                    &state_clusters /*get the cluster assignments*/);
    DeletePointers(&states);
  
    // For each cluster of states, create a pool of all the Gaussians in those
    // states, weighted by the state occupancies. This is done so that initially
    // only the Gaussians corresponding to "similar" states (similarity as
    // determined by the previous clustering) are merged.
    vector< vector<Clusterable*> > state_clust_gauss;
    state_clust_gauss.resize(num_clust_states);
    for (int32 pdf_index = 0; pdf_index < num_pdfs; pdf_index++) {
      int32 current_cluster = state_clusters[pdf_index];
      for (int32 num_gauss = am.GetPdf(pdf_index).NumGauss(),
          gauss_index = 0; gauss_index < num_gauss; ++gauss_index) {
        am.GetGaussianMean(pdf_index, gauss_index, &tmp_mean);
        am.GetGaussianVariance(pdf_index, gauss_index, &tmp_var);
        tmp_var.AddVec2(1.0, tmp_mean);  // make it x^2 stats.
        // adding 1.0e-10 to the weight will prevent problems later on, see
        // the line KALDI_ASSERT(weight > 0.0).
        BaseFloat this_weight =  (1.0e-10 + state_occs(pdf_index)) *
            (am.GetPdf(pdf_index).weights())(gauss_index);
        tmp_mean.Scale(this_weight);
        tmp_var.Scale(this_weight);
        state_clust_gauss[current_cluster].push_back(new GaussClusterable(
            tmp_mean, tmp_var, opts.cluster_varfloor, this_weight));
      }
    }
  
    // This is an unlikely operating scenario, no need to handle this in a more
    // optimized fashion.
    if (opts.intermediate_num_gauss > am.NumGauss()) {
      KALDI_WARN << "Intermediate num_gauss " << opts.intermediate_num_gauss
                 << " is more than num-gauss " << am.NumGauss()
                 << ", reducing it to " << am.NumGauss();
      opts.intermediate_num_gauss = am.NumGauss();
    }
  
    // The compartmentalized clusterer used below does not merge compartments.
    if (opts.intermediate_num_gauss < num_clust_states) {
      KALDI_WARN << "Intermediate num_gauss " << opts.intermediate_num_gauss
                 << " is less than # of preclustered states " << num_clust_states
                 << ", increasing it to " << num_clust_states;
      opts.intermediate_num_gauss = num_clust_states;
    }
  
    KALDI_VLOG(1) << "Merging from " << am.NumGauss() << " Gaussians in the "
                  << "acoustic model, down to " << opts.intermediate_num_gauss
                  << " Gaussians.";
    vector< vector<Clusterable*> > gauss_clusters_out;
    ClusterBottomUpCompartmentalized(state_clust_gauss, std::numeric_limits<BaseFloat>::max(),
                                     opts.intermediate_num_gauss,
                                     &gauss_clusters_out, NULL);
    for (int32 clust_index = 0; clust_index < num_clust_states; clust_index++)
      DeletePointers(&state_clust_gauss[clust_index]);
  
    // Next, put the remaining clustered Gaussians into a single GMM.
    KALDI_VLOG(1) << "Putting " << opts.intermediate_num_gauss << " Gaussians "
                  << "into a single GMM for final merge step.";
    Matrix<BaseFloat> tmp_means(opts.intermediate_num_gauss, dim);
    Matrix<BaseFloat> tmp_vars(opts.intermediate_num_gauss, dim);
    Vector<BaseFloat> tmp_weights(opts.intermediate_num_gauss);
    Vector<BaseFloat> tmp_vec(dim);
    int32 gauss_index = 0;
    for (int32 clust_index = 0; clust_index < num_clust_states; clust_index++) {
      for (int32 i = gauss_clusters_out[clust_index].size()-1; i >=0; --i) {
        GaussClusterable *this_cluster = static_cast<GaussClusterable*>(
            gauss_clusters_out[clust_index][i]);
        BaseFloat weight = this_cluster->count();
        KALDI_ASSERT(weight > 0.0);
        tmp_weights(gauss_index) = weight;
        tmp_vec.CopyFromVec(this_cluster->x_stats());
        tmp_vec.Scale(1.0 / weight);
        tmp_means.CopyRowFromVec(tmp_vec, gauss_index);
        tmp_vec.CopyFromVec(this_cluster->x2_stats());
        tmp_vec.Scale(1.0 / weight);
        tmp_vec.AddVec2(-1.0, tmp_means.Row(gauss_index));  // x^2 stats to var.
        tmp_vars.CopyRowFromVec(tmp_vec, gauss_index);
        gauss_index++;
      }
      DeletePointers(&(gauss_clusters_out[clust_index]));
    }
    tmp_gmm.Resize(opts.intermediate_num_gauss, dim);
    tmp_weights.Scale(1.0/tmp_weights.Sum());
    tmp_gmm.SetWeights(tmp_weights);
    tmp_vars.InvertElements();  // need inverse vars...
    tmp_gmm.SetInvVarsAndMeans(tmp_vars, tmp_means);
  
    // Finally, cluster to the desired number of Gaussians in the UBM.
    if (opts.ubm_num_gauss < tmp_gmm.NumGauss()) {
      tmp_gmm.Merge(opts.ubm_num_gauss);
      KALDI_VLOG(1) << "Merged down to " << tmp_gmm.NumGauss() << " Gaussians.";
    } else {
      KALDI_WARN << "Not merging Gaussians since " << opts.ubm_num_gauss
                 << " < " << tmp_gmm.NumGauss();
    }
    ubm_out->CopyFromDiagGmm(tmp_gmm);
  }
  
  }  // namespace kaldi