// nnet3/nnet-example-utils.h
  
  // 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.
  
  #ifndef KALDI_NNET3_NNET_EXAMPLE_UTILS_H_
  #define KALDI_NNET3_NNET_EXAMPLE_UTILS_H_
  
  #include "nnet3/nnet-example.h"
  #include "nnet3/nnet-computation.h"
  #include "nnet3/nnet-compute.h"
  #include "util/kaldi-table.h"
  
  namespace kaldi {
  namespace nnet3 {
  
  
  
  /** Merge a set of input examples into a single example (typically the size of
      "src" will be the minibatch size).  Will crash if "src" is the empty vector.
      If "compress" is true, it will compress any non-sparse features in the output.
   */
  void MergeExamples(const std::vector<NnetExample> &src,
                     bool compress,
                     NnetExample *dest);
  
  
  /** Shifts the time-index t of everything in the "eg" by adding "t_offset" to
      all "t" values.  This might be useful in things like clockwork RNNs that are
      not invariant to time-shifts, to ensure that we see different shifts of each
      example during training.  "exclude_names" is a vector (not necessarily
      sorted) of names of nnet inputs that we avoid shifting the "t" values of--
      normally it will contain just the single string "ivector" because we always
      leave t=0 for any ivector. */
  void ShiftExampleTimes(int32 t_offset,
                         const std::vector<std::string> &exclude_names,
                         NnetExample *eg);
  
  /**  This function takes a NnetExample (which should already have been
       frame-selected, if desired, and merged into a minibatch) and produces a
       ComputationRequest.  It assumes you don't want the derivatives w.r.t. the
       inputs; if you do, you can create/modify the ComputationRequest manually.
       Assumes that if need_model_derivative is true, you will be supplying
       derivatives w.r.t. all outputs.
  */
  void GetComputationRequest(const Nnet &nnet,
                             const NnetExample &eg,
                             bool need_model_derivative,
                             bool store_component_stats,
                             ComputationRequest *computation_request);
  
  
  // Writes as unsigned char a vector 'vec' that is required to have
  // values between 0 and 1.
  void WriteVectorAsChar(std::ostream &os,
                         bool binary,
                         const VectorBase<BaseFloat> &vec);
  
  // Reads data written by WriteVectorAsChar.
  void ReadVectorAsChar(std::istream &is,
                        bool binary,
                        Vector<BaseFloat> *vec);
  
  
  // Warning: after reading in the values from the command line
  // (Register() and then then po.Read()), you should then call ComputeDerived()
  // to set up the 'derived values' (parses 'num_frames_str').
  struct ExampleGenerationConfig {
    int32 left_context;
    int32 right_context;
    int32 left_context_initial;
    int32 right_context_final;
    int32 num_frames_overlap;
    int32 frame_subsampling_factor;
    std::string num_frames_str;
  
  
    // The following parameters are derived parameters, computed by
    // ComputeDerived().
  
    // the first element of the 'num_frames' vector is the 'principal' number of
    // frames; the remaining elements are alternatives to the principal number of
    // frames, to be used at most once or twice per file.
    std::vector<int32> num_frames;
  
    ExampleGenerationConfig():
        left_context(0), right_context(0),
        left_context_initial(-1), right_context_final(-1),
        num_frames_overlap(0), frame_subsampling_factor(1),
        num_frames_str("1") { }
  
    /// This function decodes 'num_frames_str' into 'num_frames', and ensures that
    /// the members of 'num_frames' are multiples of 'frame_subsampling_factor'.
    void ComputeDerived();
  
    void Register(OptionsItf *po) {
      po->Register("left-context", &left_context, "Number of frames of left "
                   "context of input features that are added to each "
                   "example");
      po->Register("right-context", &right_context, "Number of frames of right "
                   "context of input features that are added to each "
                   "example");
      po->Register("left-context-initial", &left_context_initial, "Number of "
                   "frames of left context of input features that are added to "
                   "each example at the start of the utterance (if <0, this "
                   "defaults to the same as --left-context)");
      po->Register("right-context-final", &right_context_final, "Number of "
                   "frames of right context of input features that are added "
                   "to each example at the end of the utterance (if <0, this "
                   "defaults to the same as --right-context)");
      po->Register("num-frames", &num_frames_str, "Number of frames with labels "
                  "that each example contains (i.e. the left and right context "
                  "are to be added to this).  May just be an integer (e.g. "
                  "--num-frames=8), or a principal value followed by "
                  "alternative values to be used at most once for each utterance "
                  "to deal with odd-sized input, e.g. --num-frames=40,25,50 means "
                  "that most of the time the number of frames will be 40, but to "
                  "deal with odd-sized inputs we may also generate egs with these "
                  "other sizes.  All these values will be rounded up to the "
                  "closest multiple of --frame-subsampling-factor.  As a special case, "
                  "--num-frames=-1 means 'don't do any splitting'.");
      po->Register("num-frames-overlap", &num_frames_overlap, "Number of frames of "
                   "overlap between adjacent eamples (applies to chunks of size "
                   "equal to the primary [first-listed] --num-frames value... "
                   "will be adjusted for different-sized chunks).  Advisory; "
                   "will not be exactly enforced.");
      po->Register("frame-subsampling-factor", &frame_subsampling_factor, "Used "
                   "if the frame-rate of the output labels in the generated "
                   "examples will be less than the frame-rate at the input");
    }
  };
  
  
  
  /**
     struct ChunkTimeInfo is used by class UtteranceSplitter to output
     information about how we split an utterance into chunks.
   */
  struct ChunkTimeInfo {
    int32 first_frame;
    int32 num_frames;
    int32 left_context;
    int32 right_context;
    // The 'output_weights' member is a vector of length equal to the
    // num_frames divided by frame_subsampling_factor from the config.
    // It contains values 0 < x <= 1 that represent weightings of
    // output-frames.  The idea is that if (because of overlaps) a
    // frame appears in multiple chunks, we want to downweight it
    // so that the total weight remains 1.  (Of course, the calling
    // code is free to ignore these weights if desired).
    std::vector<BaseFloat> output_weights;
  };
  
  
  class UtteranceSplitter {
   public:
  
    UtteranceSplitter(const ExampleGenerationConfig &config);
  
  
    const ExampleGenerationConfig& Config() const { return config_; }
  
    // Given an utterance length, this function creates for you a list of chunks
    // into which to split the utterance.  Note: this is partly random (will call
    // srand()).
    // Accumulates some stats which will be printed out in the destructor.
    void GetChunksForUtterance(int32 utterance_length,
                               std::vector<ChunkTimeInfo> *chunk_info);
  
  
    // This function returns true if 'supervision_length' (e.g. the length of the
    // posterior, lattice or alignment) is what we expect given
    // config_.frame_subsampling_factor.  If not, it prints a warning (which is
    // why the function needs 'utt', and returns false.  Note: we round up, so
    // writing config_.frame_subsampling_factor as sf, we expect
    // supervision_length = (utterance_length + sf - 1) / sf.
    bool LengthsMatch(const std::string &utt,
                      int32 utterance_length,
                      int32 supervision_length,
                      int32 length_tolerance = 0) const;
  
    ~UtteranceSplitter();
  
    int32 ExitStatus() { return (total_frames_in_chunks_ > 0 ? 0 : 1); }
  
   private:
  
  
    void InitSplitForLength();
  
    // This function returns the 'default duration' in frames of a split, which if
    // config_.num_frames_overlap is zero is just the sum of chunk sizes in the
    // split (i.e. the sum of the vector's elements), but otherwise, we subtract
    // the recommended overlap (see code for details).
    float DefaultDurationOfSplit(const std::vector<int32> &split) const;
  
  
    // Used in InitSplitForLength(), returns the maximum utterance-length considered
    // separately in split_for_length_.  [above this, we'll assume that the additional
    // length is consumed by multiples of the 'principal' chunk size.]  It returns
    // the primary chunk-size (config_.num_frames[0]) plus twice the largest of
    // any of the allowed chunk sizes (i.e. the max of config_.num_frames)
    int32 MaxUtteranceLength() const;
  
    // Used in InitSplitForLength(), this function outputs the set of allowed
    // splits, represented as a sorted list of nonempty vectors (each split is a
    // sorted list of chunk-sizes).
    void InitSplits(std::vector<std::vector<int32> > *splits) const;
  
  
    // Used in GetChunksForUtterance, this function selects the list of
    // chunk-sizes for that utterance (later on, the positions and and left/right
    // context information for the chunks will be added to this).  We don't call
    // this a 'split', although it's also a list of chunk-sizes, because we
    // randomize the order in which the chunk sizes appear, whereas for a 'split'
    // we sort the chunk-sizes because a 'split' is conceptually an
    // order-independent representation.
    void GetChunkSizesForUtterance(int32 utterance_length,
                                   std::vector<int32> *chunk_sizes) const;
  
  
    // Used in GetChunksForUtterance, this function selects the 'gap sizes'
    // before each of the chunks.  These 'gap sizes' may be positive (representing
    // a gap between chunks, or a number of frames at the beginning of the file that
    // don't correspond to a chunk), or may be negative, corresponding to overlaps
    // between adjacent chunks.
    //
    // If config_.frame_subsampling_factor > 1 and enforce_subsampling_factor is
    // true, this function will ensure that all elements of 'gap_sizes' are
    // multiples of config_.frame_subsampling_factor.  (we always enforce this,
    // but we set it to false inside a recursion when we recurse).  Note: if
    // config_.frame_subsampling_factor > 1, it's possible for the last chunk to
    // go over 'utterance_length' by up to config_.frame_subsampling_factor - 1
    // frames (i.e. it would require that many frames past the utterance end).
    // This will be dealt with when generating egs, by duplicating the last frame.
    void GetGapSizes(int32 utterance_length,
                     bool enforce_subsampling_factor,
                     const std::vector<int32> &chunk_sizes,
                     std::vector<int32> *gap_sizes) const;
  
    // this static function, used in GetGapSizes(), writes random values to a
    // vector 'vec' such the sum of those values equals n (n may be positive or
    // negative).  It tries to make those values as similar as possible (they will
    // differ by at most one), and the location of the larger versus smaller
    // values is random. 'vec' must be nonempty.
    static void DistributeRandomlyUniform(int32 n,
                                          std::vector<int32> *vec);
  
    // this static function, used in GetGapSizes(), writes values to a vector
    // 'vec' such the sum of those values equals n (n may be positive or
    // negative).  It tries to make those values, as exactly as it can,
    // proportional to the values in 'magnitudes', which must be positive.  'vec'
    // must be nonempty, and 'magnitudes' must be the same size as 'vec'.
    static void DistributeRandomly(int32 n,
                                   const std::vector<int32> &magnitudes,
                                   std::vector<int32> *vec);
  
    // This function is responsible for setting the 'output_weights'
    // members of the chunks.
    void SetOutputWeights(int32 utterance_length,
                          std::vector<ChunkTimeInfo> *chunk_info) const;
  
    // Accumulate stats for diagnostics.
    void AccStatsForUtterance(int32 utterance_length,
                              const std::vector<ChunkTimeInfo> &chunk_info);
  
  
    const ExampleGenerationConfig &config_;
  
    // The vector 'splits_for_length_' is indexed by the num-frames of a file, and
    // gives us a list of alternative splits that we can use if the utternace has
    // that many frames.  For example, if split_for_length[100] = ( (25, 40, 40),
    // (40, 65) ), it means we could either split as chunks of size (25, 40, 40)
    // or as (40, 65).  (we'll later randomize the order).  should use one chunk
    // of size 25 and two chunks of size 40.  In general these won't add up to
    // exactly the length of the utterance; we'll have them overlap (or have small
    // gaps between them) to account for this, and the details of this will be
    // randomly decided per file.  If splits_for_length_[u] is empty, it means the
    // utterance was shorter than the smallest possible chunk size, so
    // we will have to discard the utterance.
  
    // If an utterance's num-frames is >= split_for_length.size(), the way to find
    // the split to use is to keep subtracting the primary num-frames (==
    // config_.num_frames[0]) minus the num-frames-overlap, from the utterance
    // length, until the resulting num-frames is < split_for_length_.size(),
    // chunks, and then add the subtracted number of copies of the primary
    // num-frames to the split.
    std::vector<std::vector<std::vector<int32> > > splits_for_length_;
  
    // Below are stats used for diagnostics.
    int32 total_num_utterances_;  // total input utterances.
    int64 total_input_frames_;  // total num-frames over all utterances (before
                                // splitting)
    int64 total_frames_overlap_;  // total number of frames that overlap between
                                  // adjacent egs.
    int64 total_num_chunks_;
    int64 total_frames_in_chunks_;  // total of chunk-size times count of that
                                    // chunk.  equals the num-frames in all the
                                    // output chunks, added up.
    std::map<int32, int32> chunk_size_to_count_;  // for each chunk size, gives
                                                  // the number of chunks with
                                                  // that size.
  
  };
  
  
  class ExampleMergingConfig {
  public:
    // The following configuration values are registered on the command line.
    bool compress;
    std::string measure_output_frames;  // for back-compatibility, not used.
    std::string minibatch_size;
    std::string discard_partial_minibatches;   // for back-compatibility, not used.
  
    ExampleMergingConfig(const char *default_minibatch_size = "256"):
        compress(false),
        measure_output_frames("deprecated"),
        minibatch_size(default_minibatch_size),
        discard_partial_minibatches("deprecated") { }
  
    void Register(OptionsItf *po) {
      po->Register("compress", &compress, "If true, compress the output examples "
                   "(not recommended unless you are writing to disk)");
      po->Register("measure-output-frames", &measure_output_frames, "This "
                   "value will be ignored (included for back-compatibility)");
      po->Register("discard-partial-minibatches", &discard_partial_minibatches,
                   "This value will be ignored (included for back-compatibility)");
      po->Register("minibatch-size", &minibatch_size,
                   "String controlling the minibatch size.  May be just an integer, "
                   "meaning a fixed minibatch size (e.g. --minibatch-size=128). "
                   "May be a list of ranges and values, e.g. --minibatch-size=32,64 "
                   "or --minibatch-size=16:32,64,128.  All minibatches will be of "
                   "the largest size until the end of the input is reached; "
                   "then, increasingly smaller sizes will be allowed.  Only egs "
                   "with the same structure (e.g num-frames) are merged.  You may "
                   "specify different minibatch sizes for different sizes of eg "
                   "(defined as the maximum number of Indexes on any input), in "
                   "the format "
                   "--minibatch-size='eg_size1=mb_sizes1/eg_size2=mb_sizes2', e.g. "
                   "--minibatch-size=128=64:128,256/256=32:64,128.  Egs are given "
                   "minibatch-sizes based on the specified eg-size closest to "
                   "their actual size.");
    }
  
  
    // this function computes the derived (private) parameters; it must be called
    // after the command-line parameters are read and before MinibatchSize() is
    // called.
    void ComputeDerived();
  
    /// This function tells you what minibatch size should be used for this eg.
  
    ///  @param [in] size_of_eg   The "size" of the eg, as obtained by
    ///                           GetNnetExampleSize() or a similar function (up
    ///                           to the caller).
    ///  @param [in] num_available_egs   The number of egs of this size that are
    ///                            currently available; should be >0.  The
    ///                            value returned will be <= this value, possibly
    ///                            zero.
    ///  @param [in] input_ended   True if the input has ended, false otherwise.
    ///                            This is important because before the input has
    ///                            ended, we will only batch egs into the largest
    ///                            possible minibatch size among the range allowed
    ///                            for that size of eg.
    ///  @return                   Returns the minibatch size to use in this
    ///                            situation, as specified by the configuration.
    int32 MinibatchSize(int32 size_of_eg,
                        int32 num_available_egs,
                        bool input_ended) const;
  
  
   private:
    // struct IntSet is a representation of something like 16:32,64, which is a
    // nonempty list of either positive integers or ranges of positive integers.
    // Conceptually it represents a set of positive integers.
    struct IntSet {
      // largest_size is the largest integer in any of the ranges (64 in this
      // example).
      int32 largest_size;
      // e.g. would contain ((16,32), (64,64)) in this example.
      std::vector<std::pair<int32, int32> > ranges;
      // Returns the largest value in any range (i.e. in the set of
      // integers that this struct represents), that is <= max_value,
      // or 0 if there is no value in any range that is <= max_value.
      // In this example, this function would return the following:
      // 128->64, 64->64, 63->32, 31->31, 16->16, 15->0, 0->0
      int32 LargestValueInRange(int32 max_value) const;
    };
    static bool ParseIntSet(const std::string &str, IntSet *int_set);
  
    // 'rules' is derived from the configuration values above by ComputeDerived(),
    // and are not set directly on the command line.  'rules' is a list of pairs
    // (eg-size, int-set-of-minibatch-sizes); If no explicit eg-sizes were
    // specified on the command line (i.e. there was no '=' sign in the
    // --minibatch-size option), then we just set the int32 to 0.
    std::vector<std::pair<int32, IntSet> > rules;
  };
  
  
  /// This function returns the 'size' of a nnet-example as defined for purposes
  /// of merging egs, which is defined as the largest number of Indexes in any of
  /// the inputs or outputs of the example.
  int32 GetNnetExampleSize(const NnetExample &a);
  
  
  
  
  
  /// This class is responsible for storing, and displaying in log messages,
  /// statistics about how examples of different sizes (c.f. GetNnetExampleSize())
  /// were merged into minibatches, and how many examples were left over and
  /// discarded.
  class ExampleMergingStats {
   public:
    /// Users call this function to inform this class that one minibatch has been
    /// written aggregating 'minibatch_size' separate examples of original size
    /// 'example_size' (e.g. as determined by GetNnetExampleSize(), but the caller
    /// does that.
    /// The 'structure_hash' is provided so that this class can distinguish
    /// between egs that have the same size but different structure.  In the
    /// extremely unlikely eventuality that there is a hash collision, it will
    /// cause misleading stats to be printed out.
    void WroteExample(int32 example_size, size_t structure_hash,
                      int32 minibatch_size);
  
    /// Users call this function to inform this class that after processing all
    /// the data, for examples of original size 'example_size', 'num_discarded'
    /// examples could not be put into a minibatch and were discarded.
    void DiscardedExamples(int32 example_size, size_t structure_hash,
                           int32 num_discarded);
  
    /// Calling this will cause a log message with information about the
    /// examples to be printed.
    void PrintStats() const;
  
   private:
    // this struct stores the stats for examples of a particular size and
    // structure.
    struct StatsForExampleSize {
      int32 num_discarded;
      // maps from minibatch-size (i.e. number of egs that were
      // aggregated into that minibatch), to the number of such
      // minibatches written.
      unordered_map<int32, int32> minibatch_to_num_written;
      StatsForExampleSize(): num_discarded(0) { }
    };
  
  
    typedef unordered_map<std::pair<int32, size_t>, StatsForExampleSize,
                          PairHasher<int32, size_t> > StatsType;
  
    // this maps from a pair (example_size, structure_hash) to to the stats for
    // examples with those characteristics.
    StatsType stats_;
  
    void PrintAggregateStats() const;
    void PrintSpecificStats() const;
  
  };
  
  
  /// This class is responsible for arranging examples in groups
  /// that have the same strucure (i.e. the same input and output
  /// indexes), and outputting them in suitable minibatches
  /// as defined by ExampleMergingConfig.
  class ExampleMerger {
   public:
    ExampleMerger(const ExampleMergingConfig &config,
                  NnetExampleWriter *writer);
  
    // This function accepts an example, and if possible, writes a merged example
    // out.  The ownership of the pointer 'a' is transferred to this class when
    // you call this function.
    void AcceptExample(NnetExample *a);
  
    // This function announces to the class that the input has finished, so it
    // should flush out any smaller-sized minibatches, as dictated by the config.
    // This will be called in the destructor, but you can call it explicitly when
    // all the input is done if you want to; it won't repeat anything if called
    // twice.  It also prints the stats.
    void Finish();
  
    // returns a suitable exit status for a program.
    int32 ExitStatus() { Finish(); return (num_egs_written_ > 0 ? 0 : 1); }
  
    ~ExampleMerger() { Finish(); };
   private:
    // called by Finish() and AcceptExample().  Merges, updates the
    // stats, and writes.
    void WriteMinibatch(const std::vector<NnetExample> &egs);
  
    bool finished_;
    int32 num_egs_written_;
    const ExampleMergingConfig &config_;
    NnetExampleWriter *writer_;
    ExampleMergingStats stats_;
  
    // Note: the "key" into the egs is the first element of the vector.
    typedef unordered_map<NnetExample*, std::vector<NnetExample*>,
                          NnetExampleStructureHasher,
                          NnetExampleStructureCompare> MapType;
     MapType eg_to_egs_;
  };
  
  } // namespace nnet3
  } // namespace kaldi
  
  #endif // KALDI_NNET3_NNET_EXAMPLE_UTILS_H_