"""
---
    Copyright (c) 2018 Baskar Ganapathysubramanian, Balaji Sesha Sarath Pokuri

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
---
"""
"""
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 = '../logs/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)
