#!/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 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

def actions_eta(actions):
	time = 0
	for a in actions:
		i = 0
		while i < len(a):
			if a[i] == '-d':
				i += 1
				if i != len(a):
					time += int(a[i])
			i += 1
	return time

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('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
	commands = ["./" + options.command[0] + "-" + c for c in options.candidates]
	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)

	# ================================================================================
	# Figure out all the combinations to run
	if options.list:
		for a in actions:
			print(" ".join(a))
		sys.exit(0)


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

	# 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)))) ))

	random.shuffle(actions)

	# ================================================================================
	# Run
	options.file.write("[")
	first = True
	for i, a in enumerate(actions):
		sa = " ".join(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:
					fields[match.group(1).strip()] = float(match.group(2).strip().replace(',',''))

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

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

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