// Copyright 2019 The MediaPipe 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
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef MEDIAPIPE_UTIL_TRACKING_STREAMING_BUFFER_H_
#define MEDIAPIPE_UTIL_TRACKING_STREAMING_BUFFER_H_

#include <deque>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

#include "absl/container/node_hash_map.h"
#include "absl/types/any.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/tool/type_util.h"

namespace mediapipe {

// Streaming buffer to store arbitrary data over a chunk of frames with overlap
// between chunks.
// Useful to compute solutions that require as input buffered inputs I_ij, of
// type T_j for all frames i of a chunk. Output solutions S_ik of type T_k
// for each frame i (T_k is not associated with input types T_i in any way)
// can than be stored in the buffer as well and made available to the next
// chunk.
// After solution S_ik has been computed, buffered results (I_ij, S_ik) can be
// output and buffer is truncated to discard all elements minus the overlap.
// Remaining elements (I_ij, S_ik) form the basis for the next chunk.
// Again, arbitrary input data I_ij is buffered until buffer is full.
// However, previous solutions S_ik from the overlap still remain in the buffer
// and can be used to compute new solutions in a temporally constrained manner.

// Detailed usage example:
// Setup streaming buffer, that will buffer cv::Mat's and AffineModel's over a
// set of 100 frames. Some algorithm will then compute for each input pair the
// foreground saliency in the form of a SaliencyPointList proto.
//
// // Prepare streaming buffer with 2 inputs and one solution, i.e. 3 different
// // data types.
// vector<TaggedType> data_config = {
//     TaggedPointerType<cv::Mat>("frame"),
//     TaggedPointerType<AffineModel>("motion"),
//     TaggedPointerType<SaliencyPointList>("saliency") };
//
// // Create new buffer with overlap of 10 frames between chunks.
// StreamingBuffer streaming_buffer(data_config, 10);
//
// // Buffer inputs.
// for (int k = 0; k < 100; ++k) {
//   std::unique_ptr<cv::Mat> input_frame = ...
//   std::unique_ptr<AffineModel> affine_model = ...
//   streaming_buffer.AddDatum("frame", std::move(input_frame));
//   streaming_buffer.AddDatum("motion", std::move(affine_model);
//   // OR:
//   streaming_buffer.AddData({"frame", "motion"},
//                              std::move(input_frame),
//                              std::move(affine_model));
//  // Note: WrapUnique from util::gtl is imported into this namespace when
//  // including this file.
//
//  // Get maximum buffer size.
//  int buffer_size = streaming_buffer.MaxBufferSize("frame");
//
//  // Reached chunk boundary?
//  if (buffer_size == 100) {
//    // Check that we buffered one frame for each motion.
//    CHECK(streaming_buffer.HaveEqualSize({"frame", "motion"}));
//
//    // Compute saliency.
//    for (int k = 0; k < 100; ++k) {
//      cv::Mat& frame = streaming_buffer.GetDatum<cv::Mat>("frame", k);
//      AffineModel* model = streaming_buffer.GetMutableDatum<AffineModel>(
//          "motion", k);
//
//      // OR:
//      std::tuple<cv::Mat*, AffineModel*> data =
//          streaming_buffer.GetMutableData({"frame", "motion"}, k);
//
//      // Resulting saliency.
//      std::unique_ptr<SaliencyPointList> saliency(new SaliencyPointList);
//
//      // Some Computation using frame and model from above ...
//
//      // Store computed saliency.
//      streaming_buffer.AddDatum("saliency", saliency.release());
//    }
//
//    // Output elements, transfers ownership.
//    streaming_buffer.OutputDatum<cv::Mat>(
//      false /*is_flush*/, "frame", [](int frame_idx,
//                                      std::unique_ptr<cv::Mat> frame) {
//    });
//
//    // OR:
//    streaming_buffer.OutputData(/*is_flush*/, {"frame", "motion", "saliency"},
//      [](int frame_idx,
//         std::unique_ptr<cv::Mat> frame,
//         std::unique_ptr<AffineModel> model,
//         std::unique_ptr<SaliencyPointList> saliency) {
//    });
//
//    // Truncate buffers for next chunk.
//    streaming_buffer.TruncateBuffer(false /* is_flush */);
//
//    // End chunk boundary processing.
//  }

// Stores pair (tag, TypeId of type).
typedef std::pair<std::string, size_t> TaggedType;

// Returns TaggedType for type T* tagged with passed string.
template <class T>
TaggedType TaggedPointerType(const std::string& tag);

// Helper function to create unique_ptr's until std::make_unique is allowed.
template <class T>
std::unique_ptr<T> MakeUnique(T* ptr);

// Note: If any of the function below are called with a tag not registered by
// the constructor, the function will fail with CHECK.
// Also, if any of the functions below is called with an existing tag but
// incompatible type, the function will fail with CHECK.
class StreamingBuffer {
 public:
  // Constructs a new buffer with passed mappings (TAG_NAME, DATA_TYPE) of
  // maximum size buffer_size.
  // Data_configuration must have unique tag for each type.
  StreamingBuffer(const std::vector<TaggedType>& data_configuration,
                  int overlap);

  // Call will transfer ownership to StreamingBuffer.
  // Returns true if datum was successfully stored, false otherwise.
  template <class T>
  void AddDatum(const std::string& tag, std::unique_ptr<T> pointer);

  // Same as above but forwards pointer T* to a unique_ptr<T>.
  // Transfers ownership.
  template <class T>
  void EmplaceDatum(const std::string& tag, T* pointer);

  // Creates a deep copy and stores it in the StreamingBuffer.
  template <class T>
  void AddDatumCopy(const std::string& tag, const T& datum);

  // Convenience function to store multiple pointers at once.
  // Returns true if all data was successfully stored, false if at least one
  // failed.
  template <class... Types>
  void AddData(const std::vector<std::string>& tags,
               std::unique_ptr<Types>... pointers);

  // Convenience function to add a whole vector of objects to the buffer.
  // For each datum a copy will be created.
  template <class T>
  void AddDatumVector(const std::string& tag, const std::vector<T>& datum_vec);

  // Retrieves datum with specified tag and frame index. Returns nullptr if
  // datum does not exist.
  template <class T>
  const T* GetDatum(const std::string& tag, int frame_index) const;

  // Gets a reference on the datum.
  template <class T>
  T& GetDatumRef(const std::string& tag, int frame_index) const;

  // Same as above for mutable pointer.
  template <class T>
  T* GetMutableDatum(const std::string& tag, int frame_index) const;

  // Access all elements for a tag as vector, recommended usage:
  // for (auto ptr : buffer.GetDatumVector<SomeType>("some_tag")) { ... }
  template <class T>
  std::vector<const T*> GetDatumVector(const std::string& tag) const;
  template <class T>
  std::vector<T*> GetMutableDatumVector(const std::string& tag) const;

  // Gets a vector of references.
  template <class T>
  std::vector<std::reference_wrapper<T>> GetReferenceVector(
      const std::string& tag) const;

  template <class T>
  std::vector<std::reference_wrapper<const T>> GetConstReferenceVector(
      const std::string& tag) const;

  // Returns number of buffered inputs for specified tag.
  int BufferSize(const std::string& tag) const;

  // Returns maximum over all tags.
  int MaxBufferSize() const;

  // Returns true if the buffers for all passed tags have equal size.
  // Call with HaveEqualSize(AllTags()) to check if all buffers have equal size.
  bool HaveEqualSize(const std::vector<std::string>& tags) const;

  // Check if all items buffered for the specified tag are initialized (i.e. not
  // equal to nullptr).
  template <class T>
  bool IsInitialized(const std::string& tag) const;

  // Returns all tags.
  std::vector<std::string> AllTags() const;

  // Output function to be called after output values have been computed from
  // input values and stored in the buffer as well.
  // Use to transfer ownership for buffered content out of the streaming buffer,
  // by iteratively calling an output functor for each frame.
  //
  // Specifically, functor is repeatedly called with:
  // void operator()(int frame_index, std::unique_ptr<T> pointer)
  //
  // If flush is set, functor will be called with all frames in the half-open
  // interval [0, MaxBufferSize()), otherwise interval is limited to
  // [0, MaxBufferSize() - overlap).
  // Note that the interval is independent of the tag, i.e. for consistency
  // the interval is constant for all tags (if items are not buffered or
  // exists a nullptr is passed instead).
  // Ownership is transferred to caller via unique_ptr.
  // Element storred in buffer is reset to nullptr.
  // Note: Does not truncate the actual buffer, use TruncateBuffer after
  // outputting all desired data.
  template <class T, class Functor>
  void OutputDatum(bool flush, const std::string& tag, const Functor& functor);

  // Same as above for multiple elements.
  // Functor is repeatedly called with:
  // void operator()(int frame_index, std::unique_ptr<Pointers>... pointers)
  template <class... Pointers, class Functor>
  void OutputData(bool flush, const std::vector<std::string>& tags,
                  const Functor& functor);

  // Releases and returns the input at the specified tag and frame_index.
  // Returns nullptr if datum does not exists.
  // Element stored in buffer is reset to nullptr.
  // It is recommended that above OutputFunctions are used instead.
  template <class T>
  std::unique_ptr<T> ReleaseDatum(const std::string& tag, int frame_index);

  // Truncates the buffer by discarding all its elements within the interval
  // [0, MaxBufferSize() - overlap] if flush is set to false, or
  // [0, MaxBufferSize()] otherwise.
  // Should be called after chunk boundary has been reached, computation of
  // outputs from inputs is done and data has been Output via
  // Output[Datum|Data].
  // Returns true if each element within the truncated interval exists,
  // and all buffers have the same number of remaining elements
  // (#overlap if flush is false, zero otherwise).
  bool TruncateBuffer(bool flush);

  // Discards first num_frames of data for specified tag.
  // This is different from ReleaseDatum, as it will shrink the buffer for
  // the specified tag.
  void DiscardDatum(const std::string& tag, int num_frames);

  /// Same as above, but removes num_frames items from the end of the buffer.
  void DiscardDatumFromEnd(const std::string& tag, int num_frames);

  // Same as above for a list of tags.
  void DiscardData(const std::vector<std::string>& tags, int num_frames);

  // Returns true if tag exist.
  bool HasTag(const std::string& tag) const;
  // Returns true if all passed tags exist.
  bool HasTags(const std::vector<std::string>& tags) const;

  // Returns frame index of the first item in the buffer. This is useful to
  // determine how many frames in total went through the buffer after several
  // *successfull* calls of Truncate().
  int FirstFrameIndex() const { return first_frame_index_; }

  template <class T>
  using PointerType = std::shared_ptr<std::unique_ptr<T>>;
  template <class T>
  PointerType<T> CreatePointer(T* t);

 private:
  // Implementation function with aquiring ownership in case failure while
  // adding single argument pointer occurs.
  template <class Pointer, class... Pointers>
  void AddDataImpl(const std::vector<std::string>& tags,
                   std::unique_ptr<Pointer> pointer,
                   std::unique_ptr<Pointers>... pointers);
  // Terminates recursive template expansion for AddDataImpl. Will never be
  // called.
  void AddDataImpl(const std::vector<std::string>& tags) {
    CHECK(tags.empty());
  }

 private:
  int overlap_ = 0;
  int first_frame_index_ = 0;
  absl::node_hash_map<std::string, std::deque<absl::any>> data_;

  // Stores tag, TypeId of corresponding type.
  absl::node_hash_map<std::string, size_t> data_config_;
};

//// Implementation details.
template <class T>
TaggedType TaggedPointerType(const std::string& tag) {
  return std::make_pair(tag,
                        kTypeId<StreamingBuffer::PointerType<T>>.hash_code());
}

template <class T>
StreamingBuffer::PointerType<T> StreamingBuffer::CreatePointer(T* t) {
  return PointerType<T>(new std::unique_ptr<T>(t));
}

template <class T>
void StreamingBuffer::AddDatum(const std::string& tag,
                               std::unique_ptr<T> pointer) {
  CHECK(HasTag(tag));
  CHECK_EQ(data_config_[tag], kTypeId<PointerType<T>>.hash_code());
  auto& buffer = data_[tag];
  absl::any packet(PointerType<T>(CreatePointer(pointer.release())));
  buffer.push_back(packet);
}

template <class T>
void StreamingBuffer::EmplaceDatum(const std::string& tag, T* pointer) {
  std::unique_ptr<T> forwarded(pointer);
  AddDatum(tag, std::move(forwarded));
}

template <class T>
void StreamingBuffer::AddDatumCopy(const std::string& tag, const T& datum) {
  AddDatum(tag, std::unique_ptr<T>(new T(datum)));
}

template <class... Types>
void StreamingBuffer::AddData(const std::vector<std::string>& tags,
                              std::unique_ptr<Types>... pointers) {
  CHECK_EQ(tags.size(), sizeof...(pointers))
      << "Number of tags and data pointers is inconsistent";
  return AddDataImpl(tags, std::move(pointers)...);
}

template <class T>
void StreamingBuffer::AddDatumVector(const std::string& tag,
                                     const std::vector<T>& datum_vec) {
  for (const auto& datum : datum_vec) {
    AddDatumCopy(tag, datum);
  }
}

template <class Pointer, class... Pointers>
void StreamingBuffer::AddDataImpl(const std::vector<std::string>& tags,
                                  std::unique_ptr<Pointer> pointer,
                                  std::unique_ptr<Pointers>... pointers) {
  AddDatum(tags[0], std::move(pointer));
  if (sizeof...(pointers) > 0) {
    return AddDataImpl(std::vector<std::string>(tags.begin() + 1, tags.end()),
                       std::move(pointers)...);
  }
}

template <class T>
std::unique_ptr<T> MakeUnique(T* t) {
  return std::unique_ptr<T>(t);
}

template <class T>
const T* StreamingBuffer::GetDatum(const std::string& tag,
                                   int frame_index) const {
  return GetMutableDatum<T>(tag, frame_index);
}

template <class T>
T& StreamingBuffer::GetDatumRef(const std::string& tag, int frame_index) const {
  return *GetMutableDatum<T>(tag, frame_index);
}

template <class T>
T* StreamingBuffer::GetMutableDatum(const std::string& tag,
                                    int frame_index) const {
  CHECK_GE(frame_index, 0);
  CHECK(HasTag(tag));
  auto& buffer = data_.find(tag)->second;
  if (frame_index > buffer.size()) {
    return nullptr;
  } else {
    const absl::any& packet = buffer[frame_index];
    if (absl::any_cast<PointerType<T>>(&packet) == nullptr) {
      LOG(ERROR) << "Stored item is not of requested type. "
                 << "Check data configuration.";
      return nullptr;
    }

    // Unpack and return.
    const PointerType<T>& pointer =
        *absl::any_cast<const PointerType<T>>(&packet);
    return pointer->get();
  }
}

template <class T>
std::vector<const T*> StreamingBuffer::GetDatumVector(
    const std::string& tag) const {
  std::vector<T*> result(GetMutableDatumVector<T>(tag));
  return std::vector<const T*>(result.begin(), result.end());
}

template <class T>
std::vector<std::reference_wrapper<T>> StreamingBuffer::GetReferenceVector(
    const std::string& tag) const {
  std::vector<T*> ptrs(GetMutableDatumVector<T>(tag));
  std::vector<std::reference_wrapper<T>> refs;
  refs.reserve(ptrs.size());
  for (T* ptr : ptrs) {
    refs.push_back(std::ref(*ptr));
  }
  return refs;
}

template <class T>
std::vector<std::reference_wrapper<const T>>
StreamingBuffer::GetConstReferenceVector(const std::string& tag) const {
  std::vector<const T*> ptrs(GetDatumVector<T>(tag));
  std::vector<std::reference_wrapper<const T>> refs;
  refs.reserve(ptrs.size());
  for (const T* ptr : ptrs) {
    refs.push_back(std::cref(*ptr));
  }
  return refs;
}

template <class T>
bool StreamingBuffer::IsInitialized(const std::string& tag) const {
  CHECK(HasTag(tag));
  const auto& buffer = data_.find(tag)->second;
  int idx = 0;
  for (const auto& item : buffer) {
    const PointerType<T>* pointer = absl::any_cast<const PointerType<T>>(&item);
    CHECK(pointer != nullptr);
    if (*pointer == nullptr) {
      LOG(ERROR) << "Data for " << tag << " at frame " << idx
                 << " is not initialized.";
      return false;
    }
  }
  return true;
}

template <class T>
std::vector<T*> StreamingBuffer::GetMutableDatumVector(
    const std::string& tag) const {
  CHECK(HasTag(tag));
  auto& buffer = data_.find(tag)->second;
  std::vector<T*> result;
  for (const auto& packet : buffer) {
    if (absl::any_cast<PointerType<T>>(&packet) == nullptr) {
      LOG(ERROR) << "Stored item is not of requested type. "
                 << "Check data configuration.";
      result.push_back(nullptr);
    } else {
      result.push_back(
          absl::any_cast<const PointerType<T>>(&packet)->get()->get());
    }
  }
  return result;
}

template <class T, class Functor>
void StreamingBuffer::OutputDatum(bool flush, const std::string& tag,
                                  const Functor& functor) {
  CHECK(HasTag(tag));
  const int end_frame = MaxBufferSize() - (flush ? 0 : overlap_);
  for (int k = 0; k < end_frame; ++k) {
    functor(k, ReleaseDatum<T>(tag, k));
  }
}

template <class T>
std::unique_ptr<T> StreamingBuffer::ReleaseDatum(const std::string& tag,
                                                 int frame_index) {
  CHECK(HasTag(tag));
  CHECK_GE(frame_index, 0);

  auto& buffer = data_.find(tag)->second;
  if (frame_index >= buffer.size()) {
    return nullptr;
  } else {
    const absl::any& packet = buffer[frame_index];
    if (absl::any_cast<PointerType<T>>(&packet) == nullptr) {
      LOG(ERROR) << "Stored item is not of requested type. "
                 << "Check data configuration.";
      return nullptr;
    }

    // Unpack and return.
    const PointerType<T>& pointer =
        *absl::any_cast<const PointerType<T>>(&packet);
    return std::move(*pointer);
  }
}

}  // namespace mediapipe

#endif  // MEDIAPIPE_UTIL_TRACKING_STREAMING_BUFFER_H_