import os

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

import tensorflow as tf

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


def print_numerics_error(label, error):
    for line in str(error).splitlines():
        cleaned = line.strip()
        wrapper_end = "}" * 2 + " "
        if cleaned.startswith("{" * 2 + "function_node") and wrapper_end in cleaned:
            cleaned = cleaned.split(wrapper_end, 1)[1]
        if "Detected Infinity or NaN" in cleaned or "Tensor had" in cleaned:
            cleaned = cleaned.replace("!!! ", "").replace(" !!!", "")
            cleaned = cleaned.replace(" (# of outputs: 1)", "")
            cleaned = cleaned.split(" [Op:", 1)[0]
            print(f"{label}: {cleaned}")
            return
    print(f"{label}: {error.__class__.__name__}")


print(f"TensorFlow {tf.__version__}")

safe_tensor = tf.debugging.check_numerics(
    tf.constant([1.0, 2.0], dtype=tf.float32),
    "input batch",
)
print(f"Safe tensor: {safe_tensor.numpy().tolist()}")

try:
    tf.debugging.check_numerics(
        tf.constant([1.0, float("inf")], dtype=tf.float32),
        "manual tensor check",
    )
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("Manual check caught", error)

tf.debugging.enable_check_numerics()

try:
    tf.math.sqrt(tf.constant([4.0, -1.0], dtype=tf.float32))
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("Eager check caught", error)


@tf.function
def unstable_log(values):
    return tf.math.log(values)


try:
    unstable_log(tf.constant([1.0, 0.0], dtype=tf.float32))
except tf.errors.InvalidArgumentError as error:
    print_numerics_error("tf.function check caught", error)

tf.debugging.disable_check_numerics()