import os

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

import numpy as np
import tensorflow as tf


SEED = 2026

tf.keras.utils.set_random_seed(SEED)
tf.config.experimental.enable_op_determinism()


def train_once():
    tf.keras.backend.clear_session()
    tf.keras.utils.set_random_seed(SEED)

    features = np.array(
        [
            [0.0, 0.0],
            [0.0, 1.0],
            [1.0, 0.0],
            [1.0, 1.0],
            [2.0, 1.0],
            [1.0, 2.0],
            [2.0, 2.0],
            [3.0, 1.0],
        ],
        dtype=np.float32,
    )
    labels = np.array(
        [[0.0], [1.0], [1.0], [2.0], [3.0], [3.0], [4.0], [4.0]],
        dtype=np.float32,
    )

    dataset = (
        tf.data.Dataset.from_tensor_slices((features, labels))
        .shuffle(buffer_size=len(features), seed=SEED, reshuffle_each_iteration=False)
        .batch(4)
    )

    model = tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(2,)),
            tf.keras.layers.Dense(4, activation="relu"),
            tf.keras.layers.Dense(1),
        ]
    )
    model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=0.05), loss="mse")
    model.fit(dataset, epochs=5, verbose=0, shuffle=False)

    weights = [weight.numpy().copy() for weight in model.weights]
    predictions = model(features, training=False).numpy()
    return weights, predictions


first_weights, first_predictions = train_once()
second_weights, second_predictions = train_once()

weights_match = all(
    np.array_equal(left, right) for left, right in zip(first_weights, second_weights)
)
predictions_match = np.array_equal(first_predictions, second_predictions)

print("TensorFlow:", tf.__version__)
print("Deterministic operations enabled")
print("Weights match:", weights_match)
print("Predictions match:", predictions_match)
print("Prediction sample:", np.round(first_predictions[:3, 0], 6).tolist())