#!/usr/bin/python

#
# runptsim.py: Converts XML modules into MPI-C and runs MPICH2
#              Adopted from Eric Work's mod2asap.py
#

import sys, os
import getopt, re
from xml.sax import make_parser
from xml.sax.handler import ContentHandler


# compiled regular expressions
reComment = re.compile(r'//(.*)')
reSpacing = re.compile(r'\s+')


# prints an error when a file can't be opened
def printOpenError(fileName):

	# print error message
	sys.stderr.write('Error: Unable to open %s\n' % fileName)
	sys.exit(1)


# prints an error when a buffer is found unconnected
def printMissingLink(bufferType, number, processName):

	# print error message
	sys.stderr.write(
	                'Error: "%s%d" in module "%s" is not connected\n'
	                % (bufferType, number, processName))
	sys.exit(2)


# prints an error when a module is found unconnected
def printMissingModule(processName):

	# print error message
	sys.stderr.write(
	                'Error: Module "%s" is missing or is not connected\n'
	                % processName)
	sys.exit(3)


# prints an error when a Begin-End mismatch is found
def printBeginEndError(lineNumber, keyword):

	# print error message
	sys.stderr.write(
	                'Error at line %d: Misplaced %s keyword\n'
	                % (lineNumber, keyword))
	sys.exit(4)


# cleans line formatting
def cleanLine(lineIn):

	# convert everything but comments to lowercase
	matComment = reComment.search(lineIn)
	lineOut = reComment.sub('', lineIn).lower()
	if matComment:
		lineOut += '//' + matComment.group(1)

	# collapse whitespace
	lineOut = reSpacing.sub(' ', lineOut).strip()

	return lineOut


# process a file line-by-line
def processFile(modHand, cFileObj, mpiFileObj):

	mpiDoc = MpiCreator()
	mpiDoc.createDocument(modHand, mpiFileObj)
	
	# process line by line
	lineNumber = 1
	for lineIn in cFileObj:
		
		# clean and parse line
		lineOut = cleanLine(lineIn)
		if lineOut:
			mpiDoc.parseCLine(modHand, mpiFileObj, lineOut, lineNumber)

		lineNumber += 1

	# append function test_gen
	mpiDoc.appendTestGen(modHand, mpiFileObj)

	# append function test_read
	mpiDoc.appendTestRead(modHand, mpiFileObj)

	# append function main
	mpiDoc.appendMain(modHand, mpiFileObj)

	# close file objects
	mpiFileObj.close()
	cFileObj.close()


# linkin data type
class LinkIn:

	def __init__(self, src, src_num, dst_num):

		self.datatype = None
		self.src      = src
		self.src_num  = src_num
		self.dst_num  = dst_num


# linkout data type
class LinkOut:

	def __init__(self, src_num, dest, dst_num):

		self.datatype = None
		self.src_num  = src_num
		self.dest     = dest
		self.dst_num  = dst_num


# module data type
class Module:

	def __init__(self, id):

		self.id      = id
		self.path    = None
		self.linkIn  = {}
		self.linkOut = {}

	def addLinkIn(self, src, src_num, dst_num):

		# add input link
		self.linkIn[dst_num] = LinkIn(src, src_num, dst_num)

	def addLinkOut(self, src_num, dest, dst_num):

		# add output link
		self.linkOut[src_num] = LinkOut(src_num, dest, dst_num)


# creates an mpi file
class MpiCreator:

	reParseInclude   = re.compile(r'#include.*')
	reParseBegin     = re.compile(r'begin (\w+)')
	reParseEnd       = re.compile(r'end')
	reParseIbufRead  = re.compile(r'(\w+) = ibuf(\d+)')
	reParseObufWrite = re.compile(r'obuf(\d+)')
	reParseBufType   = re.compile(r'([\w\s]+) (ibuf|obuf)(\d+);')

	bufferDatatype = {'signed char'    : 'MPI_CHAR',
	                  'double'         : 'MPI_DOUBLE',
	                  'float'          : 'MPI_FLOAT',
	                  'int'            : 'MPI_INT',
	                  'long'           : 'MPI_LONG',
	                  'long long'      : 'MPI_LONG_LONG_INT',
	                  'long double'    : 'MPI_LONG_DOUBLE',
	                  'short'          : 'MPI_SHORT',
	                  'unsigned char'  : 'MPI_UNSIGNED_CHAR',
	                  'unsigned int'   : 'MPI_UNSIGNED',
	                  'unsigned long'  : 'MPI_UNSIGNED_LONG',
	                  'unsigned short' : 'MPI_UNSIGNED_SHORT'}

	def __init__(self):

		self.processName  = None
		self.processId    = None
		self.ibuf         = None
		self.ibufDatatype = None
		self.obuf         = None
		self.obufDatatype = None
		self.inFunction   = False

	def createDocument(self, modHand, mpiFileObj):

		# define constants
		moduleKeyList = modHand.module.keys()
		for moduleKey in moduleKeyList:
			moduleId    = moduleKey
			processRank = modHand.module.keys().index(moduleKey)
			mpiFileObj.write(
			                '#define ID_%d %d\n'
			                % (moduleId, processRank))
							
		mpiFileObj.write('\n')

		# include mpi header file
		mpiFileObj.write('#include <mpi.h>\n')

	def parseCLine(self, modHand, mpiFileObj, lineIn, lineNumber):

		# find headers
		matParseInclude = self.reParseInclude.match(lineIn)
		if matParseInclude:
			mpiFileObj.write(lineIn + '\n')
			return
					
		# find begin line
		matParseBegin = self.reParseBegin.match(lineIn)
		if matParseBegin:
			if not self.inFunction:
				self.processName = matParseBegin.group(1)
				self.processId   = self.getProcessId(modHand, self.processName)
				mpiFileObj.write('\nvoid ' + self.processName + ' () {\n')
				self.appendDeclarations(mpiFileObj)
				self.inFunction = True
				return
			else:
				printBeginEndError(lineNumber, 'BEGIN')				
				
		# find end line
		matParseEnd = self.reParseEnd.match(lineIn)
		if matParseEnd:
			if self.inFunction:
				mpiFileObj.write('}\n\n')
				self.inFunction = False
				return
			else:
				printBeginEndError(lineNumber, 'END')	

		# find input buffer read
		matParseIbufRead = self.reParseIbufRead.match(lineIn)
		if matParseIbufRead:
			self.ibuf = int(matParseIbufRead.group(2))
			self.appendRecv(modHand, lineIn, mpiFileObj)
			mpiFileObj.write(lineIn + '\n')
			return

		# find output buffer write
		matParseObufWrite = self.reParseObufWrite.match(lineIn)
		if matParseObufWrite:
			mpiFileObj.write(lineIn + '\n')
			self.obuf = int(matParseObufWrite.group(1)) 
			self.appendSend(modHand, lineIn, mpiFileObj)
			return

		# get buffer data type
		matParseBufType = self.reParseBufType.match(lineIn)
		if matParseBufType:
			cDatatype    = matParseBufType.group(1)
			bufferType   = matParseBufType.group(2)
			bufferNumber = int(matParseBufType.group(3))
			mpiDatatype  = self.bufferDatatype[cDatatype]

			# set MPI datatype of buffer
			if bufferType == 'ibuf':
				try:
					modHand.module[self.processId].linkIn[bufferNumber].datatype = mpiDatatype
				except:
					printMissingLink(bufferType, bufferNumber, self.processName);
			elif bufferType == 'obuf':
				try:
					modHand.module[self.processId].linkOut[bufferNumber].datatype = mpiDatatype	
				except:
					printMissingLink(bufferType, bufferNumber, self.processName);

			mpiFileObj.write(lineIn + '\n')

			return
			
		# do nothing
		mpiFileObj.write(lineIn + '\n')
	
	def getProcessId(self, modHand, processName):

		moduleKeyList = modHand.module.keys()
		for moduleKey in moduleKeyList:
			if modHand.module[moduleKey].path == processName:
				return moduleKey

		printMissingModule(processName)

	# add test_gen function
	def appendTestGen(self, modHand, mpiFileObj):

		dest     = modHand.module[0].linkOut[0].dest
		tag      = modHand.module[0].linkOut[0].dst_num
		datatype = modHand.module[dest].linkIn[tag].datatype
	
		test_input_file = mpiFileObj.name.replace('MPI_', '')
		test_input_file = test_input_file.replace('.c', '_input.dat')

		# append test_gen() function
		mpiFileObj.write(
		  'void test_gen() {\n' + \
		  '  int in_data;\n' + \
		  '  int buffer_size;\n' + \
		  '  double *buffer;\n' + \
		  '  FILE *input;\n' + \
		  '  int retno;\n' + \
		  '  input = fopen("%s", "r");\n'
		     % (test_input_file) + \
		  '  if (input == NULL) {\n' + \
		  '  printf("\\nERROR: Could not open input test file, EXITING\\n");\n' + \
		  '  }\n' + \
		  '  while (retno != EOF) {\n' + \
		  '    retno = fscanf(input, "%d", &in_data);\n' + \
		  '    if (retno != EOF) {\n' + \
		  '      buffer_size = (sizeof(in_data) + MPI_BSEND_OVERHEAD);\n' + \
		  '      buffer = (double *)malloc(buffer_size);\n' + \
		  '      MPI_Buffer_attach(buffer, buffer_size);\n' + \
		  '      MPI_Bsend(&in_data, 1, %s, ID_%d, %d, MPI_COMM_WORLD);\n'
		         % (datatype, dest, tag) + \
		  '      MPI_Buffer_detach(buffer, &buffer_size);\n' + \
		  '      }\n' + \
		  '  }\n' + \
		  '  fclose(input);\n' + \
		  '  printf("test_gen: finished scanning all of input\\n");\n' + \
		  '}\n\n')

	# add test_read function
	def appendTestRead(self, modHand, mpiFileObj):

		self.processId  = max(modHand.module.keys())
		src      = modHand.module[self.processId].linkIn[0].src
		src_num  = modHand.module[self.processId].linkIn[0].src_num
		tag      = modHand.module[self.processId].linkIn[0].dst_num
		datatype = modHand.module[src].linkOut[src_num].datatype
 
		# append test_read() function	
		mpiFileObj.write(
		  'void test_read() {\n' + \
		  '  int out_data;\n' + \
		  '  int i = 1, flag;\n' + \
		  '  MPI_Status status;\n' + \
		  '  double t1, t2;\n' + \
		  '  FILE *output;\n' + \
		  '  output = fopen("chip_output_data.m", "w");\n' + \
		  '  t1 = MPI_Wtime();\n' + \
		  '  for (; !flag;) {\n' + \
		  '    MPI_Iprobe(ID_%d, %d, MPI_COMM_WORLD, &flag, &status);\n'
		       % (src, tag) + \
		  '    t2 = MPI_Wtime();\n' + \
		  '    if (flag) {\n' + \
		  '      MPI_Recv(&out_data, 1, %s, ID_%d, %d, MPI_COMM_WORLD, &status);\n'
		         % (datatype, src, tag) + \
		  '      fprintf(output, "out_code_gen(%6d)=[%6d];\\n", i, out_data);\n' + \
		  '      i++;\n' + \
		  '      t1 = MPI_Wtime();\n' + \
		  '      flag = 0;\n' + \
		  '    }\n' + \
		  '    else if (!flag && t2-t1 > 5) flag = 1;\n' + \
		  '  }\n' + \
		  '  fclose(output);\n' + \
		  '  printf("test_read: finished writing all of output\\n");\n' + \
		  '  MPI_Abort(MPI_COMM_WORLD, 0);\n' + \
		  '}\n\n')

	# add main function
	def appendMain(self, modHand, mpiFileObj):
		
		mpiFileObj.write(
		  'int main(int argc, char **argv) {\n' + \
		  '  int rank, size;\n' + \
		  '  MPI_Status status;\n' + \
		  '  MPI_Request request;\n' + \
		  '  MPI_Init(&argc, &argv);\n' + \
		  '  MPI_Comm_rank(MPI_COMM_WORLD, &rank);\n' + \
		  '  MPI_Comm_size(MPI_COMM_WORLD, &size);\n')

		moduleKeyList = modHand.module.keys()
		for moduleKey in moduleKeyList:
			mpiFileObj.write(
			  '  if (ID_%d == rank) %s();\n'
			     % (moduleKey, modHand.module[moduleKey].path))

		mpiFileObj.write(
		  '  MPI_Finalize();\n' + \
		  '  return 0;\n' + \
		  '}\n\n')

	def appendDeclarations(self, mpiFileObj):

		mpiFileObj.write(
		  'int buffer_size;\n' + \
		  'double *buffer;\n' + \
		  'MPI_Status status;\n' + \
		  'MPI_Request request;\n')

	def appendRecv(self, modHand, lineIn, mpiFileObj):

		datatype = modHand.module[self.processId].linkIn[self.ibuf].datatype
		source   = modHand.module[self.processId].linkIn[self.ibuf].src
		tag      = modHand.module[self.processId].linkIn[self.ibuf].dst_num

		mpiFileObj.write(
		  'MPI_Recv(&ibuf%d, 1, %s, ID_%d, %d, MPI_COMM_WORLD, &status);\n'
		   % (self.ibuf, datatype, source, tag))

	def appendSend(self, modHand, lineIn, mpiFileObj):
	
		datatype = modHand.module[self.processId].linkOut[self.obuf].datatype
		dest     = modHand.module[self.processId].linkOut[self.obuf].dest
		tag      = modHand.module[self.processId].linkOut[self.obuf].dst_num

		mpiFileObj.write(
		  'buffer_size = (sizeof(obuf%d) + MPI_BSEND_OVERHEAD);\n'
		   % self.obuf + \
		  'buffer = (double *)malloc(buffer_size);\n' + \
		  'MPI_Buffer_attach(buffer, buffer_size);\n' + \
		  'MPI_Bsend(&obuf%d, 1, %s, ID_%d, %d, MPI_COMM_WORLD);\n'
		   % (self.obuf, datatype, dest, tag) + \
		  'MPI_Buffer_detach(buffer, &buffer_size);\n') 

 
# sax handler for parsing module files
class ModuleHandler(ContentHandler):

	def __init__(self):
	
		ContentHandler.__init__(self)
		self.module     = {}
		self.moduleId   = None
		self.inPath   	= False
	
	def getModule(self, moduleId):

		# create module object if non-existent
		if not self.module.has_key(moduleId):
			self.module[moduleId] = Module(moduleId)
		
		return self.module[moduleId]

	def startElement(self, name, attrs):

		# set module id
		if name == 'module':
			self.moduleId = int(attrs['id'])
			moduleObj = self.getModule(self.moduleId)

		elif name == 'path':
			self.inPath = True

		# set input and output links
		elif name == 'link':
			src     = self.moduleId
			src_num = int(attrs['src_num'])
			dest    = int(attrs['dest'])
			dst_num = int(attrs['dst_num'])
			
			# set output link of source module
			moduleObj1 = self.getModule(src)
			moduleObj1.addLinkOut(src_num, dest, dst_num)
			
			# set input link of destination module
			moduleObj2 = self.getModule(dest)
			moduleObj2.addLinkIn(src, src_num, dst_num)

		# set input module
		elif name == 'input':
			dest    = int(attrs['id'])
			dst_num = int(attrs['num'])
			src     = 0
			src_num = 0
			
			# set input link of destination module
			moduleObj1 = self.getModule(dest)
			moduleObj1.addLinkIn(src, src_num, dst_num)

			# set output link of source module
			moduleObj2 = self.getModule(0)
			moduleObj2.path = 'test_gen'
			moduleObj2.addLinkOut(src, dest, dst_num)
	
		# set output module
		elif name == 'output':
			src     = int(attrs['id'])
			src_num = int(attrs['num'])
			dest    = max(self.module.keys()) + 1 
			dst_num = 0
			
			# set output link of source module
			moduleObj1 = self.getModule(src)
			moduleObj1.addLinkOut(src_num, dest, dst_num)

			# set input link of destination module
			moduleObj2 = self.getModule(dest)
			moduleObj2.path = 'test_read'
			moduleObj2.addLinkIn(src, src_num, dst_num)

	def characters(self, content):

		if self.inPath:

			# remove leading/tailing whitespace
			content = content.strip()

			# get module path
			moduleObj = self.getModule(self.moduleId)
			moduleObj.path = content.replace('.mod', '')

	def endElement(self, name):	

		# reset path flag
		if name == 'path':
			self.inPath = False


# check if being executed
if __name__ == '__main__':

	# prints the program usage message
	def printUsage():

		# extract application short name
		appName = os.path.basename(sys.argv[0])

		print '\nUsage:', appName, '[OPTION] [FILE] ...\n'
		print 'Options:'
		print '  -h       prints this application usage message'
		print '  FILE     path to C file\n'

		# exit application
		sys.exit(0)

	# parse command-line arguments
	try:
		cmdOpts, cmdArgs = getopt.getopt(sys.argv[1:], 'h')
	except getopt.error, errMsg:
		print 'Error:', str(errMsg).capitalize()
		printUsage()

	# process command-line arguments
	for opt, arg in cmdOpts:
		if opt == '-h': printUsage()

	# check that we are given a filename
	if not cmdArgs: printUsage()
	cFileIn    = cmdArgs[0]
	xmlFileIn  = cFileIn.replace('.c', '.xml')
	mpiFileOut = 'MPI_' + cFileIn

	# open c file
	try:
		cFileObj = open(cFileIn, 'r')
	except IOError:
		printOpenError(cFileIn)

	# open xml file
	try:
		xmlFileObj = open(xmlFileIn, 'r')
	except IOError:
		printOpenError(xmlFileIn)

	# open output file
	try:
		mpiFileObj = open(mpiFileOut, 'w')
	except IOError:
		printOpenError(mpiFileOut)

	# instantiate a new sax parser
	saxParser = make_parser()

	# set the handler to use for parsing
	modHand = ModuleHandler()
	saxParser.setContentHandler(modHand)

	# parse xml file
	try:
		saxParser.parse(xmlFileObj)
	except:
		sys.stderr.write('Error: Failed to parse XML file\n')
		sys.exit(1)

	# create mpi program 
	processFile(modHand, cFileObj, mpiFileObj)

	# run mpi program
	testName = cFileIn.replace('.c', '')
	commSize = len(modHand.module.keys())
	os.system('mpd &')
	os.system('mpicc -c MPI_%s.c' % testName)
	os.system('mpicc -o MPI_%s MPI_%s.o' % (testName, testName))
	os.system('mpiexec -n %d MPI_%s' % (commSize, testName))
	os.system('mpdallexit')
	os.system('rm -rf MPI_%s'   % testName)
	os.system('rm -rf MPI_%s.o' % testName)