#!/usr/bin/env python

# Matthieu Brucher
# Last Change : 2008-11-19 19:05

import subprocess
import threading
import datetime

names = [("pid", int),
         ("comm", str),
         ("state", str),
         ("ppid", int),
         ("pgrp", int),
         ("session", int),
         ("tty_nr", int),
         ("tpgid", int),
         ("flags", int),
         ("minflt", int),
         ("cminflt", int),
         ("majflt", int),
         ("cmajflt", int),
         ("utime", int),
         ("stime", int),
         ("cutime", int),
         ("cstime", int),
         ("priority", int),
         ("nice", int),
         ("0", int),
         ("itrealvalue", int),
         ("starttime", int),
         ("vsize", int),
         ("rss", int),
         ("rlim", int),
         ("startcode", int),
         ("endcode", int),
         ("startstack", int),
         ("kstkesp", int),
         ("kstkeip", int),
         ("signal", int),
         ("blocked", int),
         ("sigignore", int),
         ("sigcatch", int),
         ("wchan", int),
         ("nswap", int),
         ("cnswap", int),
         ("exit_signal", int),
         ("processor", int),]

colours = ['b', 'g', 'r', 'c', 'm', 'y']

def getPageSize():
  import resource
  f = open("/proc/meminfo")
  mem = f.readline()
  f.close()
  return resource.getpagesize() / (1024 * float(mem[10:-3].strip()))

pagesizepercent = getPageSize()

def collectData(pid, task):
  """
  Collect process list
  """
  f1 = open("/proc/%d/task/%s/stat"%(pid,task))
  f2 = open("/proc/%d/task/%s/statm"%(pid,task))
  t = datetime.datetime.now()
  stat = f1.readline().split()
  mem = f2.readline().split()
  d = dict([(name[0], name[1](el)) for (name, el) in zip(names, stat)])
  d["pmem"] = 100 * float(mem[1]) * pagesizepercent
  return t, d

def getTime(key):
  """
  Returns the time in microseconds
  """
  return (((key.weekday() * 24 + key.hour) * 60 + key.minute) * 60 + key.second) * 1000000 + key.microsecond
  
class MonitorThread(threading.Thread):
  """
  The monitor thread saves the process info every 5 seconds
  """
  def __init__(self, pid):
    import collections

    self.pid = pid
    threading.Thread.__init__(self)
    self.data = collections.defaultdict(dict)
    self.process = True
    
  def run(self):
    import os
    import time

    while self.process:
      threads = os.listdir("/proc/%d/task/" % self.pid)
      for thread in threads:
        t, d = collectData(self.pid, thread)
        d["current_time"] = t
        
        if "now" in self.data[thread]:
          now = self.data[thread]["now"]
          d['pcpu'] = 1e6 * ((d['utime'] + d['stime']) - (now['utime'] + now['stime'])) / float((getTime(t) - getTime(now["current_time"])))

        self.data[thread][getTime(t)] = d
        self.data[thread]["now"] = d
      time.sleep(1)

def displayCPU(data, pid):
  """
  Displays and saves the graph
  """
  import pylab
  import numpy
  
  spid = str(pid)
  
  c = 0
  threads = data.keys()
  threads.sort()
  for thread in threads:
    d = data[thread]
    keys = d.keys()
    keys.remove("now")
    keys.sort()
    mykeys = numpy.array(keys)/1e6
    mykeys -= mykeys[0]
  
    pylab.plot(mykeys[2:], [d[key]['pcpu'] for key in keys[2:]], colours[c], label = thread)
    c = c+1
    if spid == thread:
      pylab.plot(mykeys[2:], [d[key]['pmem'] for key in keys[2:]], 'k', label = 'MEM')

  pylab.ylim([-5, 105])
  pylab.legend(loc=6)
  
  pylab.savefig('%d.svg' % pid)
  pylab.savefig('%d.png' % pid)
  pylab.close()

if __name__ == "__main__":
  import sys
  import os
  import pickle
  
  stdin = open(sys.argv[1])
  stdout = open(sys.argv[2], "w")
  
  process = subprocess.Popen(sys.argv[3:], stdin = stdin, stdout = stdout)
  
  thread = MonitorThread(process.pid)
  thread.start()

  process.wait()

  thread.process = False
  thread.join()
  
  f = open('%d.data' % process.pid, 'w')
  pickle.dump(thread.data, f)
  
  try:
    displayCPU(thread.data, process.pid)
  except:
    pass