#pragma once

#include <torch/arg.h>
#include <torch/types.h>

#include <chrono>
#include <cstddef>

namespace torch {
namespace data {

/// Options to configure a `DataLoader`.
struct DataLoaderOptions {
  DataLoaderOptions() = default;
  /* implicit */ DataLoaderOptions(size_t batch_size)
      : batch_size_(batch_size) {}

  /// The size of each batch to fetch.
  TORCH_ARG(size_t, batch_size) = 1;

  /// The number of worker threads to launch. If zero, the main thread will
  /// synchronously perform the data loading.
  TORCH_ARG(size_t, workers) = 0;

  /// The maximum number of jobs to enqueue for fetching by worker threads.
  /// Defaults to two times the number of worker threads.
  TORCH_ARG(optional<size_t>, max_jobs);

  /// An optional limit on the time to wait for the next batch.
  TORCH_ARG(optional<std::chrono::milliseconds>, timeout);

  /// Whether to enforce ordering of batches when multiple are loaded
  /// asynchronously by worker threads. Set to `false` for better performance if
  /// you do not care about determinism.
  TORCH_ARG(bool, enforce_ordering) = true;

  /// Whether to omit the last batch if it contains less than `batch_size`
  /// examples.
  TORCH_ARG(bool, drop_last) = false;
};

/// Like `DataLoaderOptions`, but without any unconfigured state.
/// `DataLoaderOptions` has some options that depend on other options
/// (`max_jobs` => `2 * workers`). In the spirit of properly using the C++ type
/// system, `DataLoaderOptions` allows only setting values. To access values,
/// you must create a `FullDataLoaderOptions` from a `DataLoaderOptions`
/// instance, which will do any necessary coalescing.
struct FullDataLoaderOptions {
  explicit FullDataLoaderOptions(DataLoaderOptions options)
      : batch_size(options.batch_size()),
        workers(options.workers()),
        max_jobs(options.max_jobs().value_or(2 * workers)),
        timeout(options.timeout()),
        enforce_ordering(options.enforce_ordering()),
        drop_last(options.drop_last()) {}

  size_t batch_size;
  size_t workers;
  size_t max_jobs;
  optional<std::chrono::milliseconds> timeout;
  bool enforce_ordering;
  bool drop_last;
};
} // namespace data
} // namespace torch