#!/usr/bin/env python

# Copyright (c) 2014 by Pavel Sountsov
#
# All rights reserved. Distributed under GPL 3.0. For full terms see the file LICENSE.

from ConfigParser import SafeConfigParser, NoOptionError
from argparse import ArgumentParser
from subprocess import Popen, PIPE

import fileinput
import shlex
import re

from optimizer2.differential_evolution import DifferentialEvolutionOptimizer
from optimizer2.cont_differential_evolution import ContDifferentialEvolutionOptimizer
from optimizer2.array_parser import parse_array

# pop is an array of vectors with the first element being the fitness
# and the remaining elements being the parameters
def runner(cmd_str, res_re, max_launches, pop, verbose):
	idx = 0
	
	while idx < len(pop):
		num_launches = min(len(pop) - idx, max_launches)
		procs = []
		cmds = []
		for i in range(num_launches):
			cmds.append(cmd_str.format(*pop[idx + i][1:]))
			if verbose:
				print 'Running:', cmds[i]
			args = shlex.split(cmds[i])
			procs.append(Popen(args, stdout=PIPE))
		
		for i in range(num_launches):
			out, _ = procs[i].communicate()
			if procs[i].returncode != 0:
				print out
				raise Exception('\'' + cmds[i] + '\' returned a non-0 status!')
			if res_re:
				match = res_re.search(out)
				if match:
					l = match.group(1)
				else:
					raise Exception('Could not parse the output!')
			else:
				l = out.splitlines()[-1]
				
			pop[idx + i][0] = float(l)
		idx += num_launches

def main():
	parser = ArgumentParser()
	parser.add_argument('--cfg', dest='cfg_name', help='Configuration file name. If omitted, the file will be read from stdin.')
	
	args = parser.parse_args()
	
	cfg = SafeConfigParser();
	
	if args.cfg_name:
		cfg.read(args.cfg_name)
	else:
		cfg.readfp(fileinput.input())
	
	cmd_str = cfg.get('options', 'command')
	try:
		max_launches = cfg.getint('options', 'max_launches')
	except NoOptionError:
		max_launches = 1000

	try:
		pat = cfg.get('options', 'result_re')
		res_re = re.compile(pat)
	except NoOptionError:
		res_re = None
	
	try:
		num_args = cfg.getint('options', 'num_args')
	except NoOptionError:
		num_args = 1

	limits = []
	for i in range(num_args):
		try:
			lim_str = cfg.get('options', 'limit%d' % i)
			lim = parse_array(lim_str)
			if len(lim) != 2:
				raise Exception('limit%d invalid: (expected "[low, hi]" but got "%s")' % (i, lim_str))
		except NoOptionError:
			lim = [0.0, 1.0]
		limits.append(lim)
	
	alg_str = cfg.get('options', 'algorithm')
	if alg_str == 'de':
		opt = DifferentialEvolutionOptimizer(cfg, limits, lambda pop: runner(cmd_str, res_re, max_launches, pop, False))	
	elif alg_str == 'cont_de':
		opt = ContDifferentialEvolutionOptimizer(cfg, limits, cmd_str, res_re, max_launches)
	else:
		raise Exception('Unknown algorithm \'' + alg_str + '\'')
	
	pop = opt.run()
	print 'Final population:'
	for ind in pop:
		print ind[1:], 'fit:', ind[0]

if __name__ == '__main__':
    main()
