// nnet/nnet-matrix-buffer.h

// Copyright 2016  Brno University of Technology (author: Karel Vesely)

// 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_NNET_NNET_MATRIX_BUFFER_H_
#define KALDI_NNET_NNET_MATRIX_BUFFER_H_

#include <utility>
#include <vector>
#include <list>
#include <string>

#include "itf/options-itf.h"
#include "util/common-utils.h"
#include "matrix/kaldi-matrix.h"


namespace kaldi {
namespace nnet1 {

struct MatrixBufferOptions {
  int32 matrix_buffer_size;

  MatrixBufferOptions():
    matrix_buffer_size(3 * 1024)  // 3 x 1GB,
  { }

  void Register(OptionsItf *opts) {
    opts->Register("matrix-buffer-size", &matrix_buffer_size,
       "Capacity of buffer for feature matrices, in MB.");
  }
};


/**
 * A buffer for caching (utterance-key, feature-matrix) pairs.
 * Typically, it reads 'matrix_buffer_size' megabytes of data,
 * and returns records with similar number of speech frames
 * through the standard Key(), Value(), Next(), Done() interface.
 *
 * The preferred length is reset by ResetLength().
 * The buffer gets refilled after having less
 * data than 50% of 'matrix_buffer_size'.
 */
class MatrixBuffer {
 public:
  MatrixBuffer():
    reader_(NULL),
    current_(NULL),
    preferred_length_(0)
  { }

  ~MatrixBuffer()
  { }

  void Init(SequentialBaseFloatMatrixReader* reader,
            MatrixBufferOptions opts = MatrixBufferOptions()) {
    KALDI_ASSERT(SizeInBytes() == 0);
    reader_ = reader;
    opts_ = opts;

    Read();
  }

  bool Done() {
    return (reader_->Done() && NumPairs() <= 1);
  }

  void Next();

  void ResetLength() {
    preferred_length_ = 0;
  }

  std::string Key() {
    return current_->first;
  }
  Matrix<BaseFloat> Value() {
    return current_->second;
  }

  /// Total amount of features in the buffer (bytes),
  size_t SizeInBytes() const;

  /// Total amount of features in the buffer (Mega-bytes),
  size_t SizeInMegaBytes() const;

  /// Total number of (key,matrix) pairs in the buffer,
  size_t NumPairs() const;

 private:

  void Read();  ///< fills the buffer,
  void DisposeValue();  ///< removes 'current_' from data structure,

  SequentialBaseFloatMatrixReader* reader_;

  typedef std::pair<std::string,Matrix<BaseFloat> > PairType;
  typedef std::list<PairType> ListType;
  typedef std::map<size_t, ListType> BufferType;
  BufferType buffer_;  ///< Buffer indexed by 'NumRows()',

  PairType* current_;  ///< The currently active (key,value) pair,

  MatrixBufferOptions opts_;

  size_t preferred_length_;
};

void MatrixBuffer::Next() {
  KALDI_ASSERT(!buffer_.empty());

  // remove old 'Value()' matrix,
  DisposeValue();

  // start re-filling,
  if (SizeInMegaBytes() < 0.5 * opts_.matrix_buffer_size) {
    Read();
  }

  KALDI_ASSERT(!buffer_.empty());

  // randomly select 'length' present in the 'map',
  // (weighted by total #frames in the bin),
  if (preferred_length_ == 0) {
    int32 longest = (--buffer_.end())->first;
    // pre-fill the vector of 'keys',
    std::vector<int32> keys;
    BufferType::iterator it;
    for (it = buffer_.begin(); it != buffer_.end(); ++it) {
      int32 key = it->first; // i.e. NumRows() of matrices in the bin,
      int32 frames_in_bin = it->second.size() * key;
      for (int32 i = 0; i < frames_in_bin; i += longest) {
        keys.push_back(key); // keys are repeated,
      }
    }
    // choose the key,
    std::vector<int32>::iterator it2 = keys.begin();
    std::advance(it2, rand() % keys.size());
    preferred_length_ = (*it2);  // NumRows(), key of the 'map',
  }

  // select list by 'preferred_length_',
  BufferType::iterator it = buffer_.lower_bound(preferred_length_);
  if (it == buffer_.end()) { --it; } // or the last one,

  // take a front element 'ptr' from that list,
  current_ = &(it->second.front());
}

size_t MatrixBuffer::SizeInBytes() const {
  size_t ans = 0;
  for (BufferType::const_iterator it = buffer_.begin(); it != buffer_.end(); ++it) {
    for (ListType::const_iterator it2 = it->second.begin(); it2 != it->second.end(); ++it2) {
      ans += it2->second.SizeInBytes();
    }
  }
  return ans;
}

size_t MatrixBuffer::SizeInMegaBytes() const {
  return (SizeInBytes() / (1024 * 1024));
}

size_t MatrixBuffer::NumPairs() const {
  size_t ans = 0;
  for (BufferType::const_iterator it = buffer_.begin(); it != buffer_.end(); ++it) {
    ans += it->second.size();
  }
  return ans;
}

void MatrixBuffer::Read() {
  if (!reader_->Done())
    KALDI_LOG << "Read() started... Buffer size in MB: "
              << SizeInMegaBytes() << ", max " << opts_.matrix_buffer_size
              << ", having " << NumPairs() << " utterances.";
  for ( ; !reader_->Done(); reader_->Next()) {
    // see if we are full,
    if (SizeInMegaBytes() > opts_.matrix_buffer_size) {
      KALDI_LOG << "Read() finished... Buffer size in MB: "
                << SizeInMegaBytes() << ", max " << opts_.matrix_buffer_size
                << ", having " << NumPairs() << " utterances.";
      break;
    }
    // get matrix,
    const std::string& key = reader_->Key();
    const Matrix<BaseFloat>& mat = reader_->Value();
    size_t num_rows = mat.NumRows();
    // see if 'num_rows' already in keys,
    if (buffer_.find(num_rows) == buffer_.end()) {
      buffer_[num_rows] = ListType();  // add empty list,
    }
    // add matrix to the buffer,
    buffer_[num_rows].push_back(PairType(key, mat));
  }
}

void MatrixBuffer::DisposeValue() {
  // remove old 'Value()' matrix,
  if (current_ != NULL) {
    size_t r = current_->second.NumRows();
    KALDI_ASSERT(current_ == &(buffer_[r].front()));
    // remove the (key,value) pair,
    buffer_[r].pop_front();
    // eventually remove the 'NumRows()' key,
    if (buffer_[r].empty()) { buffer_.erase(r); }
    current_ = NULL;
  }
}


}  // namespace nnet1
}  // namespace kaldi

#endif  // KALDI_NNET_NNET_MATRIX_BUFFER_H_