import os
import pathlib
import shutil

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

import tensorflow as tf

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

checkpoint_dir = pathlib.Path("training_checkpoints")
if checkpoint_dir.exists():
    shutil.rmtree(checkpoint_dir)

features = tf.constant(
    [
        [0.2, 0.8, 0.4, 0.6],
        [0.9, 0.1, 0.7, 0.8],
        [0.1, 0.7, 0.3, 0.4],
        [0.8, 0.2, 0.6, 0.7],
        [0.3, 0.9, 0.2, 0.3],
        [0.7, 0.3, 0.8, 0.9],
    ],
    dtype=tf.float32,
)
labels = tf.constant([[0.0], [1.0], [0.0], [1.0], [0.0], [1.0]], dtype=tf.float32)
probe = tf.constant([[0.15, 0.75, 0.35, 0.45], [0.85, 0.25, 0.65, 0.75]], dtype=tf.float32)


def build_model():
    return tf.keras.Sequential(
        [
            tf.keras.layers.Input(shape=(4,)),
            tf.keras.layers.Dense(8, activation="relu"),
            tf.keras.layers.Dense(1, activation="sigmoid"),
        ],
        name="support_score",
    )


model = build_model()
optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
model.compile(optimizer=optimizer, loss="binary_crossentropy")
model.fit(features, labels, epochs=8, verbose=0)

training_step = tf.Variable(int(optimizer.iterations), dtype=tf.int64)
checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer, training_step=training_step)
manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)
save_path = manager.save()

before = model(probe, training=False)

restored_model = build_model()
restored_optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
restored_optimizer.build(restored_model.trainable_variables)
restored_step = tf.Variable(0, dtype=tf.int64)
restored_checkpoint = tf.train.Checkpoint(
    model=restored_model,
    optimizer=restored_optimizer,
    training_step=restored_step,
)
latest_path = tf.train.latest_checkpoint(checkpoint_dir)
status = restored_checkpoint.restore(latest_path)
status.assert_consumed()

after = restored_model(probe, training=False)
tf.debugging.assert_near(before, after)

print(f"Saved checkpoint: {save_path}")
print(f"Latest checkpoint: {latest_path}")
print(f"Checkpoint files: {', '.join(sorted(path.name for path in checkpoint_dir.iterdir()))}")
print(f"Restored training step: {int(restored_step.numpy())}")
print(f"Predictions before save: {[round(float(value), 4) for value in before[:, 0]]}")
print(f"Predictions after restore: {[round(float(value), 4) for value in after[:, 0]]}")
print("Restore assertion: consumed")
print("Prediction match: True")

shutil.rmtree(checkpoint_dir)
print(f"Cleaned up: {not checkpoint_dir.exists()}")