import keras
import numpy as np


keras.utils.set_random_seed(42)

x_train = np.linspace(-1.0, 1.0, 160, dtype="float32").reshape(80, 2)
y_train = (0.75 * x_train[:, :1]) - (0.25 * x_train[:, 1:2])
x_val = np.linspace(-0.8, 0.8, 40, dtype="float32").reshape(20, 2)
y_val = (0.75 * x_val[:, :1]) - (0.25 * x_val[:, 1:2])

model = keras.Sequential(
    [
        keras.layers.Input(shape=(2,)),
        keras.layers.Dense(8, activation="relu"),
        keras.layers.Dense(1),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.05),
    loss="mse",
    metrics=["mae"],
)

checkpoint_path = "checkpoints/best.keras"
checkpoint = keras.callbacks.ModelCheckpoint(
    filepath=checkpoint_path,
    monitor="val_loss",
    mode="min",
    save_best_only=True,
    verbose=1,
)

history = model.fit(
    x_train,
    y_train,
    validation_data=(x_val, y_val),
    epochs=4,
    batch_size=16,
    callbacks=[checkpoint],
    verbose=0,
)

loaded = keras.models.load_model(checkpoint_path)
loaded_loss, loaded_mae = loaded.evaluate(x_val, y_val, verbose=0)

print(f"Saved checkpoint: {checkpoint_path}")
print(f"Best val_loss: {min(history.history['val_loss']):.4f}")
print(f"Loaded val_loss: {loaded_loss:.4f}")
print(f"Loaded val_mae: {loaded_mae:.4f}")