#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''
This script uses the Mapper to map simulations with evaluation functions.
'''

import argparse
from math import floor, log10
import tabulate

from hateno.explorer import Explorer, Mapper, MapperUI
from hateno.utils import string, jsonfiles, utils

def addArguments(parser):
	parser.add_argument('--config', type = str, help = 'name of the config to use')
	parser.add_argument('--add-to-manager', action = 'store_true', help = 'add the generated simulations to the manager')
	parser.add_argument('--find', action = 'store_true', help = 'find the true tests after mapping')
	parser.add_argument('--search', type = int, nargs = '?', default = argparse.SUPPRESS, help = 'search for optimal values')
	parser.add_argument('--search-tolerance', type = float, default = 1E-5, help = 'search tolerance')
	parser.add_argument('--search-itermax', type = int, default = 20, help = 'maximal number of iterations')
	parser.add_argument('--map-output', type = str, help = 'path to the file where the map output should be written')
	parser.add_argument('tree', type = str, help = 'path to the file describing the tree to map')
	parser.add_argument('output', type = str, help = 'path to the file where the output should be written')

def main(args):
	folder = utils.findFolder()

	if folder is None:
		print('No Hateno folder found')
		return

	with Mapper(folder, args.config, generate_only = not(args.add_to_manager)) as mapper:
		ui = MapperUI(mapper)

		output = mapper.mapTree(jsonfiles.read(args.tree, allow_generator = True))

		if 'search' in args:
			if args.map_output:
				jsonfiles.write(output, args.map_output)

			explorer = Explorer(mapper)
			found = explorer.search(None, args.search, tolerance = args.search_tolerance, itermax = args.search_itermax)

			jsonfiles.write(found, args.output)

			headers = [''] + [f'{setting["set"]}[{setting["set_index"]}].{setting["name"]}' for setting in found['previous_settings'] + [found['setting']]] + ['Evaluation', 'Criterion']
			values_fmt = f'.{max(-floor(log10(args.search_tolerance)), 0) + 3}g'
			lines_fmt = ('s',) + ('g',)*len(found['previous_settings']) + (values_fmt, 'g')

			for search in found['searches']:
				latest = search['iterations'][-1]

				stopping_criterions = [i['stopping_criterion'] for i in search['iterations']]
				k_best = stopping_criterions.index(min(stopping_criterions))
				best = search['iterations'][k_best]

				print(f'Found between {search["interval"]["bounds"][0]} and {search["interval"]["bounds"][1]} in {string.plural(len(search["iterations"]), "iteration", "iterations")} (criterion: {latest["stopping_criterion"]})')

				line_lower = ['Lower bound', *search['previous_values'], latest['interval']['bounds'][0], latest['interval']['evaluations'][0], '']
				line_upper = ['Upper bound', *search['previous_values'], latest['interval']['bounds'][1], latest['interval']['evaluations'][1], '']
				line_sol = ['Last iteration', *search['previous_values'], latest['iterate'], latest['evaluation'], latest['stopping_criterion']]
				line_best = [f'Minimum (iteration {k_best+1})', *search['previous_values'], best['iterate'], best['evaluation'], best['stopping_criterion']]
				lines = [line_lower, line_upper, line_sol, line_best]

				print(tabulate.tabulate(lines, headers = headers, tablefmt = 'grid', floatfmt = lines_fmt))

		elif args.find:
			if args.map_output:
				jsonfiles.write(output, args.map_output)

			explorer = Explorer(mapper)
			found = explorer.find()

			jsonfiles.write(found, args.output)

			if found:
				for depth, result in found.items():
					print(f'Test at depth {depth}: {result["test"]}')

					if result['true_tests']:
						headers = ['Index'] + [f'{setting["set"]}[{setting["set_index"]}].{setting["name"]}' for setting in result['settings']] + ['Evaluation']

						for test in result['true_tests']:
							lines = [[j, *evaluated['values'], evaluated['evaluation']] for j, evaluated in test.items()]
							print(tabulate.tabulate(lines, headers = headers, tablefmt = 'grid'))

					else:
						print('Nothing found')

			else:
				print('No test defined')

		else:
			jsonfiles.write(output, args.output)

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description = 'Map simulations.')
	addArguments(parser)
	main(parser.parse_args())
