import os
import struct
from pathlib import Path

os.environ["KERAS_BACKEND"] = "jax"

import keras
from keras import layers


inputs = keras.Input(shape=(12,), name="features")
x = layers.Dense(16, activation="relu", name="feature_encoder")(inputs)
x = layers.Dropout(0.2, name="regularizer")(x)
outputs = layers.Dense(3, activation="softmax", name="risk_score")(x)
model = keras.Model(inputs=inputs, outputs=outputs, name="risk_score_model")

output_path = Path("model-plot.png")
keras.utils.plot_model(
    model,
    to_file=output_path,
    show_shapes=True,
    show_layer_names=True,
    show_layer_activations=True,
    rankdir="TB",
    dpi=160,
)

png_bytes = output_path.read_bytes()
width, height = struct.unpack(">II", png_bytes[16:24])

print(f"backend: {keras.backend.backend()}")
print("layers:", ", ".join(layer.name for layer in model.layers))
print(f"plot file: {output_path}")
print(f"image size: {width}x{height}")
print(f"file bytes: {output_path.stat().st_size}")