import os
from pathlib import Path

os.environ["KERAS_BACKEND"] = "jax"

import keras
import numpy as np


def build_model(compile_model=True):
    model = keras.Sequential(
        [
            keras.Input(shape=(4,), name="features"),
            keras.layers.Dense(8, activation="relu", name="hidden"),
            keras.layers.Dense(1, activation="sigmoid", name="score"),
        ]
    )
    if compile_model:
        model.compile(
            optimizer=keras.optimizers.Adam(learning_rate=0.04),
            loss=keras.losses.BinaryCrossentropy(),
        )
    return model


keras.utils.set_random_seed(29)
weights_path = Path("risk-score.weights.h5")
if weights_path.exists():
    weights_path.unlink()

x_train = np.linspace(0.0, 1.0, 64, dtype="float32").reshape(16, 4)
y_train = ((x_train[:, 1] + x_train[:, 3]) > 0.9).astype("float32")
sample = np.array([[0.2, 0.7, 0.4, 0.8]], dtype="float32")

model = build_model()
model.fit(x_train, y_train, epochs=10, batch_size=4, verbose=0)
before = model.predict(sample, verbose=0)
model.save_weights(weights_path)

clone = build_model(compile_model=False)
clone(np.zeros((1, 4), dtype="float32"))
clone.load_weights(weights_path)
after = clone.predict(sample, verbose=0)

print(f"backend: {keras.backend.backend()}")
print(f"weights file: {weights_path}")
print(f"weights file exists: {weights_path.exists()}")
print(f"prediction before save: {before[0, 0]:.6f}")
print(f"prediction after load: {after[0, 0]:.6f}")
print(f"predictions match: {np.allclose(before, after, atol=1e-6)}")