#pragma once

#include <torch/data/datasets/base.h>
#include <torch/data/example.h>
#include <torch/types.h>

#include <cstddef>
#include <vector>

namespace torch {
namespace data {
namespace datasets {

/// A dataset of tensors.
/// Stores a single tensor internally, which is then indexed inside `get()`.
struct TensorDataset : public Dataset<TensorDataset, TensorExample> {
  /// Creates a `TensorDataset` from a vector of tensors.
  explicit TensorDataset(const std::vector<Tensor>& tensors)
      : TensorDataset(torch::stack(tensors)) {}

  explicit TensorDataset(torch::Tensor tensor) : tensor(std::move(tensor)) {}

  /// Returns a single `TensorExample`.
  TensorExample get(size_t index) override {
    return tensor[index];
  }

  /// Returns the number of tensors in the dataset.
  optional<size_t> size() const override {
    return tensor.size(0);
  }

  Tensor tensor;
};

} // namespace datasets
} // namespace data
} // namespace torch