#!/usr/bin/env python

import argparse
import arrow
import blessings
import collections
import numpy
import socket
import subprocess
import time

parser = argparse.ArgumentParser(description="Show GPU usage.")
parser.add_argument("--timezone", default="US/Mountain", help="Timezone to be used for display.  Default: %(default)s")
arguments = parser.parse_args()

temperatures = collections.defaultdict(list)
terminal = blessings.Terminal()
with terminal.fullscreen():
    try:
        while True:
            command = ["nvidia-smi", "--query-gpu=index,name,temperature.gpu,utilization.gpu,memory.used,memory.total", "--format=csv,noheader,nounits"]
            results = subprocess.check_output(command).decode("utf-8")

            print(terminal.clear())
            print(u"{t.bold_bright_white}{host}{t.normal}  {t.white}{timestamp}{t.normal}".format(
                host=socket.gethostname(),
                timestamp=arrow.utcnow().to(arguments.timezone).format("ddd MMM DD, YYYY HH:mm:ss"),
                t=terminal,
                ))
            for line in results.strip().split("\n"):
                fields = [field.strip() for field in line.split(",")]
                gpu_id = fields[0]
                temperatures[gpu_id].append(float(fields[2]))
                temperatures[gpu_id] = temperatures[gpu_id][-200:]

                print(u"{t.cyan}[{id}] {t.blue}{name} {t.white}| {t.bold_bright_red}{temp:.1f}\u00B0C{t.normal}{t.white}, {t.bold_bright_green}{util:>3} % {t.normal}{t.white}| {t.bold_bright_yellow}{memory} {t.normal}{t.white}/ {t.normal}{t.yellow}{avail} {t.normal}{t.white}MB{t.normal}".format(
                    id=gpu_id,
                    name=fields[1],
                    temp=numpy.mean(temperatures[gpu_id]),
                    util=fields[3],
                    memory=fields[4],
                    avail=fields[5],
                    t=terminal,
                    ))

            time.sleep(0.1)

    except KeyboardInterrupt:
        pass

