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)

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

x_train = features[:256]
y_train = labels[:256]
x_val = features[256:]
y_val = labels[256:]

train_ds = (
    tf.data.Dataset.from_tensor_slices((x_train, y_train))
    .shuffle(256, seed=7, reshuffle_each_iteration=True)
    .batch(32)
)
val_ds = tf.data.Dataset.from_tensor_slices((x_val, y_val)).batch(32)

model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(8,)),
        tf.keras.layers.Dense(16, activation="relu"),
        tf.keras.layers.Dense(8, activation="relu"),
        tf.keras.layers.Dense(1, activation="sigmoid"),
    ]
)

loss_fn = tf.keras.losses.BinaryCrossentropy()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)

train_loss = tf.keras.metrics.Mean(name="train_loss")
train_accuracy = tf.keras.metrics.BinaryAccuracy(name="train_accuracy")
val_loss = tf.keras.metrics.Mean(name="val_loss")
val_accuracy = tf.keras.metrics.BinaryAccuracy(name="val_accuracy")


@tf.function
def train_step(features, labels):
    with tf.GradientTape() as tape:
        predictions = model(features, training=True)
        loss = loss_fn(labels, predictions)
        if model.losses:
            loss += tf.add_n(model.losses)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    train_loss.update_state(loss)
    train_accuracy.update_state(labels, predictions)


@tf.function
def val_step(features, labels):
    predictions = model(features, training=False)
    loss = loss_fn(labels, predictions)
    val_loss.update_state(loss)
    val_accuracy.update_state(labels, predictions)


for epoch in range(1, 5):
    train_loss.reset_state()
    train_accuracy.reset_state()
    val_loss.reset_state()
    val_accuracy.reset_state()

    for features_batch, labels_batch in train_ds:
        train_step(features_batch, labels_batch)

    for features_batch, labels_batch in val_ds:
        val_step(features_batch, labels_batch)

    print(
        f"Epoch {epoch}: "
        f"train_loss={train_loss.result():.4f} "
        f"train_accuracy={train_accuracy.result():.4f} "
        f"val_loss={val_loss.result():.4f} "
        f"val_accuracy={val_accuracy.result():.4f}"
    )

predictions = model(x_val[:4], training=False)[:, 0]
formatted = [round(float(value), 4) for value in predictions]
print(f"sample_predictions={formatted}")