import argparse
from collections import defaultdict
from datetime import datetime
import math

import sqlalchemy
import optuna

profile_result = defaultdict(lambda: datetime.now() - datetime.now())


class Profile:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.start = datetime.now()

    def __exit__(self, exc_type, exc_val, exc_tb):
        profile_result[self.name] += datetime.now() - self.start


def print_profile():
    for key, value in sorted(profile_result.items(), key=lambda i: i[1],
                             reverse=True):
        print(key.ljust(50) + '{}'.format(value))


def build_objective_fun(n_param):
    def objective(trial):
        return sum([
            math.sin(trial.suggest_float('param-{}'.format(i), 0, math.pi * 2))
            for i in range(n_param)
        ])

    return objective


def define_flags(parser):
    parser.add_argument('mysql_user', type=str)
    parser.add_argument('mysql_password', type=str)
    parser.add_argument('mysql_host', type=str)
    parser.add_argument('n_study', type=int)
    parser.add_argument('n_trial', type=int)
    parser.add_argument('n_param', type=int)
    return parser


if __name__ == '__main__':
    parser = define_flags(argparse.ArgumentParser())
    args = parser.parse_args()
    print(vars(args))

    storage_str = 'mysql+pymysql://{}:{}@{}/'.format(
        args.mysql_user,
        args.mysql_password,
        args.mysql_host,
    )
    engine = sqlalchemy.create_engine(storage_str)
    sampler = optuna.samplers.TPESampler()
    conn = engine.connect()
    conn.execute("commit")
    database_str = 'profile_storage_s{}_t{}_p{}'.format(
        args.n_study, args.n_trial, args.n_param)

    try:
        conn.execute(
            "drop database {}".format(database_str))
    except:
        pass
    conn.execute("create database {}".format(
        database_str
    ))
    conn.close()

    for i in range(args.n_study):
        storage = optuna.storages.get_storage(storage_str + database_str)
        study = optuna.create_study(sampler=sampler, storage=storage)
        study_id = study.study_id
        with Profile('perf:optimize-study'):
            study.optimize(build_objective_fun(args.n_param), n_trials=args.n_trial, gc_after_trial=False)

    print_profile()