import os
import time

import numpy as np

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

import tensorflow as tf


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

AUTOTUNE = tf.data.AUTOTUNE
EXAMPLES = 640
FEATURES = 3
BATCH_SIZE = 32
EPOCHS = 2
DELAY_SECONDS = 0.002
FEATURE_SCALE = np.array([1000.0, 100.0, 10.0], dtype=np.float32)


features = tf.random.uniform(
    (EXAMPLES, FEATURES),
    minval=0.0,
    maxval=100.0,
    dtype=tf.float32,
    seed=7,
)
labels = tf.cast(tf.reduce_sum(features, axis=1) > 150.0, tf.float32)


def pause_and_normalize(value):
    time.sleep(DELAY_SECONDS)
    return (value.astype(np.float32) / FEATURE_SCALE).astype(np.float32)


def slow_preprocess(feature_row, label):
    normalized = tf.numpy_function(pause_and_normalize, [feature_row], tf.float32)
    normalized.set_shape([FEATURES])
    return normalized, label


def build_pipeline(optimized):
    dataset = tf.data.Dataset.from_tensor_slices((features, labels))
    if optimized:
        dataset = dataset.map(slow_preprocess, num_parallel_calls=AUTOTUNE)
        dataset = dataset.cache()
        dataset = dataset.shuffle(
            buffer_size=EXAMPLES,
            seed=2026,
            reshuffle_each_iteration=False,
        )
        dataset = dataset.batch(BATCH_SIZE)
        dataset = dataset.prefetch(AUTOTUNE)
        return dataset

    dataset = dataset.map(slow_preprocess)
    dataset = dataset.shuffle(
        buffer_size=EXAMPLES,
        seed=2026,
        reshuffle_each_iteration=False,
    )
    dataset = dataset.batch(BATCH_SIZE)
    return dataset


def consume(dataset):
    batches = 0
    checksum = 0.0
    started = time.perf_counter()
    for _ in range(EPOCHS):
        for batch_features, _ in dataset:
            checksum += float(tf.reduce_sum(batch_features).numpy())
            batches += 1
    return time.perf_counter() - started, batches, checksum


baseline_seconds, baseline_batches, baseline_checksum = consume(build_pipeline(False))
optimized_dataset = build_pipeline(True)
optimized_seconds, optimized_batches, optimized_checksum = consume(optimized_dataset)
batch_features, batch_labels = next(iter(optimized_dataset))
speedup = baseline_seconds / optimized_seconds

print(f"TensorFlow {tf.__version__}")
print(f"examples={EXAMPLES}")
print(f"batch_size={BATCH_SIZE}")
print(f"baseline_seconds={baseline_seconds:.3f}")
print(f"optimized_seconds={optimized_seconds:.3f}")
print(f"speedup={speedup:.2f}x")
print(f"optimized_batches={optimized_batches}")
print(f"train_batch_shape={tuple(batch_features.shape)}")
print(f"train_label_shape={tuple(batch_labels.shape)}")
print(f"checksums_match={abs(baseline_checksum - optimized_checksum) < 1e-3}")
print(f"optimized_dataset={type(optimized_dataset).__name__}")