// ivectorbin/ivector-mean.cc
  
  // Copyright 2013-2014  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"
  
  
  int main(int argc, char *argv[]) {
    using namespace kaldi;
    typedef kaldi::int32 int32;
    try {
      const char *usage =
          "With 3 or 4 arguments, averages iVectors over all the
  "
          "utterances of each speaker using the spk2utt file.
  "
          "Input the spk2utt file and a set of iVectors indexed by
  "
          "utterance; output is iVectors indexed by speaker.  If 4
  "
          "arguments are given, extra argument is a table for the number
  "
          "of utterances per speaker (can be useful for PLDA).  If 2
  "
          "arguments are given, computes the mean of all input files and
  "
          "writes out the mean vector.
  "
          "
  "
          "Usage: ivector-mean <spk2utt-rspecifier> <ivector-rspecifier> "
          "<ivector-wspecifier> [<num-utt-wspecifier>]
  "
          "or: ivector-mean <ivector-rspecifier> <mean-wxfilename>
  "
          "e.g.: ivector-mean data/spk2utt exp/ivectors.ark exp/spk_ivectors.ark exp/spk_num_utts.ark
  "
          "or: ivector-mean exp/ivectors.ark exp/mean.vec
  "
          "See also: ivector-subtract-global-mean
  ";
  
      ParseOptions po(usage);
      bool binary_write = false;
      po.Register("binary", &binary_write, "If true, write output in binary "
                  "(only applicable when writing files, not archives/tables.");
  
      po.Read(argc, argv);
  
      if (po.NumArgs() < 2 || po.NumArgs() > 4) {
        po.PrintUsage();
        exit(1);
      }
  
      if (po.NumArgs() == 2) {
        // Compute the mean of the input vectors and write it out.
        std::string ivector_rspecifier = po.GetArg(1),
            mean_wxfilename = po.GetArg(2);
        int32 num_done = 0;
        SequentialBaseFloatVectorReader ivector_reader(ivector_rspecifier);
        Vector<double> sum;
        for (; !ivector_reader.Done(); ivector_reader.Next()) {
          if (sum.Dim() == 0) sum.Resize(ivector_reader.Value().Dim());
          sum.AddVec(1.0, ivector_reader.Value());
          num_done++;
        }
        if (num_done == 0) {
          KALDI_ERR << "No iVectors read";
        } else {
          sum.Scale(1.0 / num_done);
          WriteKaldiObject(sum, mean_wxfilename, binary_write);
          return 0;
        }
      } else {
        std::string spk2utt_rspecifier = po.GetArg(1),
            ivector_rspecifier = po.GetArg(2),
            ivector_wspecifier = po.GetArg(3),
            num_utts_wspecifier = po.GetOptArg(4);
  
        double spk_sumsq = 0.0;
        Vector<double> spk_sum;
  
        int64 num_spk_done = 0, num_spk_err = 0,
            num_utt_done = 0, num_utt_err = 0;
  
        RandomAccessBaseFloatVectorReader ivector_reader(ivector_rspecifier);
        SequentialTokenVectorReader spk2utt_reader(spk2utt_rspecifier);
        BaseFloatVectorWriter ivector_writer(ivector_wspecifier);
        Int32Writer num_utts_writer(num_utts_wspecifier);
  
        for (; !spk2utt_reader.Done(); spk2utt_reader.Next()) {
          std::string spk = spk2utt_reader.Key();
          const std::vector<std::string> &uttlist = spk2utt_reader.Value();
          if (uttlist.empty()) {
            KALDI_ERR << "Speaker with no utterances.";
          }
          Vector<BaseFloat> spk_mean;
          int32 utt_count = 0;
          for (size_t i = 0; i < uttlist.size(); i++) {
            std::string utt = uttlist[i];
            if (!ivector_reader.HasKey(utt)) {
              KALDI_WARN << "No iVector present in input for utterance " << utt;
              num_utt_err++;
            } else {
              if (utt_count == 0) {
                spk_mean = ivector_reader.Value(utt);
              } else {
                spk_mean.AddVec(1.0, ivector_reader.Value(utt));
              }
              num_utt_done++;
              utt_count++;
            }
          }
          if (utt_count == 0) {
            KALDI_WARN << "Not producing output for speaker " << spk
                       << " since no utterances had iVectors";
            num_spk_err++;
          } else {
            spk_mean.Scale(1.0 / utt_count);
            ivector_writer.Write(spk, spk_mean);
            if (num_utts_wspecifier != "")
              num_utts_writer.Write(spk, utt_count);
            num_spk_done++;
            spk_sumsq += VecVec(spk_mean, spk_mean);
            if (spk_sum.Dim() == 0)
              spk_sum.Resize(spk_mean.Dim());
            spk_sum.AddVec(1.0, spk_mean);
          }
        }
  
        KALDI_LOG << "Computed mean of " << num_spk_done << " speakers ("
                  << num_spk_err << " with no utterances), consisting of "
                  << num_utt_done << " utterances (" << num_utt_err
                  << " absent from input).";
  
        if (num_spk_done != 0) {
          spk_sumsq /= num_spk_done;
          spk_sum.Scale(1.0 / num_spk_done);
          double mean_length = spk_sum.Norm(2.0),
              spk_length = sqrt(spk_sumsq),
              norm_spk_length = spk_length / sqrt(spk_sum.Dim());
          KALDI_LOG << "Norm of mean of speakers is " << mean_length
                    << ", root-mean-square speaker-iVector length divided by "
                    << "sqrt(dim) is " << norm_spk_length;
        }
  
        return (num_spk_done != 0 ? 0 : 1);
      }
    } catch(const std::exception &e) {
      std::cerr << e.what();
      return -1;
    }
  }