import os
import shutil
from pathlib import Path

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

import tensorflow as tf


logdir = Path("logs/profile-demo")
if logdir.exists():
    shutil.rmtree(logdir)
logdir.mkdir(parents=True)

features = tf.random.uniform((512, 10), seed=42)
labels = tf.cast(tf.reduce_sum(features, axis=1, keepdims=True) > 5.0, tf.float32)
dataset = (
    tf.data.Dataset.from_tensor_slices((features, labels))
    .shuffle(512, seed=42)
    .batch(32)
    .prefetch(tf.data.AUTOTUNE)
)

model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(10,)),
        tf.keras.layers.Dense(16, activation="relu"),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ]
)
model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])

tf.profiler.experimental.start(str(logdir))
try:
    model.fit(dataset, epochs=2, verbose=0, shuffle=False)
finally:
    tf.profiler.experimental.stop()

profile_files = sorted(logdir.glob("plugins/profile/*/*"))
profile_dirs = sorted({profile_file.parent for profile_file in profile_files})

print(f"TensorBoard log directory: {logdir}")
print(f"Profiler files: {len(profile_files)}")
for profile_dir in profile_dirs:
    print(f"Profile directory: {profile_dir}")
for profile_file in profile_files:
    print(f"Trace file: {profile_file.name}")