import os
import matplotlib.pyplot as plt

import GPUtil

def get_gpu_memory_usage():
    # 获取GPU使用情况
    gpus = GPUtil.getGPUs()
    gpu_memory_info = []
    for gpu in gpus:
        gpu_memory_info.append((gpu.id, gpu.name, f"{gpu.memoryUsed} MB", f"{gpu.memoryTotal} MB", f"{gpu.memoryUtil * 100:.1f}%"))

    return gpu_memory_info

def display_gpu_memory_usage():
    gpu_memory_info = get_gpu_memory_usage()

    if gpu_memory_info:
        print("GPU Memory Usage:")
        for info in gpu_memory_info:
            print(f"GPU ID: {info[0]}, Name: {info[1]}, Memory Used: {info[2]}, Total Memory: {info[3]}, Memory Utilization: {info[4]}")
    else:
        print("No GPU found.")

def color_print(content):
    print(f'\033[1;46m{content}\033[0m\n')

def save_pic_iterly(pic_name, postfix, info):
    pic_idx=1
    pic_name_full=f'{pic_name}_{pic_idx}.{postfix}'

    while os.path.exists(pic_name_full):
        print(f'File {pic_name_full} already exists.')
        pic_idx += 1
        pic_name_full=f'{pic_name}_{pic_idx}.png'

    plt.savefig(pic_name_full, dpi=300, bbox_inches='tight')

    color_print(f'!!!!! {info} is saved in file {pic_name_full}')
