"""

this script is used to render cpu usage
to get cpu usage, here use top rather than ps because ps got bad accuarcy, for top got firefox's cpu
usage, 8.3 or 8.4, while ps always got 6.2
"""

import sys
import matplotlib as mpl # render diagram to pdf rather than x window because our target machine may not allowed using x11 
mpl.use('Agg') # this must called before import pylab
import pylab
import subprocess
import time
import datetime
import getopt
import re


TotalTime = 30 # seconds
Interval = 1 # seconds
MoitorProcess = ""
LabelYTime = '0' 
ThreadName = "" 

def getProcessCPUUsage( processName ):
	global Interval

	p1 = subprocess.Popen(["top", "-n", str(Interval), "-b"], stdout=subprocess.PIPE)
	p2 = subprocess.Popen(["grep", processName], stdin=p1.stdout, stdout=subprocess.PIPE)
	p3 = subprocess.Popen(["grep", "-v", "grep"], stdin=p2.stdout, stdout=subprocess.PIPE)
	# here we just extact first process that is matched with process name
	p4 = subprocess.Popen(["awk", "{if( NR==1 ) print $9}"], stdin=p3.stdout, stdout=subprocess.PIPE) 
	usg = p4.communicate()[0]
	print "process cpu: " , usg 
	# use regular express remove none decimal charactar
	non_decimal = re.compile(r'[^\d.]+')
	cpu = non_decimal.sub('', usg )
	return float( cpu )

def getProcessId( processName ):
	p1 = subprocess.Popen(["top", "-n", "1", "-b"], stdout=subprocess.PIPE)
	p2 = subprocess.Popen(["grep", processName], stdin=p1.stdout, stdout=subprocess.PIPE)
	p3 = subprocess.Popen(["grep", "-v", "grep"], stdin=p2.stdout, stdout=subprocess.PIPE)
	# here we just extact first process that is matched with process name
	p4 = subprocess.Popen(["awk", "{if( NR==1 ) print $1}"], stdin=p3.stdout, stdout=subprocess.PIPE) 
	pid = p4.communicate()[0]
	return pid

def getProcessThread( pid, threadName ):
	pid=str(int(pid)) # force tick out none-decimal 
	p1 = subprocess.Popen(["top", "-p", pid,  "-n", str(Interval), "-H", "-b"], stdout=subprocess.PIPE)
	p2 = subprocess.Popen(["grep", threadName], stdin=p1.stdout, stdout=subprocess.PIPE)
	p3 = subprocess.Popen(["grep", "-v", "grep"], stdin=p2.stdout, stdout=subprocess.PIPE)
	p4 = subprocess.Popen(["wc", "-l"], stdin=p2.stdout, stdout=subprocess.PIPE)
	if '0' == p4.communicate()[0]:
		return False
	else:
		return True

def getProcessThreadCPUUsage( processName, threadName ):
	global Interval

	pid = getProcessId( processName )
	if( '' == pid or None == pid ):
		print "cannot find process: ", processName
		return	
	#print 'pid:', pid
	pid=str(int(pid)) # force tick out none-decimal 
	p1 = subprocess.Popen(["top", "-p", pid,  "-n", str(Interval), "-H", "-b"], stdout=subprocess.PIPE)
	p2 = subprocess.Popen(["grep", threadName], stdin=p1.stdout, stdout=subprocess.PIPE)
	p3 = subprocess.Popen(["grep", "-v", "grep"], stdin=p2.stdout, stdout=subprocess.PIPE)
	# here we just extact first process that is matched with process name
	p4 = subprocess.Popen(["awk", "{if( NR==1 ) print $9}"], stdin=p3.stdout, stdout=subprocess.PIPE) 
	usg = p4.communicate()[0]
	if usg == '':
		print "cannot find thread: ", threadName
		return 
	print "thread cpu: " , usg 
	# use regular express remove none decimal charactar
	non_decimal = re.compile(r'[^\d.]+')
	cpu = non_decimal.sub('', usg )
	return float( cpu )

def renderFromFile( fileName ):
	print "input file: ", fileName
	# generate data
	i = 0
	x = list()
	y = list()
	with open( fileName, 'r' ) as pf:
		for line in pf:
			y.append( round( float( line ), 2))
			x.append( i )
			i += 1
	pylab.plot( x, y )
	#pylab.show()
	pylab.savefig( fileName + '.png' )

def renderFromListXY( listX, listY ):
	# generate data
	#i = 0
	#x = list()
	#y = list()
	#for el in data:
	#	y.append( round( float( el ), 2))
	#	x.append( i )
	#	i += 1

	global MoitorProcess
	pylab.xlabel( "time( in seconds )" )
	pylab.ylabel( "cpu usage" )
	pylab.title( "CPU Usage of Process( " + MoitorProcess + ":" + ThreadName + " )" )
	pylab.plot( listX, listY, aa=True, label="CPU Usage of " + MoitorProcess )
	pylab.gcf().autofmt_xdate()
	#pylab.show()

def renderFromListY( listY ):
	# generate data
	i = 0
	x = list()
	y = list()
	for el in listY:
		y.append( round( float( el ), 2))
		x.append( i )
		i += 1

	global MoitorProcess
	pylab.xlabel( "time( in seconds )" )
	pylab.ylabel( "cpu usage" )
	pylab.title( "CPU Usage of Process( " + MoitorProcess + ":" + ThreadName + " )" )
	pylab.plot( x, y, aa=True )
	pylab.gcf().autofmt_xdate()
	#pylab.show()

def savePlot( processName, threadName ):
	if threadName == '':
		pylab.savefig( processName + '.png' )
	else:
		pylab.savefig( processName + '-' + threadName + '.png' )


def usage():
	print "usage: python render.py [options]"
	print "options:"
	print "	-h help"
	print "	-t total time"
	print "	-i interval to pick up cpu usage"
	print "	-p process name"
	print "	-r thread name"
	print "	-l 1: mark X axis with time, 0: without time"
	print "	-m mode, by default, use 'top', value can be top or ps"
	sys.exit( 0 )

def parseOptions( argv ):
	argv.remove( argv[0])
	optlist, args = getopt.getopt(argv, 'ht:i:p:r:l:m:')
	bTotalTime = False
	bInterval = False
	bProcessName = False
	bLabelYTime = False
	bThreadName = False

	try:
		for opt in optlist:
			if opt[0] == '-h':
				usage()
				return
			if opt[0] == '-t':
				global TotalTime 
				TotalTime = int(opt[1])
				bTotalTime = True
			if opt[0] == '-i':
				global Interval
				Interval = int(opt[1])
				bInterval = True
			if opt[0] == '-p':
				global MoitorProcess
				MoitorProcess = opt[1]
				bProcessName = True
			if opt[0] == '-r':
				global ThreadName
				ThreadName = opt[1]
				bThreadName = True
			if opt[0] == '-l':
				global LabelYTime 
				LabelYTime = opt[1]
				bLabelYTime = True
			if opt[0] == '-m':
				pass
			else:
				pass
		if bProcessName == False:
			print "process name must be provided"
			usage()
			return
		if bTotalTime == False:
			print "total time not specified, use default value - 30 seconds"
		if bInterval == False:
			print "interval not provided, use default value - 1 second"
		if bLabelYTime == False:
			print "draw diagram without time"
		else:
			print "draw diagram with time"
	except Exception as e:
		#print "got exception: " + e.message()
		usage()
    	

def main( argv ):
	# get input file
	#fileName = argv[1]
	#renderFromFile( fileName )
	parseOptions( argv )

	global TotalTime
	global Interval
	global MoitorProcess  
	global LabelYTime
	global ThreadName

	# check if process or thread exists, if not exit program
	pid = getProcessId( MoitorProcess )
	if '' == pid:
		print "cannot find process: " + MoitorProcess
		sys.exit( 0 )
	if False == getProcessThread( pid, ThreadName ):
		print "cannot find thread: " + ThreadName
		sys.exit( 0 )

	listUage = list()
	listTime = list()
	count = TotalTime / Interval
	print "capture data...."
	print "total time: ", TotalTime, " interval: ", Interval
	for it in range( count ):
		if ThreadName != "":
			usage = getProcessThreadCPUUsage( MoitorProcess, ThreadName )
		else:
			usage = getProcessCPUUsage( MoitorProcess )
		listUage.append( usage )
		#tm = time.localtime()
		#listTime.append( str(tm.tm_hour)+":"+str(tm.tm_min)+":"+str(tm.tm_sec))
		listTime.append( datetime.datetime.now())
		#time.sleep( Interval )
	
	# render
	if LabelYTime == '1':
		renderFromListXY( listTime, listUage )
	else:
		renderFromListY( listUage )
	savePlot( MoitorProcess, ThreadName  )


if __name__ == "__main__":
	main( sys.argv )