import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import tensorflow as tf

layers = tf.keras.layers

tf.get_logger().setLevel("ERROR")
tf.keras.utils.set_random_seed(7)

IMAGE_SIZE = 64
BATCH_SIZE = 4

def make_image(index):
    index = tf.cast(index, tf.float32)
    grid = tf.linspace(0.0, 1.0, 48)
    xx, yy = tf.meshgrid(grid, grid)
    xx = tf.expand_dims(xx, -1)
    yy = tf.expand_dims(yy, -1)
    diagonal = tf.cast(xx > yy, tf.float32)
    image = tf.concat(
        [
            tf.clip_by_value(xx + index * 0.04, 0.0, 1.0),
            tf.clip_by_value(1.0 - yy * 0.8, 0.0, 1.0),
            tf.clip_by_value(diagonal * 0.7 + index * 0.02, 0.0, 1.0),
        ],
        axis=-1,
    )
    return tf.cast(image * 255.0, tf.float32)

images = tf.stack([make_image(i) for i in range(8)], axis=0)
labels = tf.constant([0, 1, 0, 1, 0, 1, 0, 1], dtype=tf.int32)

resize_and_rescale = tf.keras.Sequential(
    [
        layers.Resizing(IMAGE_SIZE, IMAGE_SIZE),
        layers.Rescaling(1.0 / 255.0),
    ]
)

augment = tf.keras.Sequential(
    [
        layers.RandomFlip("horizontal", seed=7),
        layers.RandomRotation(0.15, seed=7),
        layers.RandomZoom(0.1, seed=7),
    ]
)

def prepare(image, label):
    image = resize_and_rescale(image)
    return image, label

base_ds = (
    tf.data.Dataset.from_tensor_slices((images, labels))
    .map(prepare, num_parallel_calls=tf.data.AUTOTUNE)
    .cache()
)

train_ds = (
    base_ds
    .shuffle(len(labels), seed=7, reshuffle_each_iteration=False)
    .batch(BATCH_SIZE)
)

train_aug_ds = train_ds.map(
    lambda x, y: (augment(x, training=True), y),
    num_parallel_calls=tf.data.AUTOTUNE,
).prefetch(tf.data.AUTOTUNE)

val_ds = base_ds.batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

first_train_images, first_train_labels = next(iter(train_aug_ds))
second_train_images, _ = next(iter(train_aug_ds))
val_images, _ = next(iter(val_ds))
base_images, _ = next(iter(base_ds.batch(BATCH_SIZE)))

pixel_delta = float(tf.reduce_sum(tf.abs(first_train_images - second_train_images)).numpy())
val_matches_base = bool(tf.reduce_all(tf.abs(val_images - base_images) < 1e-6).numpy())

print(f"TensorFlow {tf.__version__}")
print(f"Keras {tf.keras.__version__}")
print(f"train_batch_shape={tuple(first_train_images.shape)}")
print(f"train_labels={first_train_labels.numpy().tolist()}")
print(f"augmentation_changes_pixels={pixel_delta > 0.0}")
print(f"validation_matches_base={val_matches_base}")
print(
    "pixel_value_range="
    f"({float(tf.reduce_min(first_train_images).numpy()):.3f}, "
    f"{float(tf.reduce_max(first_train_images).numpy()):.3f})"
)