import os
from pathlib import Path

os.environ.setdefault("KERAS_BACKEND", "jax")

import keras
import numpy as np


checkpoint_dir = Path("orbax-credit-score-checkpoints").resolve()
if checkpoint_dir.exists():
    raise SystemExit(f"Remove the existing checkpoint root first: {checkpoint_dir}")

keras.utils.set_random_seed(42)

x_train = np.linspace(-1.0, 1.0, 256, dtype="float32").reshape(64, 4)
y_train = (
    0.6 * x_train[:, :1]
    - 0.2 * x_train[:, 1:2]
    + 0.1 * x_train[:, 2:3]
    + 0.3
)
x_val = np.linspace(-0.8, 0.8, 64, dtype="float32").reshape(16, 4)
y_val = (
    0.6 * x_val[:, :1]
    - 0.2 * x_val[:, 1:2]
    + 0.1 * x_val[:, 2:3]
    + 0.3
)
sample = x_val[:3]

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

checkpoint = keras.callbacks.OrbaxCheckpoint(
    directory=checkpoint_dir,
    max_to_keep=2,
    save_on_background=False,
)

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

restored = keras.saving.load_model(checkpoint_dir)
restored_prediction = restored.predict(sample, verbose=0)
checkpoint_steps = sorted(
    path.name for path in checkpoint_dir.iterdir() if path.is_dir()
)

print(f"Keras backend: {keras.config.backend()}")
print(f"Checkpoint root: {checkpoint_dir}")
print(f"Checkpoint steps kept: {checkpoint_steps}")
print(f"Final val_loss: {history.history['val_loss'][-1]:.4f}")
print(f"Restored model type: {type(restored).__name__}")
print(f"Restored prediction shape: {restored_prediction.shape}")