import tensorflow as tf
import time

def check_tensorflow_gpu():
    print("TensorFlow version:", tf.__version__)
    print("-" * 60)

    # List all physical devices
    print("Available physical devices:")
    for device in tf.config.list_physical_devices():
        print("  •", device)

    print("-" * 60)

    # List all GPU devices
    gpus = tf.config.list_physical_devices('GPU')
    if gpus:
        print(f"GPU detected: {len(gpus)} GPU(s) available.")
        for gpu in gpus:
            print("  ->", gpu)
    else:
        print("No GPU detected. TensorFlow is running on CPU.")
        return

    print("-" * 60)

    # Run a small computation to confirm GPU usage
    print("Running matrix multiplication test...")
    with tf.device('/GPU:0'):
        start = time.time()
        a = tf.random.normal([10000, 10000])
        b = tf.random.normal([10000, 10000])
        c = tf.matmul(a, b)
        tf.experimental.numpy.mean(c).numpy()  # Force computation
        end = time.time()
        print(f"Computation finished in {end - start:.3f} seconds on GPU")

    print("-" * 60)
    print("If you see no CUDA errors and fast computation, TensorFlow is using the GPU.")

if __name__ == "__main__":
    check_tensorflow_gpu()
