import os
import shutil
from pathlib import Path

os.environ.setdefault("KERAS_BACKEND", "tensorflow")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

import keras
import numpy as np
import tensorflow as tf


export_dir = Path("credit-risk-score-savedmodel")
if export_dir.exists():
    shutil.rmtree(export_dir)

keras.utils.set_random_seed(42)

model = keras.Sequential(
    [
        keras.Input(shape=(4,), name="features"),
        keras.layers.Dense(3, activation="relu", name="hidden"),
        keras.layers.Dense(2, activation="softmax", name="risk_score"),
    ]
)

sample = np.array([[0.2, 0.4, 0.1, 0.8]], dtype="float32")
keras_output = model(sample, training=False).numpy()

model.export(export_dir, format="tf_saved_model", verbose=False)

reloaded = tf.saved_model.load(export_dir)
served_output = reloaded.serve(tf.constant(sample)).numpy()
prediction = [round(float(value), 4) for value in served_output[0]]
artifact_files = sorted(path.name for path in export_dir.iterdir())
signature_names = sorted(reloaded.signatures.keys())

print(f"Exported SavedModel directory: {export_dir}")
print(f"SavedModel files: {artifact_files}")
print(f"Available signatures: {signature_names}")
print("Serve endpoint: serve")
print(f"Serve output shape: {served_output.shape}")
print(f"Prediction row: {prediction}")
print(f"Matches Keras output: {np.allclose(keras_output, served_output, atol=1e-6)}")