from msg import *
from sg_script_api import *


PORT_22 = 0

FINALIZE = 221297

slaves_count = 0
stopSim = False
slaves = []

graph = None
curve = None

task_count = 0
number_of_tasks = 0
executed_tasks = 0

def slave(argv):
	global PORT_22
	global FINALIZE
	global slaves_count
	global stopSim
	global slaves
	global executed_tasks
	
	g = PyMsgGlobal.instance()

	cur_slave = g.getCurrentHost().getName()

	print 'Start slave', cur_slave
	
	script_api = SimgridSimulationInterface.instance()
	
	script_api.incComponentRedLevel(g.getCurrentHost().getName())
	
	slaves_count = slaves_count + 1
	
	finalize = False
	
	while finalize == False or stopSim:
		task = g.getTaskFromCurrentHost(PORT_22)
		if task != None:
			print 'Slave ', cur_slave, 'Received \"', task.getName(), '\"'
			
			if task.getData() == FINALIZE:
				print 'Slave ', cur_slave, ' exiting.'
				task.destroy()
				finalize = True
				break
			else:
				print 'Slave ', cur_slave, 'Processing \"', task.getName(), '\"'
				script_api.setComponentMessage(g.getCurrentHost().getName(), "Executing task: " + task.getName())
				task.execute()
				executed_tasks = executed_tasks + 1
				curve.addPoint(g.getClock(), executed_tasks)
				
				script_api.setComponentMessage(g.getCurrentHost().getName(), "")
				print 'Slave ', cur_slave, '\"', task.getName(), '\" done'
				task.destroy()
		else:
			print "Hey? What's up?"
			return -1

	script_api.decComponentRedLevel(g.getCurrentHost().getName())
	
	print 'Slave ', cur_slave, 'Done. See you!'
	return 0


def master(argv):
	global PORT_22
	global FINALIZE
	global stopSim

	number_of_tasks = int(argv[1])
	task_comp_size = int(argv[2])
	task_comm_size = int(argv[3])
	
	todo = []
	
	g = PyMsgGlobal.instance()

	script_api = SimgridSimulationInterface.instance()
	script_api.incComponentRedLevel(g.getCurrentHost().getName())
	
	print '\nArgs:'
	print 'Tasks:', number_of_tasks
	
	for i in range(number_of_tasks):
		if stopSim:
			break

		todo.append(PyMsgTask('Task_' + str(i), task_comp_size, task_comm_size, None))
	
	slaves_count = len(argv) - 4
	
	for i in range(slaves_count):
		slave = argv[4+i]
		slaves.append(g.getHostByName(slave));
		if slaves[i] == None:
			print "Unknown host ", slave, " Stopping Now! ", slave
			return -1
	
	print "Got ", len(slaves), " slaves:"

	if not stopSim:	
		for i in range(slaves_count):
			print slaves[i].getName()
	
	for i in range(number_of_tasks):
		if stopSim:
			break

		sched_slave = i % slaves_count
		
		print "Sending ", todo[i].getName(), " to ", slaves[sched_slave].getName()
		if g.getCurrentHost().getName() == slaves[sched_slave].getName():
			print "Hey! It's me! :)\n"
		slaves[sched_slave].putTask(todo[i], PORT_22)
	
	print "All tasks have been dispatched. Let's tell everybody the computation is over."
	for i in range(slaves_count):
		print 'Sending Finalize to ', slaves[i].getName()
		slaves[i].putTask(PyMsgTask("Finalize", 0, 0, FINALIZE), PORT_22)
		
	script_api.decComponentRedLevel(g.getCurrentHost().getName())

	print "Goodbye now!\n"
	
	return 0		

def stop():
	global stopSim
	print 'Stoping Simulation'
	stopSim = True
	print 'Stoping Simulation Finished'
	return

def main(platFile, deployFile):
	global slaves_count
	global stopSim
	global slaves
	global script_api
	global curve
	global executed_tasks
	
	slaves_count = 0
	stopSim = False
	slaves = []
	executed_tasks = 0

	g = PyMsgGlobal.instance()
	script_api = SimgridSimulationInterface.instance()

	graph = script_api.createGraphic()
	graph.setTitle("Task Execution")
	graph.setXText("Time")
	graph.setYText("Tasks Executed")

	curve = graph.createCurve()
	curve.setText("Tasks Executed")
	curve.setWidth(2)
	curve.addPoint(0, 0)

	g.globalInit(['testMsg'])
	g.setChannelNumber(1)
	g.setPajeOutput("msg_test.trace")
	g.createEnvironment(platFile)
	
	g.functionRegister("master", master)
	g.functionRegister("slave", slave)
	
	g.launchApplication(deployFile)
	res = g.main()
	g.destroy()
