import os
from pathlib import Path

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

import keras
import numpy as np


keras.utils.set_random_seed(42)

model_path = Path("credit-risk-score.keras")
if model_path.exists():
    model_path.unlink()

x_train = np.linspace(0.05, 0.95, 32, dtype="float32").reshape(8, 4)
y_train = ((x_train[:, 0] + x_train[:, 1]) > 0.9).astype("float32").reshape(-1, 1)

model = keras.Sequential(
    [
        keras.Input(shape=(4,), name="features"),
        keras.layers.Dense(6, activation="relu", name="hidden"),
        keras.layers.Dense(1, activation="sigmoid", name="risk_score"),
    ]
)
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.05),
    loss=keras.losses.BinaryCrossentropy(),
    metrics=[keras.metrics.BinaryAccuracy(name="accuracy")],
)
model.fit(x_train, y_train, epochs=8, batch_size=4, verbose=0)

sample = np.array([[0.25, 0.60, 0.35, 0.45]], dtype="float32")
before_save = model.predict(sample, verbose=0)

model.save(model_path)
loaded_model = keras.models.load_model(model_path)
after_load = loaded_model.predict(sample, verbose=0)
loaded_loss, loaded_accuracy = loaded_model.evaluate(x_train, y_train, verbose=0)

print(f"Backend: {keras.backend.backend()}")
print(f"Saved model: {model_path}")
print(f"Saved file exists: {model_path.exists()}")
print(f"Loaded model type: {loaded_model.__class__.__name__}")
print(f"Loaded optimizer: {loaded_model.optimizer.__class__.__name__}")
print(f"Reloaded accuracy: {loaded_accuracy:.4f}")
print(f"Prediction before save: {before_save[0, 0]:.6f}")
print(f"Prediction after load: {after_load[0, 0]:.6f}")
print(f"Predictions match: {np.allclose(before_save, after_load, atol=1e-6)}")