// nnet3bin/nnet3-acc-lda-stats.cc
  
  // Copyright 2015  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 "base/kaldi-common.h"
  #include "util/common-utils.h"
  #include "hmm/transition-model.h"
  #include "nnet3/nnet-nnet.h"
  #include "nnet3/nnet-example-utils.h"
  #include "nnet3/nnet-optimize.h"
  #include "transform/lda-estimate.h"
  
  
  namespace kaldi {
  namespace nnet3 {
  
  class NnetLdaStatsAccumulator {
   public:
    NnetLdaStatsAccumulator(BaseFloat rand_prune,
                            const Nnet &nnet):
        rand_prune_(rand_prune), nnet_(nnet), compiler_(nnet) { }
  
    void AccStats(const NnetExample &eg) {
      ComputationRequest request;
      bool need_backprop = false, store_stats = false;
      GetComputationRequest(nnet_, eg, need_backprop, store_stats, &request);
      const NnetComputation &computation = *(compiler_.Compile(request));
      NnetComputeOptions options;
      if (GetVerboseLevel() >= 3)
        options.debug = true;
      NnetComputer computer(options, computation, nnet_, NULL);
  
      computer.AcceptInputs(nnet_, eg.io);
      computer.Run();
      const CuMatrixBase<BaseFloat> &nnet_output = computer.GetOutput("output");
      AccStatsFromOutput(eg, nnet_output);
    }
  
    void WriteStats(const std::string &stats_wxfilename, bool binary) {
      if (lda_stats_.TotCount() == 0) {
        KALDI_ERR << "Accumulated no stats.";
      } else {
        WriteKaldiObject(lda_stats_, stats_wxfilename, binary);
        KALDI_LOG << "Accumulated stats, soft frame count = "
                  << lda_stats_.TotCount() << ".  Wrote to "
                  << stats_wxfilename;
      }
    }
   private:
    void AccStatsFromOutput(const NnetExample &eg,
                            const CuMatrixBase<BaseFloat> &nnet_output) {
      BaseFloat rand_prune = rand_prune_;
      const NnetIo *output_supervision = NULL;
      for (size_t i = 0; i < eg.io.size(); i++)
        if (eg.io[i].name == "output")
          output_supervision = &(eg.io[i]);
      KALDI_ASSERT(output_supervision != NULL && "no output in eg named 'output'");
      int32 num_rows = output_supervision->features.NumRows(),
          num_pdfs = output_supervision->features.NumCols();
      KALDI_ASSERT(num_rows == nnet_output.NumRows());
      if (lda_stats_.Dim() == 0)
        lda_stats_.Init(num_pdfs, nnet_output.NumCols());
      if (output_supervision->features.Type() == kSparseMatrix) {
        const SparseMatrix<BaseFloat> &smat =
            output_supervision->features.GetSparseMatrix();
        for (int32 r = 0; r < num_rows; r++) {
          // the following, transferring row by row to CPU, would be wasteful
          // if we actually were using a GPU, but we don't anticipate doing this
          // in this program.
          CuSubVector<BaseFloat> cu_row(nnet_output, r);
          // "row" is actually just a redudant copy, since we're likely on CPU,
          // but we're about to do an outer product, so this doesn't dominate.
          Vector<BaseFloat> row(cu_row);
  
          const SparseVector<BaseFloat> &post(smat.Row(r));
          const std::pair<MatrixIndexT, BaseFloat> *post_data = post.Data(),
              *post_end = post_data + post.NumElements();
          for (; post_data != post_end; ++post_data) {
            MatrixIndexT pdf = post_data->first;
            BaseFloat weight = post_data->second;
            BaseFloat pruned_weight = RandPrune(weight, rand_prune);
            if (pruned_weight != 0.0)
              lda_stats_.Accumulate(row, pdf, pruned_weight);
          }
        }
      } else {
        Matrix<BaseFloat> output_mat;
        output_supervision->features.GetMatrix(&output_mat);
        for (int32 r = 0; r < num_rows; r++) {
          // the following, transferring row by row to CPU, would be wasteful
          // if we actually were using a GPU, but we don't anticipate doing this
          // in this program.
          CuSubVector<BaseFloat> cu_row(nnet_output, r);
          // "row" is actually just a redudant copy, since we're likely on CPU,
          // but we're about to do an outer product, so this doesn't dominate.
          Vector<BaseFloat> row(cu_row);
  
          SubVector<BaseFloat> post(output_mat, r);
          int32 num_pdfs = post.Dim();
          for (int32 pdf = 0; pdf < num_pdfs; pdf++) {
            BaseFloat weight = post(pdf);
            BaseFloat pruned_weight = RandPrune(weight, rand_prune);
            if (pruned_weight != 0.0)
              lda_stats_.Accumulate(row, pdf, pruned_weight);
          }
        }
      }
    }
  
    BaseFloat rand_prune_;
    const Nnet &nnet_;
    CachingOptimizingCompiler compiler_;
    LdaEstimate lda_stats_;
  
  };
  
  }
  }
  
  int main(int argc, char *argv[]) {
    try {
      using namespace kaldi;
      using namespace kaldi::nnet3;
      typedef kaldi::int32 int32;
      typedef kaldi::int64 int64;
  
      const char *usage =
          "Accumulate statistics in the same format as acc-lda (i.e. stats for
  "
          "estimation of LDA and similar types of transform), starting from nnet
  "
          "training examples.  This program puts the features through the network,
  "
          "and the network output will be the features; the supervision in the
  "
          "training examples is used for the class labels.  Used in obtaining
  "
          "feature transforms that help nnet training work better.
  "
          "
  "
          "Usage:  nnet3-acc-lda-stats [options] <raw-nnet-in> <training-examples-in> <lda-stats-out>
  "
          "e.g.:
  "
          "nnet3-acc-lda-stats 0.raw ark:1.egs 1.acc
  "
          "See also: nnet-get-feature-transform
  ";
  
      bool binary_write = true;
      BaseFloat rand_prune = 0.0;
  
      ParseOptions po(usage);
      po.Register("binary", &binary_write, "Write output in binary mode");
      po.Register("rand-prune", &rand_prune,
                  "Randomized pruning threshold for posteriors");
  
      po.Read(argc, argv);
  
      if (po.NumArgs() != 3) {
        po.PrintUsage();
        exit(1);
      }
  
      std::string nnet_rxfilename = po.GetArg(1),
          examples_rspecifier = po.GetArg(2),
          lda_accs_wxfilename = po.GetArg(3);
  
      Nnet nnet;
      ReadKaldiObject(nnet_rxfilename, &nnet);
  
      NnetLdaStatsAccumulator accumulator(rand_prune, nnet);
  
      int64 num_egs = 0;
  
      SequentialNnetExampleReader example_reader(examples_rspecifier);
      for (; !example_reader.Done(); example_reader.Next(), num_egs++)
        accumulator.AccStats(example_reader.Value());
  
      KALDI_LOG << "Processed " << num_egs << " examples.";
      // the next command will die if we accumulated no stats.
      accumulator.WriteStats(lda_accs_wxfilename, binary_write);
  
      return 0;
    } catch(const std::exception &e) {
      std::cerr << e.what() << '
  ';
      return -1;
    }
  }