import os
import shutil
from pathlib import Path

os.environ["KERAS_BACKEND"] = "tensorflow"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import keras
import numpy as np


keras.utils.set_random_seed(21)

export_dir = Path("score_savedmodel")
if export_dir.exists():
    shutil.rmtree(export_dir)

model = keras.Sequential(
    [
        keras.Input(shape=(3,), name="features"),
        keras.layers.Dense(4, activation="relu", name="hidden"),
        keras.layers.Dense(1, activation="sigmoid", name="score"),
    ]
)
model.compile(optimizer="adam", loss="binary_crossentropy")
model(np.zeros((1, 3), dtype="float32"))
model.export(export_dir, verbose=False)

try:
    keras.models.load_model(export_dir)
except ValueError as exc:
    load_model_error = str(exc).split(":", 1)[0]
else:
    load_model_error = "load_model unexpectedly succeeded"

layer = keras.layers.TFSMLayer(export_dir, call_endpoint="serve")
inputs = np.array([[0.2, 0.6, 0.4], [0.9, 0.1, 0.7]], dtype="float32")
outputs = layer(inputs)
output_array = np.asarray(outputs)

print(f"backend: {keras.backend.backend()}")
print(f"export directory: {export_dir}")
print(f"load_model result: {load_model_error}")
print(f"reloaded layer: {layer.__class__.__name__}")
print("call endpoint: serve")
print(f"output shape: {output_array.shape}")
print(f"first score: {output_array[0, 0]:.6f}")