import os

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

import tensorflow as tf

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

gpus = tf.config.list_physical_devices("GPU")
if len(gpus) < 2:
    raise SystemExit(f"Need at least 2 visible GPUs; found {len(gpus)}")

strategy = tf.distribute.MirroredStrategy()
if strategy.num_replicas_in_sync < 2:
    raise SystemExit(
        f"Need at least 2 synchronized replicas; found {strategy.num_replicas_in_sync}"
    )

per_replica_batch = 16
global_batch = per_replica_batch * strategy.num_replicas_in_sync

features = tf.random.stateless_normal((640, 8), seed=(7, 11))
score = (
    features[:, 0] * 0.9
    + features[:, 1] * 0.5
    - features[:, 2] * 0.4
    + features[:, 3] * 0.2
)
labels = tf.cast(score > 0, tf.float32)[:, None]

train_ds = (
    tf.data.Dataset.from_tensor_slices((features[:512], labels[:512]))
    .shuffle(512, seed=7, reshuffle_each_iteration=True)
    .batch(global_batch)
    .prefetch(tf.data.AUTOTUNE)
)
val_ds = (
    tf.data.Dataset.from_tensor_slices((features[512:], labels[512:]))
    .batch(global_batch)
    .prefetch(tf.data.AUTOTUNE)
)

with strategy.scope():
    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(8,)),
            tf.keras.layers.Dense(32, activation="relu"),
            tf.keras.layers.Dense(16, activation="relu"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ]
    )
    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=0.01),
        loss=tf.keras.losses.BinaryCrossentropy(),
        metrics=[tf.keras.metrics.BinaryAccuracy(name="binary_accuracy")],
    )

print(f"visible_gpus={len(gpus)}")
print(f"replicas={strategy.num_replicas_in_sync}")
print(f"global_batch={global_batch}")

history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=3,
    shuffle=False,
    verbose=2,
)

print(f"final_val_binary_accuracy={history.history['val_binary_accuracy'][-1]:.4f}")