import os
import pathlib
import shutil

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

import tensorflow as tf

tf.get_logger().setLevel("ERROR")
tf.keras.utils.set_random_seed(7)

export_dir = pathlib.Path("exported/number_classifier/1")
model_root = export_dir.parents[1]
if model_root.exists():
    shutil.rmtree(model_root)

features = tf.constant(
    [
        [0.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 1.0],
        [1.0, 0.0, 1.0, 0.0],
        [1.0, 1.0, 1.0, 1.0],
        [0.2, 0.9, 0.1, 0.8],
        [0.8, 0.2, 0.9, 0.1],
    ],
    dtype=tf.float32,
)
labels = tf.constant([[0.0], [0.0], [1.0], [1.0], [0.0], [1.0]], dtype=tf.float32)

inputs = tf.keras.Input(shape=(4,), name="features")
x = tf.keras.layers.Dense(8, activation="relu")(inputs)
outputs = tf.keras.layers.Dense(1, activation="sigmoid")(x)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name="number_classifier")

model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
model.fit(features, labels, epochs=12, verbose=0)
model.export(export_dir)

artifact = tf.saved_model.load(str(export_dir))
signature = artifact.signatures["serving_default"]
_, keyword_inputs = signature.structured_input_signature
input_name, input_spec = next(iter(keyword_inputs.items()))

sample_batch = tf.constant(
    [
        [0.1, 0.9, 0.2, 0.8],
        [0.9, 0.1, 0.8, 0.2],
    ],
    dtype=input_spec.dtype,
)
prediction = signature(**{input_name: sample_batch})

print(f"TensorFlow {tf.__version__}")
print(f"Keras {tf.keras.__version__}")
print(f"SavedModel dir: {export_dir}")
print(f"Contains SavedModel: {tf.saved_model.contains_saved_model(str(export_dir))}")
print(f"Signature names: {sorted(artifact.signatures.keys())}")
print(f"Input key: {input_name} shape={tuple(input_spec.shape)} dtype={input_spec.dtype.name}")
print(f"Output keys: {sorted(prediction.keys())}")
print(f"Output shape: {tuple(next(iter(prediction.values())).shape)}")