"""
Example 2: Restart from previous state
1 - dimensional

same optimization as example 0
What you will learn:
* saving and retrieving data from previous run
"""
import logging
import time
import traceback
from typing import List

import numpy as np
import examples.examples_all_functions as exf
from PARyOpt import BayesOpt
from PARyOpt.evaluators import FunctionEvaluator


# custom initialization
def my_init_strategy(n_dim: int, n_init: int, l_bound: np.array, u_bound: np.array) -> List[np.array]:
    """
    User customized initialization -- since the example is one-dimensional
    we shall uniformly divide the domain into n_init parts and then return the values
    :param n_dim: dimensionality
    :param n_init: number of initial points
    :param l_bound: lower bound
    :param u_bound: upper bound
    :return:
    """
    if n_dim == 1:
        lin_spaced_intervals = np.linspace(l_bound, u_bound, n_init+2)
        return [np.asarray([lin_spaced_intervals[i+1]]) for i in range(n_init)]
    else:
        raise NotImplementedError("init strategy not implemented for more than 1 dimension")


if __name__ == "__main__":
    # logging setup -- if this is not done, it will be streamed to stdout
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)  # either NOTSET, INFO, DEBUG, WARNING, ERROR, CRITICAL -- different levels of log
    log_file_name = 'example2_{}.log'.format(time.strftime("%Y.%m.%d-%H%M%S"))
    fh = logging.FileHandler(log_file_name, mode='a')

    # log file format
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    fh.setFormatter(formatter)
    logger.addHandler(fh)

    # first line of the log
    logger.info('Example run for Bayesian optimization')

    # let us define basic parameters here

    # dimensionality of problem
    n_dim = 1

    # bounds, as a numpy array
    # l_bound[0] is the lower bound for dimension 0, etc.
    l_bound = np.asarray([-12.])
    u_bound = np.asarray([12.])

    # number of cost function calls to create the initial surrogate
    n_init = 2

    # maximum number of iterations
    max_iter = 10

    # cost function : simple parabola
    cost_function = exf.parabolic_cost_function
    # how to evaluate the cost function - we will use the serial, in-script evaluator
    evaluator = FunctionEvaluator(cost_function)

    # kappa strategy : defines exploration vs exploitation .. we will use a pure exploration here
    # large values of kappa are exploratory (prioritize areas where confidence is low),
    # small values are exploitatory (prioritize areas where the surrogate seems optimal)
    # the following definition returns a value of 1000 for any iteration value
    def my_kappa(iteration: int) -> float:
        """
        return a constant kappa value of 1000.
        :param iteration: iteration
        :return:
        """
        return 1000.0


    # instantiate the optimizer with cost function, bounds, initial evaluation points,
    # type of kernel function (squared exponential), acquisition function (lower confidence bound)
    # and a kappa strategy (constant here)

    # A sample run will be performed which will get the optimizer state.
    b_opt = BayesOpt(cost_function=evaluator,
                     l_bound=l_bound, u_bound=u_bound, n_dim=n_dim,
                     n_init=2, init_strategy=my_init_strategy,
                     kern_function='sqr_exp',
                     acq_func='LCB',
                     kappa_strategy=my_kappa,
                     if_restart=False)

    logger.info('BO initialized')

    # visualize the surrogate generated by the optimizer, with no hyper-parameter optimization
    exf.visualize_fit(b_opt)
    # estimate best kernel parameters: perform hyper parameter optimization
    b_opt.estimate_best_kernel_parameters(theta_bounds=[[0.1, 10]])
    exf.visualize_fit(b_opt)

    # update iterations
    for curr_iter in range(max_iter):
        # update_iter finds acquisition function optima and updates the prior to get the posterior
        # (optimize acquisition function to pick the next set of points,
        # evaluate the cost function at those points, and update the surrogate with the new values)
        b_opt.update_iter()
        # estimate and set best kernel parameters
        b_opt.estimate_best_kernel_parameters(theta_bounds=[[0.1, 10]])
        # visualize it every other iteration
        if curr_iter % 3:
            exf.visualize_fit(b_opt)

    # get current best evaluated value
    best_location, best_value = b_opt.get_current_best()

    result_txt = 'Optimization done for {} iterations, best evaluation is at {} with cost: {}'. \
        format(b_opt.get_current_iteration(), best_location, best_value)

    logger.info(result_txt)
    print(result_txt)

    # restart the optimization using the `opt_state.dat` file from the previous instance
    restarted_bo = BayesOpt(cost_function=evaluator,
                            l_bound=l_bound, u_bound=u_bound, n_dim=n_dim,
                            n_init=2,
                            kern_function='sqr_exp',
                            acq_func='LCB',
                            kappa_strategy=my_kappa,
                            if_restart=True, restart_filename='opt_state.dat')

    logger.info('Restarted optimization started, iterations completed: {}'.format(restarted_bo.get_current_iteration()))
    exf.visualize_fit(restarted_bo)

    # continue the iterations for another 5 iterations
    # note the new syntax for simultaneous update without a loop
    restarted_bo.update_iter(5)
    exf.visualize_fit(restarted_bo)
    restarted_bo.estimate_best_kernel_parameters(theta_bounds=[[0.005, 10]])
    exf.visualize_fit(restarted_bo)

    # get current best evaluated value
    best_location, best_value = restarted_bo.get_current_best()
    result_txt = 'Optimization done for {} iterations in total, best evaluation is at {} with cost: {}'. \
        format(restarted_bo.get_current_iteration(), best_location, best_value)
    logger.info(result_txt)
    print(result_txt)
