import tensorflow as tf


def setup_gpu(gpu_option, device_number=0):
    import os
    if gpu_option:
        """
        Level | Level for Humans | Level Description                  
        -------|------------------|------------------------------------ 
        0     | DEBUG            | [Default] Print all messages       
        1     | INFO             | Filter out INFO messages           
        2     | WARNING          | Filter out INFO & WARNING messages 
        3     | ERROR            | Filter out all messages    
        """
        os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # remove tensorflow INFO messages
        gpus = tf.config.experimental.list_physical_devices('GPU')
        print("Num GPUs Available: ", len(gpus))
        if gpus:
            # Restrict TensorFlow to only use the first GPU
            try:
                tf.config.experimental.set_visible_devices(gpus[device_number], 'GPU')
                tf.config.experimental.set_memory_growth(gpus[device_number], True)
            except RuntimeError as e:
                # Visible devices must be set at program startup
                print(e)
    else:
        import os
        print('Using CPU')
        os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
        tf.config.experimental.set_visible_devices([], 'GPU')

    print("using GPUs:")
    print(tf.config.experimental.get_visible_devices('GPU'))

    return tf.config.experimental.get_visible_devices('GPU')