# Standard library imports
import subprocess

def query_nvsmi(properties, index=None):
	query = ['nvidia-smi', '--query-gpu={}'.format(properties), '--format=csv,noheader,nounits']

	if index is not None:
		query.append('--id={}'.format(index))

	process = subprocess.Popen(query, stdout=subprocess.PIPE)
	output = process.stdout.read().decode()
	rows = []

	for line in output.splitlines():
		rows.append(line.rstrip().split(', '))

	return rows

class GPU():
	def __init__(self, index, name, total_memory):
		self.name = name
		self.total_memory = int(total_memory)
		self.index = int(index)

	def get_clock_speeds(self):
		row = query_nvsmi('clocks.gr,clocks.mem', self.index)[0]

		return {
			'core_clock_speed': int(row[0]),
			'memory_clock_speed': int(row[1])
		}

	def get_max_clock_speeds(self):
		row = query_nvsmi('clocks.max.gr,clocks.max.mem', self.index)[0]

		return {
			'max_core_clock_speed': int(row[0]),
			'max_memory_clock_speed': int(row[1])
		}

	def get_memory_details(self):
		row = query_nvsmi('memory.used,memory.free', self.index)[0]

		return {
			'used_memory': int(row[0]),
			'free_memory': int(row[1])
		}

def get_gpus():
	rows = query_nvsmi('index,name,memory.total')
	gpus = []

	for row in rows:
		gpus.append(GPU(*row))

	return gpus

if __name__ == '__main__':
	for gpu in get_gpus():
		print(gpu.__dict__)
		print(gpu.get_max_clock_speeds())
		print(gpu.get_clock_speeds())
		print(gpu.get_memory_details())
