# Path: fedjax/core/client_datasets.py
# class BatchPreprocessor:
# class PaddedBatchHParams:
# class ShuffleRepeatBatchHParams:
# class BatchHParams:
# class ClientDataset:
# class PaddedBatchView:
# class ShuffleRepeatBatchView:
# class BatchView:
#   def __init__(self, fns: Iterable[Callable[[Examples], Examples]] = ()):
#   def __call__(self, examples: Examples) -> Examples:
#   def append(self, fn: Callable[[Examples], Examples]) -> 'BatchPreprocessor':
#   def __str__(self) -> str:
#   def __repr__(self) -> str:
#   def __init__(self,
#                raw_examples: Examples,
#                preprocessor: BatchPreprocessor = NoOpBatchPreprocessor):
#   def __len__(self) -> int:
#   def __getitem__(self, index: slice) -> 'ClientDataset':
#   def all_examples(self) -> Examples:
#   def padded_batch(self,
#                    hparams: Optional[PaddedBatchHParams] = None,
#                    **kwargs) -> Iterable[Examples]:
#   def shuffle_repeat_batch(self,
#                            hparams: Optional[ShuffleRepeatBatchHParams] = None,
#                            **kwargs) -> Iterable[Examples]:
#   def batch(self,
#             hparams: Optional[BatchHParams] = None,
#             **kwargs) -> Iterable[Examples]:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: PaddedBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def _pick_final_batch_size(data_size: int, batch_size: int,
#                            num_batch_size_buckets: int) -> int:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: ShuffleRepeatBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
#   def __init__(self, client_dataset: ClientDataset, hparams: BatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def padded_batch_client_datasets(datasets: Iterable[ClientDataset],
#                                  hparams: Optional[PaddedBatchHParams] = None,
#                                  **kwargs) -> Iterator[Examples]:
# def buffered_shuffle(source: Iterable[Any], buffer_size: int,
#                      rng: np.random.RandomState) -> Iterator[Any]:
# def buffered_shuffle_batch_client_datasets(
#     datasets: Iterable[ClientDataset], batch_size: int, buffer_size: int,
#     rng: np.random.RandomState) -> Iterator[Examples]:
#   def gen_items():
# def assert_consistent_rows(examples: Examples):
# def num_examples(examples: Examples, validate: bool = True) -> int:
# def slice_examples(examples: Examples, index: slice) -> Examples:
# def concat_examples(many_examples: Iterable[Examples]) -> Examples:
# def attach_mask(examples: Examples, mask: np.ndarray) -> Examples:
# def pad_examples(examples: Examples, size: int) -> Examples:
# EXAMPLE_MASK_KEY = '__mask__'
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


import timeit
import numpy as np
from absl import app
from fedjax.core import client_datasets
from fedjax.core import util

def main(_):
  for mode in ['train', 'eval']:
    for preprocess in [False, True]:
      assert (bench_client_dataset(preprocess,
                                   mode) == bench_tf_dataset(preprocess, mode))
      print(
          f'ClientDataset\tmode={mode}\tpreprocess={preprocess}\t',
          timeit.timeit(
              lambda: bench_client_dataset(preprocess, mode), number=100))
      print(
          f'TF Dataset\tmode={mode}\tpreprocess={preprocess}\t',
          timeit.timeit(lambda: bench_tf_dataset(preprocess, mode), number=100))


FAKE_MNIST = {
    'pixels': np.random.uniform(size=(1000, 28, 28)),
    'label': np.random.randint(10, size=(1000,))
}


def f(x):
  return {**x, 'binary_label': x['label'] % 2}


def bench_client_dataset(preprocess, mode, batch_size=128, num_steps=100):
  """Benchmarks ClientDataset."""
  preprocessor = client_datasets.NoOpBatchPreprocessor
  if preprocess:
    preprocessor = preprocessor.append(f)

# Path: fedjax/core/client_datasets.py
# class BatchPreprocessor:
# class PaddedBatchHParams:
# class ShuffleRepeatBatchHParams:
# class BatchHParams:
# class ClientDataset:
# class PaddedBatchView:
# class ShuffleRepeatBatchView:
# class BatchView:
#   def __init__(self, fns: Iterable[Callable[[Examples], Examples]] = ()):
#   def __call__(self, examples: Examples) -> Examples:
#   def append(self, fn: Callable[[Examples], Examples]) -> 'BatchPreprocessor':
#   def __str__(self) -> str:
#   def __repr__(self) -> str:
#   def __init__(self,
#                raw_examples: Examples,
#                preprocessor: BatchPreprocessor = NoOpBatchPreprocessor):
#   def __len__(self) -> int:
#   def __getitem__(self, index: slice) -> 'ClientDataset':
#   def all_examples(self) -> Examples:
#   def padded_batch(self,
#                    hparams: Optional[PaddedBatchHParams] = None,
#                    **kwargs) -> Iterable[Examples]:
#   def shuffle_repeat_batch(self,
#                            hparams: Optional[ShuffleRepeatBatchHParams] = None,
#                            **kwargs) -> Iterable[Examples]:
#   def batch(self,
#             hparams: Optional[BatchHParams] = None,
#             **kwargs) -> Iterable[Examples]:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: PaddedBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def _pick_final_batch_size(data_size: int, batch_size: int,
#                            num_batch_size_buckets: int) -> int:
#   def __init__(self, client_dataset: ClientDataset,
#                hparams: ShuffleRepeatBatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
#   def __init__(self, client_dataset: ClientDataset, hparams: BatchHParams):
#   def __iter__(self) -> Iterator[Examples]:
# def padded_batch_client_datasets(datasets: Iterable[ClientDataset],
#                                  hparams: Optional[PaddedBatchHParams] = None,
#                                  **kwargs) -> Iterator[Examples]:
# def buffered_shuffle(source: Iterable[Any], buffer_size: int,
#                      rng: np.random.RandomState) -> Iterator[Any]:
# def buffered_shuffle_batch_client_datasets(
#     datasets: Iterable[ClientDataset], batch_size: int, buffer_size: int,
#     rng: np.random.RandomState) -> Iterator[Examples]:
#   def gen_items():
# def assert_consistent_rows(examples: Examples):
# def num_examples(examples: Examples, validate: bool = True) -> int:
# def slice_examples(examples: Examples, index: slice) -> Examples:
# def concat_examples(many_examples: Iterable[Examples]) -> Examples:
# def attach_mask(examples: Examples, mask: np.ndarray) -> Examples:
# def pad_examples(examples: Examples, size: int) -> Examples:
# EXAMPLE_MASK_KEY = '__mask__'
# 
# Path: fedjax/core/util.py
# def safe_div(a, b):
#   def __getattr__(self, _):
# def import_tf():
# class FakeTensorFlowModule:


# Path: fedjax/core/client_datasets_benchmark.py
import timeit
import numpy as np
from absl import app
from fedjax.core import client_datasets
from fedjax.core import util


def main(_):
  for mode in ['train', 'eval']:
    for preprocess in [False, True]:
      assert (bench_client_dataset(preprocess,
                                   mode) == bench_tf_dataset(preprocess, mode))
      print(
          f'ClientDataset\tmode={mode}\tpreprocess={preprocess}\t',
          timeit.timeit(
              lambda: bench_client_dataset(preprocess, mode), number=100))
      print(
          f'TF Dataset\tmode={mode}\tpreprocess={preprocess}\t',
          timeit.timeit(lambda: bench_tf_dataset(preprocess, mode), number=100))


FAKE_MNIST = {
    'pixels': np.random.uniform(size=(1000, 28, 28)),
    'label': np.random.randint(10, size=(1000,))
}


def f(x):
  return {**x, 'binary_label': x['label'] % 2}


def bench_client_dataset(preprocess, mode, batch_size=128, num_steps=100):
  """Benchmarks ClientDataset."""
  preprocessor = client_datasets.NoOpBatchPreprocessor
  if preprocess:
    preprocessor = preprocessor.append(f)

  dataset = client_datasets.ClientDataset(FAKE_MNIST, preprocessor)