#!/bin/bash
# A script to check for common errors running CUDA/CuDNN/Tensorflow/et al on Ubuntu

WARNINGS=""

function fail_msg() {
    echo -e "[31m$*[39m"
    exit 1
}

function warn_msg() {
    echo -e "[33m$*[39m"
    WARNINGS="${*}\n${WARNINGS}"
}

function green() {
    echo -e "[32m$*[39m"
}



(lspci | grep NVIDIA) && green "Found one or more NVIDIA hardware devices" \
    || warn_msg "Warning: No NVIDIA hardware devices detected. Check that GPU cards are connected and properly powered."


nvidia-smi && green "NVIDIA drivers are installed" \
    || warn_msg "Warning: Could not run nvidia-smi. Check that NVIDIA drivers are installed."


NVIDIA_DRIVER_VERSION=$(nvidia-smi --help | head -1)
green "NVIDIA driver version is: $NVIDIA_DRIVER_VERSION"


CUDA_VERSION=$(ldconfig -p | grep libcuda.so)
if [ "$CUDA_VERSION" ]; then
    green "CUDA is installed with compiler version: $CUDA_VERSION"
else
    warn_msg "Warning: Could not find nvcc."
fi


CUDNN_VERSION=$(ldconfig -p | grep libcudnn)
if [ "$CUDNN_VERSION" ]; then
    green "CUDNN is installed at: $CUDNN_VERSION"
else
    warn_msg "Warning: Could not find libcudnn.so; Tensorflow and Keras may fail"
fi

if [ ! -z "$WARNINGS" ]; then
    echo "Found the following potential issues:"
    echo -e ${WARNINGS}
    exit 1
else
    green "No warnings found"
fi
