#!/usr/bin/python3
"""
Python Script to implement R.M.I.T. testing : Randomized Multiple Interleaved Trials

./rmit.py run COMMAND CANDIDATES
-t trials
-o option:values
"""


import argparse
import datetime
import itertools
import json
import os
import random
import re
import socket
import subprocess
import sys


def parse_range(x):
    result = []
    for part in x.split(','):
        if '-' in part:
            a, b = part.split('-')
            a, b = int(a), int(b)
            result.extend(range(a, b + 1))
        else:
            a = int(part)
            result.append(a)
    return result

class DependentOpt:
	def __init__(self, key, value):
		self.key = key
		self.value = value
		self.vars = re.findall("[a-zA-Z]", value)

def parse_option(key, values):
	try:
		num = int(values)
		return key, [num]
	except:
		pass

	if re.search("^[0-9-,]+$", values):
		values = parse_range(values)
		return key, [v for v in values]
	else:
		return key, DependentOpt(key, values)

def eval_one(fmt, vals):
	orig = fmt
	for k, v in vals:
		fmt = fmt.replace(k, str(v))

	if not re.search("^[0-9-/*+ ]+$", fmt):
		print('ERROR: pattern option {} (interpreted as {}) could not be evaluated'.format(orig, fmt), file=sys.stderr)
		sys.exit(1)

	return eval(fmt)

def eval_options(opts):
	dependents = [d for d in opts.values() if type(d) is DependentOpt]
	processed = []
	nopts = []
	for d in dependents:
		processed.append(d.key)
		lists = []
		for dvar in d.vars:
			if not dvar in opts.keys():
				print('ERROR: extra pattern option {}:{} uses unknown key {}'.format(d.key,d.value,dvar), file=sys.stderr)
				sys.exit(1)

			lists.append([(dvar, o) for o in opts[dvar]])
			processed.append(dvar)

		kopt = []
		for vals in list(itertools.product(*lists)):
			res = ['-{}'.format(d.key), "{}".format(eval_one(d.value, vals))]
			for k, v in vals:
				res.extend(['-{}'.format(k), "{}".format(v)])
			kopt.append(res)
		nopts.append(kopt)


	for k, vals in opts.items():
		if k not in processed:
			kopt = []
			for v in vals:
				kopt.append(['-{}'.format(k), "{}".format(v)])
			nopts.append(kopt)

	return nopts

# returns the first option with key 'opt'
def search_option(action, opt):
	i = 0
	while i < len(action):
		if action[i] == opt:
			i += 1
			if i != len(action):
				return action[i]
		i += 1

	return None

def actions_eta(actions):
	time = 0
	for a in actions:
		o = search_option(a, '-d')
		if o :
			time += int(o)
	return time

taskset_maps = None

def init_taskset_maps():
	global taskset_maps
	known_hosts = {
		"jax": {
			range(  1,  24) : "48-71",
			range( 25,  48) : "48-71,144-167",
			range( 49,  96) : "48-95,144-191",
			range( 97, 144) : "24-95,120-191",
			range(145, 192) : "0-95,96-191",
		},
	}

	if (host := socket.gethostname()) in known_hosts:
		taskset_maps = known_hosts[host]
		return True

	print("Warning unknown host '{}', disable taskset usage".format(host), file=sys.stderr)
	return False


def settaskset_one(action):
	o = search_option(action, '-p')
	if not o:
		return action
	try:
		oi = int(o)
	except ValueError:
		return action

	m = "Not found"
	for key in taskset_maps:
		if oi in key:
			return ['taskset', '-c', taskset_maps[key], *action]

	print("Warning no mapping for {} cores".format(oi), file=sys.stderr)
	return action

def settaskset(actions):
	return [settaskset_one(a) for a in actions]

if __name__ == "__main__":
	# ================================================================================
	# parse command line arguments
	formats = ['raw', 'csv', 'json']
	parser = argparse.ArgumentParser(description='Python Script to implement R.M.I.T. testing : Randomized Multiple Interleaved Trials')
	parser.add_argument('--list', help='List all the commands that would be run', action='store_true')
	parser.add_argument('--file', nargs='?', type=argparse.FileType('w'), default=sys.stdout)
	parser.add_argument('--trials', help='Number of trials to run per combinaison', type=int, default=3)
	parser.add_argument('--notaskset', help='If specified, the trial will not use taskset to match the -p option', action='store_true')
	parser.add_argument('command', metavar='command', type=str, nargs=1, help='the command prefix to run')
	parser.add_argument('candidates', metavar='candidates', type=str, nargs='*', help='the candidate suffix to run')

	try:
		options, unknown =  parser.parse_known_args()

		options.option = []
		while unknown:
			key = unknown.pop(0)
			val = unknown.pop(0)

			if key[0] != '-':
				raise ValueError

			options.option.append((key[1:], val))

	except:
		print('ERROR: invalid arguments', file=sys.stderr)
		parser.print_help(sys.stderr)
		sys.exit(1)

	# ================================================================================
	# Identify the commands to run
	command = './' + options.command[0]
	if options.candidates:
		commands = [command + "-" + c for c in options.candidates]
	else:
		commands = [command]
	for c in commands:
		if not os.path.isfile(c):
			print('ERROR: invalid command {}, file does not exist'.format(c), file=sys.stderr)
			sys.exit(1)

		if not os.access(c, os.X_OK):
			print('ERROR: invalid command {}, file not executable'.format(c), file=sys.stderr)
			sys.exit(1)


	# ================================================================================
	# Identify the options to run
	opts = dict([parse_option(k, v) for k, v in options.option])

	# Evaluate the options (options can depend on the value of other options)
	opts = eval_options(opts)

	# ================================================================================
	# Figure out all the combinations to run
	actions = []
	for p in itertools.product(range(options.trials), commands, *opts):
		act = [p[1]]
		for o in p[2:]:
			act.extend(o)
		actions.append(act)

	# ================================================================================
	# Fixup the different commands

	# Add tasksets
	withtaskset = False
	if not options.notaskset and init_taskset_maps():
		withtaskset = True
		actions = settaskset(actions)

	# ================================================================================
	# Now that we know what to run, print it.
	# find expected time
	time = actions_eta(actions)
	print("Running {} trials{}".format(len(actions), "" if time == 0 else " (expecting to take {})".format(str(datetime.timedelta(seconds=int(time)))) ))

	# dry run if options ask for it
	if options.list:
		for a in actions:
			print(" ".join(a))
		sys.exit(0)


	# ================================================================================
	# Prepare to run

	random.shuffle(actions)

	# ================================================================================
	# Run
	options.file.write("[")
	first = True
	for i, a in enumerate(actions):
		sa = " ".join(a[3:] if withtaskset else a)
		if first:
			first = False
		else:
			options.file.write(",")
		if options.file != sys.stdout:
			print("{}/{} : {}          \r".format(i, len(actions), sa), end = '')
		fields = {}
		with subprocess.Popen( a, stdout  = subprocess.PIPE, stderr  = subprocess.PIPE) as proc:
			out, err = proc.communicate()
			if proc.returncode != 0:
				print("ERROR: command '{}' encountered error, returned code {}".format(sa, proc.returncode), file=sys.stderr)
				print(err.decode("utf-8"))
				sys.exit(1)
			for s in out.decode("utf-8").splitlines():
				match = re.search("^(.*):(.*)$", s)
				if match:
					try:
						fields[match.group(1).strip()] = float(match.group(2).strip().replace(',',''))
					except:
						pass

		options.file.write(json.dumps([a[3 if withtaskset else 0][2:], sa, fields]))
		options.file.flush()

	options.file.write("]\n")

	if options.file != sys.stdout:
		print("Done                                                                                ")