#!/usr/bin/env python3
import concurrent.futures
import glob
import json
import multiprocessing
import os
import shutil
import subprocess
import sys
import tempfile
import uuid

class Executer(object):
    def __init__(self, tempdir, source, testcases):
        self.tempdir = os.path.abspath(tempdir)
        self.source = self.tempdir + '/' + os.path.basename(source)
        shutil.copyfile(source, self.source)
        self.testcases = testcases

    def compile(self, params):
        defines = ' '.join(f'-D{k.upper()}={v}' for k, v in params.items())
        binary = self.tempdir + '/' + str(uuid.uuid4())
        command = ' '.join([
            os.environ.get('CXX', 'g++'),
            os.environ.get('CXXFLAGS', '-std=c++14 -Wall -O2'),
            defines,
            '-o', binary,
            self.source,
        ]).split(' ')
        result = subprocess.run(command, stderr=subprocess.DEVNULL)
        result.check_returncode()
        return binary

    def execute(self, binary, testcase):
        with open(testcase) as fh:
            result = subprocess.run(binary, stdin=fh, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
        result.check_returncode()
        line = result.stderr.splitlines()[-1].decode()
        assert line.startswith('raw score = ')
        return int(line.split()[-1])

    def evaluate(self, raw_scores):
        return sum(raw_scores) / len(raw_scores) * 32

    def __call__(self, **kwargs):
        binary = self.compile(kwargs)
        with concurrent.futures.ThreadPoolExecutor(max_workers=multiprocessing.cpu_count()) as executor:
            futures = []
            for testcase in self.testcases:
                futures.append(executor.submit(self.execute, binary, testcase))
            return self.evaluate([ f.result() for f in futures ])

def plot(optimizer, name, path):
    import matplotlib.pyplot as plt
    import numpy as np
    import seaborn as sns
    xs = np.array([ params[name] for params in optimizer.res['all']['params'] ])
    ys = np.array(optimizer.res['all']['values'])
    plt.figure(figsize=(16, 9))
    sns.regplot(xs, ys, order=4)  # order=2 is weak, order=3 is 
    plt.savefig(path)

def main():
    # input
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('testcase', nargs='*', default=glob.glob('test/random-*.in'))
    parser.add_argument('--source', default='main.cpp')
    parser.add_argument('--n-iter', type=int, default=90)
    parser.add_argument('--init-points', type=int, default=10)
    parser.add_argument('--acq', choices=[ 'ucb', 'ei', 'poi' ], default='ucb')
    parser.add_argument('--plot')
    args = parser.parse_args()

    # optimize
    from bayes_opt import BayesianOptimization  # https://pypi.org/project/bayesian-optimization/
    param_bounds = {
        'LOG_BOLTZMANN': (-5.0, 5.0),
    }
    kwargs = {
        'init_points': args.init_points,
        'n_iter': args.n_iter,
        'acq': args.acq,
        'kappa': 1.5,
    }
    with tempfile.TemporaryDirectory() as tempdir:
        execute = Executer(tempdir=tempdir, source=args.source, testcases=args.testcase)
        optimizer = BayesianOptimization(f=execute, pbounds=param_bounds)
        optimizer.maximize(**kwargs)

    # output
    print(json.dumps(optimizer.res, indent=4))
    if args.plot:
        plot(optimizer, 'LOG_BOLTZMANN', args.plot)

if __name__ == '__main__':
    main()