#!/usr/bin/env python
#
# Generates a makefile that executes several simulation runs, and
# optionally invokes "make" as well.
#
# Author: Andras Varga
#

from __future__ import absolute_import, division, print_function, unicode_literals
import argparse
import os
import sys
import re
import subprocess
import tempfile
import multiprocessing
import math

description = """\
Execute a number of OMNeT++ simulation runs, making use of multiple CPUs and
multiple processes. In the simplest case, you would invoke it like this:

% opp_runall ./simprog -c Foo

where "simprog" is an OMNeT++ simulation program, and "Foo" is an omnetpp.ini
configuration, containing iteration variables and/or specifying multiple
repetitions.

The first positional (non-option) argument and all following arguments are
treated as the simulation command (simulation program and its arguments).
Options intended for opp_runall should come before the the simulation command.

To limit the set of simulation runs to be performed, add a -r <runfilter>
argument to the simulation command.

opp_run runs simulations in several batches, making sure to keep all CPUs
busy. Runs of a batch execute sequentially, inside the same Cmdenv process.
The batch size as well as the number of CPUs to use can be overridden.

Command-line options:
"""

epilog = """\
Operation: opp_runall invokes "./simprog -c Foo" with the "-q runnumbers"
extra command-line arguments to figure out how many (and which) simulation
runs it needs to perform, then runs them using multiple Cmdenv processes.
opp_runall exploits GNU Make's -j option to support multiple CPUs or cores:
the commands to perform the simulation runs are written into a temporary
Makefile, and run using "make" with the appropriate -j option. The Makefile
may be exported for inspection.
"""

class Runall:
	def __init__(self):
		self.verbose = False

	def run(self):
		opts = self.parseArgs()

		origwd = os.getcwd()
		if opts.directory != None:
			self.changeDir(opts.directory)

		runNumbers = self.resolveRunNumbers(opts.simProgArgs)

		batchSize = opts.batchsize if opts.batchsize != None else min(5, int(math.ceil(len(runNumbers) / opts.jobs)))
		self.chatter("batch size: " + str(batchSize))
		makefileContent = self.createMakefileContent(opts.simProgArgs, runNumbers, batchSize)

		if opts.export != None:
			self.changeDir(origwd)
			self.saveFile(opts.export, makefileContent)
		else:
			self.runMakefile(makefileContent, opts.jobs)
		sys.exit(0)

	def parseArgs(self):
		parser = argparse.ArgumentParser(description=description, epilog=epilog, formatter_class=argparse.RawDescriptionHelpFormatter)
		parser.add_argument('simProgArgs', metavar='SIMULATION_COMMAND',  nargs=argparse.REMAINDER, help='Simulation command to execute. It should specify the name of the configuration that contains the parameter study (-c CONFIGNAME), and optionally a run filter (-r FILTEREXPR)')
		parser.add_argument('-e', '--export', metavar='MAKEFILE', nargs='?', const='Runfile', default=None, help='Export a makefile that runs the simulations.')
		parser.add_argument('-b', '--batchsize', metavar='N', type=int, default=None, help='Number of simulation runs per Cmdenv instance. Defaults to approximately #runs/#jobs but maximum 5.')
		parser.add_argument('-j', '--jobs', metavar='N', type=int, help='Allow N processes to run at once. Defaults to the number of CPU cores.')
		parser.add_argument('-V', '--verbose', action='store_true', help='Print extra information useful for debugging.')
		parser.add_argument('-C', '--directory', metavar='DIR', help='Change to the given directory before doing anything.')

		try:
			opts = parser.parse_args()
			self.verbose = opts.verbose
		except (IOError, OSError) as e:
			fail(e)
		if not opts.simProgArgs:
			parser.print_help()
			sys.exit(0)

		if opts.jobs == None:
			try:
				opts.jobs = multiprocessing.cpu_count()
				self.chatter("detected " + str(opts.jobs) + " cpus")
			except NotImplementedError:
				self.chatter("could not detect number of cpus (not implemented in this python instance)")

		return opts

	def changeDir(self, directory):
			self.chatter("changing into directory: " + directory)
			try:
				os.chdir(directory)
			except (IOError, OSError) as e:
				self.fail("Cannot change directory: " + str(e))

	def resolveRunNumbers(self, simProgArgs):
		tmpArgs = simProgArgs + ["-s", "-q", "runnumbers"]
		self.chatter("running: " + " ".join(tmpArgs))
		try:
			output = subprocess.check_output(tmpArgs)
		except subprocess.CalledProcessError as e:
			self.fail(simProgArgs[0] + " [...] -q runnumbers returned nonzero exit status")
		except (IOError, OSError) as e:
			self.fail("Cannot execute " + simProgArgs[0] + " [...] -q runnumbers: " + str(e))

		try:
			output = output.decode("utf-8")
			runNumbers = [int(num) for num in output.split()]
			self.chatter("run numbers: " + str(runNumbers))
			return runNumbers
		except:
			self.fail("Error parsing output of " + simProgArgs[0] + " [...] -q runnumbers")


	def createMakefileContent(self, simProgArgs, runNumbers, batchSize):
		tmpArgs = list(simProgArgs)
		if "-r" in tmpArgs:
			i = tmpArgs.index("-r")
			del tmpArgs[i:i+2]
		tmpArgs = [arg for arg in tmpArgs if not arg.startswith("-r")]
		if "-s" not in tmpArgs:
			tmpArgs += ["-s"]
		if "-u" not in tmpArgs:
			tmpArgs += ["-u", "Cmdenv"]
		if not any(arg.startswith("--cmdenv-redirect-output=") for arg in tmpArgs):
			tmpArgs += ["--cmdenv-redirect-output=true"]
		simulationCommand = " ".join(tmpArgs)

		batches = self.chunks(runNumbers, batchSize)
		self.chatter("number of batches: " + str(len(batches)))

		targets = ""
		rules = ""
		i = 0
		for batch in batches:
			target = "run" + str(batch[0]) if len(batch)==1 else "batch" + str(i)
			targets += " " + target
			rules += target + ":\n\t$(SIMULATIONCMD) -r " + ",".join([str(x) for x in batch]) + "\n\n"
			i += 1

		argv = " ".join(sys.argv)

		content = """
#
# This makefile was generated with the following command:
# %(argv)s
#

SIMULATIONCMD = %(simulationCommand)s
TARGETS = %(targets)s

.PHONY: $(TARGETS)

all: $(TARGETS)
	@echo All runs completed.

%(rules)s
""" % locals()

		self.chatter("generated makefile content, " + str(len(content)) + " chars")
		return content

	def runMakefile(self, makefileContent, jobs):
		f = tempfile.NamedTemporaryFile(mode='w', prefix="Makefile.", dir=os.getcwd())
		f.write(makefileContent)
		f.flush()
		makeProg = os.getenv("MAKE", "make")
		makeCmd = [makeProg, "-f", f.name]
		if jobs != None:
			makeCmd += ["-j", str(jobs)]
		self.chatter("running: " + " ".join(makeCmd))
		try:
			exitCode = subprocess.call(makeCmd)
			f.close()
			if exitCode != 0:
				exit(1)  # no need to print an error message, "make" already does that
		except KeyboardInterrupt:
			self.fail("interrupted")

	def chunks(self, lst, n):
		"""Yield successive n-sized chunks from l."""
		n = max(1, n)
		return list(lst[i:i+n] for i in range(0, len(lst), n))

	def saveFile(self, filename, content):
		self.chatter("saving " + filename)
		try:
			f = open(filename, "w")
			f.write(content)
			f.close()
		except (IOError, OSError) as e:
			self.fail("Cannot save makefile into " + filename + ": " + str(e))

	def chatter(self, msg):
		if self.verbose:
			print(msg)

	def fail(self, msg):
		print("opp_runall: " + msg, file=sys.stderr)
		sys.exit(1)

tool = Runall()
tool.run()