import os
import pathlib
import shutil

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

import tensorflow as tf

tf.get_logger().setLevel("ERROR")

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

model = tf.keras.Sequential(
    [
        tf.keras.layers.Input(shape=(4,), name="features"),
        tf.keras.layers.Dense(
            1,
            activation="sigmoid",
            name="priority_score",
        ),
    ],
    name="support_priority",
)

model(tf.zeros((1, 4), dtype=tf.float32))
model.get_layer("priority_score").set_weights(
    [
        tf.constant(
            [[1.2], [0.8], [-0.6], [1.0]],
            dtype=tf.float32,
        ).numpy(),
        tf.constant([-0.4], dtype=tf.float32).numpy(),
    ]
)

@tf.function(
    input_signature=[
        tf.TensorSpec(
            shape=(None, 4),
            dtype=tf.float32,
            name="features",
        )
    ]
)
def serve(features):
    scores = model(features, training=False)
    return {"scores": scores}

@tf.function(
    input_signature=[
        tf.TensorSpec(
            shape=(None, 4),
            dtype=tf.float32,
            name="features",
        )
    ]
)
def serve_labels(features):
    scores = model(features, training=False)
    labels = tf.cast(scores >= 0.5, tf.int32)
    return {"scores": scores, "labels": labels}

export_archive = tf.keras.export.ExportArchive()
export_archive.track(model)
export_archive.add_endpoint(name="serve", fn=serve)
export_archive.add_endpoint(
    name="labels",
    fn=serve_labels,
)
export_archive.write_out(export_dir, verbose=False)

artifact = tf.saved_model.load(str(export_dir))
sample_batch = tf.constant(
    [
        [0.2, 0.9, 0.1, 0.8],
        [0.8, 0.3, 0.9, 0.2],
    ],
    dtype=tf.float32,
)
serve_signature = artifact.signatures["serve"]
labels_signature = artifact.signatures["labels"]
serve_outputs = serve_signature(features=sample_batch)
labels_outputs = labels_signature(features=sample_batch)
signature_names = sorted(artifact.signatures.keys())

print("Primary signature: serve")
print("Additional signature: labels")
print(f"Default alias present: {'serving_default' in signature_names}")
print(f"Serve keys: {sorted(serve_outputs.keys())}")
print(f"Labels keys: {sorted(labels_outputs.keys())}")
print(f"SavedModel dir: {export_dir}")