import os
from pathlib import Path

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

import keras
import numpy as np
import onnxruntime as ort


export_path = Path("credit-risk-score.onnx")
if export_path.exists():
    export_path.unlink()

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_path, format="onnx", verbose=False)

session = ort.InferenceSession(
    str(export_path),
    providers=["CPUExecutionProvider"],
)
input_meta = session.get_inputs()[0]
output_meta = session.get_outputs()[0]
onnx_output = session.run(None, {input_meta.name: sample})[0]
prediction = [round(float(value), 4) for value in onnx_output[0]]

print(f"Exported ONNX file: {export_path}")
print(f"ONNX file size: {export_path.stat().st_size} bytes")
print(f"ONNX input: {input_meta.name} {input_meta.shape} {input_meta.type}")
print(f"ONNX output: {output_meta.name} {output_meta.shape} {output_meta.type}")
print(f"ONNX output shape: {onnx_output.shape}")
print(f"Prediction row: {prediction}")
print(f"Matches Keras output: {np.allclose(keras_output, onnx_output, atol=1e-5)}")