# Path: init2winit/dataset_lib/autoaugment.py
# _MAX_LEVEL = 10.
# NAME_TO_FUNC = {
#     'AutoContrast': autocontrast,
#     'Equalize': equalize,
#     'Invert': invert,
#     'Rotate': rotate,
#     'Posterize': posterize,
#     'Solarize': solarize,
#     'SolarizeAdd': solarize_add,
#     'Color': color,
#     'Contrast': contrast,
#     'Brightness': brightness,
#     'Sharpness': sharpness,
#     'ShearX': shear_x,
#     'ShearY': shear_y,
#     'TranslateX': translate_x,
#     'TranslateY': translate_y,
#     'Cutout': cutout,
# }
# def policy_v0():
# def policy_vtest():
# def blend(image1, image2, factor):
# def cutout(image, pad_size, replace=0):
# def solarize(image, threshold=128):
# def solarize_add(image, addition=0, threshold=128):
# def color(image, factor):
# def contrast(image, factor):
# def brightness(image, factor):
# def posterize(image, bits):
# def rotate(image, degrees, replace):
# def translate_x(image, pixels, replace):
# def translate_y(image, pixels, replace):
# def shear_x(image, level, replace):
# def shear_y(image, level, replace):
# def autocontrast(image):
#   def scale_channel(image):
#     def scale_values(im):
# def sharpness(image, factor):
# def equalize(image):
#   def scale_channel(im, c):
#     def build_lut(histo, step):
# def invert(image):
# def wrap(image):
# def unwrap(image, replace):
# def _randomly_negate_tensor(tensor):
# def _rotate_level_to_arg(level):
# def _shrink_level_to_arg(level):
# def _enhance_level_to_arg(level):
# def _shear_level_to_arg(level):
# def _translate_level_to_arg(level, translate_const):
# def level_to_arg(hparams):
# def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
# def _apply_func_with_prob(func, image, args, prob):
# def select_and_apply_random_policy(policies, image):
# def build_and_apply_nas_policy(policies, image,
#                                augmentation_hparams):
#     def make_final_policy(tf_policy_):
#       def final_policy(image_):
# def distort_image_with_autoaugment(image, augmentation_name):
# def distort_image_with_randaugment(image, num_layers, magnitude, key):
# 
# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/dataset_lib/image_preprocessing.py
# def _crop(key, image, hps):
# def crop(key, images, hps):
# def mixup_general(key,
#                   *args_to_mix,
#                   alpha=0.1,
#                   fold_in=None,
#                   n=2,
#                   **kwargs_to_mix):
#   def mix(batch):
#     def mul(a, b):  # B * BHWC -> B111 * BHWC
# def mixup(key, alpha, images, labels):
# def augment_cifar10(key, images, labels, hps):
# def identity_augment(key, images, labels):


from collections import abc
from absl import logging
from init2winit.dataset_lib import autoaugment
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from ml_collections.config_dict import config_dict
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds


def preprocess_for_train(image_bytes,
                         key,
                         dtype=tf.float32,
                         image_size=224,
                         hps=None):
  """Preprocesses the given image for training.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    key: tf random key.
    dtype: data type of the image.
    image_size: image size.
    hps: ConfigDict of options.

  Returns:
    A preprocessed image `Tensor`.
  """
  if hps.get('use_inception_crop'):
    image = _decode_and_inception_crop(image_bytes, image_size)
  else:
    image = _decode_and_random_crop(image_bytes, image_size)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.image.random_flip_left_right(image)

  if hps.get('use_randaug'):
    # NOTE(dsuo): autoaugment code expects uint8 image; not sure why we use
    # float32[0, 255], but just making sure pipeline runs.
    image = tf.cast(tf.clip_by_value(image, 0, 255), tf.uint8)

# Path: init2winit/dataset_lib/autoaugment.py
# _MAX_LEVEL = 10.
# NAME_TO_FUNC = {
#     'AutoContrast': autocontrast,
#     'Equalize': equalize,
#     'Invert': invert,
#     'Rotate': rotate,
#     'Posterize': posterize,
#     'Solarize': solarize,
#     'SolarizeAdd': solarize_add,
#     'Color': color,
#     'Contrast': contrast,
#     'Brightness': brightness,
#     'Sharpness': sharpness,
#     'ShearX': shear_x,
#     'ShearY': shear_y,
#     'TranslateX': translate_x,
#     'TranslateY': translate_y,
#     'Cutout': cutout,
# }
# def policy_v0():
# def policy_vtest():
# def blend(image1, image2, factor):
# def cutout(image, pad_size, replace=0):
# def solarize(image, threshold=128):
# def solarize_add(image, addition=0, threshold=128):
# def color(image, factor):
# def contrast(image, factor):
# def brightness(image, factor):
# def posterize(image, bits):
# def rotate(image, degrees, replace):
# def translate_x(image, pixels, replace):
# def translate_y(image, pixels, replace):
# def shear_x(image, level, replace):
# def shear_y(image, level, replace):
# def autocontrast(image):
#   def scale_channel(image):
#     def scale_values(im):
# def sharpness(image, factor):
# def equalize(image):
#   def scale_channel(im, c):
#     def build_lut(histo, step):
# def invert(image):
# def wrap(image):
# def unwrap(image, replace):
# def _randomly_negate_tensor(tensor):
# def _rotate_level_to_arg(level):
# def _shrink_level_to_arg(level):
# def _enhance_level_to_arg(level):
# def _shear_level_to_arg(level):
# def _translate_level_to_arg(level, translate_const):
# def level_to_arg(hparams):
# def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
# def _apply_func_with_prob(func, image, args, prob):
# def select_and_apply_random_policy(policies, image):
# def build_and_apply_nas_policy(policies, image,
#                                augmentation_hparams):
#     def make_final_policy(tf_policy_):
#       def final_policy(image_):
# def distort_image_with_autoaugment(image, augmentation_name):
# def distort_image_with_randaugment(image, num_layers, magnitude, key):
# 
# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/dataset_lib/image_preprocessing.py
# def _crop(key, image, hps):
# def crop(key, images, hps):
# def mixup_general(key,
#                   *args_to_mix,
#                   alpha=0.1,
#                   fold_in=None,
#                   n=2,
#                   **kwargs_to_mix):
#   def mix(batch):
#     def mul(a, b):  # B * BHWC -> B111 * BHWC
# def mixup(key, alpha, images, labels):
# def augment_cifar10(key, images, labels, hps):
# def identity_augment(key, images, labels):


# Path: init2winit/dataset_lib/imagenet_dataset.py
from collections import abc
from absl import logging
from init2winit.dataset_lib import autoaugment
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from ml_collections.config_dict import config_dict
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds



def preprocess_for_train(image_bytes,
                         key,
                         dtype=tf.float32,
                         image_size=224,
                         hps=None):
  """Preprocesses the given image for training.

  Args:
    image_bytes: `Tensor` representing an image binary of arbitrary size.
    key: tf random key.
    dtype: data type of the image.
    image_size: image size.
    hps: ConfigDict of options.

  Returns:
    A preprocessed image `Tensor`.
  """
  if hps.get('use_inception_crop'):
    image = _decode_and_inception_crop(image_bytes, image_size)
  else:
    image = _decode_and_random_crop(image_bytes, image_size)
  image = tf.reshape(image, [image_size, image_size, 3])
  image = tf.image.random_flip_left_right(image)

  if hps.get('use_randaug'):
    # NOTE(dsuo): autoaugment code expects uint8 image; not sure why we use
    # float32[0, 255], but just making sure pipeline runs.
    image = tf.cast(tf.clip_by_value(image, 0, 255), tf.uint8)

    image = autoaugment.distort_image_with_randaugment(image,