try:
    from pySOT.experimental_design import SymmetricLatinHypercube, LatinHypercube, TwoFactorial
    using_pysot = True
except ImportError:
    using_pysot = False

if using_pysot:

    from poap.controller import BasicWorkerThread, ThreadController
    from pySOT.experimental_design import SymmetricLatinHypercube, LatinHypercube, TwoFactorial
    from pySOT.strategy import SRBFStrategy, EIStrategy, DYCORSStrategy,RandomStrategy, LCBStrategy
    from pySOT.surrogate import CubicKernel, LinearTail, RBFInterpolant, GPRegressor
    from pySOT.optimization_problems.optimization_problem import OptimizationProblem
    import numpy as np
    import logging

    logging.getLogger('pySOT').setLevel(logging.ERROR)


    class GenericProblem(OptimizationProblem):

        def __init__(self, dim, objective):
            self.dim = dim
            self.min = 0
            self.minimum = np.zeros(dim)
            self.lb = 0 * np.ones(dim)
            self.ub = 1 * np.ones(dim)
            self.int_var = np.array([])
            self.cont_var = np.arange(0, dim)
            self.objective = objective
            self.info = str(dim) + "-dimensional objective function " + objective.__name__
            self.feval_count = 0

        def eval(self, x):
            """Evaluate the objective x

            :param x: Data point
            :type x: numpy.array
            :return: Value at x
            :rtype: float
            """
            self.__check_input__(x)
            d = float(self.dim)
            self.feval_count = self.feval_count+1
            return self.objective(list(x))


    def pysot_cube( objective, n_trials, n_dim, with_count=False, method=None, design=None):
        """ Minimize
        :param objective:
        :param n_trials:
        :param n_dim:
        :param with_count:
        :return:
        """
        logging.getLogger('pySOT').setLevel(logging.ERROR)

        num_threads = 1
        asynchronous = True

        max_evals = n_trials
        gp = GenericProblem(dim=n_dim, objective=objective)

        if design=='latin':
            exp_design = LatinHypercube(dim=n_dim, num_pts=2 * (n_dim + 1))
        elif design=='symmetric':
            exp_design = SymmetricLatinHypercube(dim=n_dim, num_pts=2 * (n_dim + 1))
        elif design=='factorial':
            exp_design = TwoFactorial(dim=n_dim)
        else:
            raise ValueError('design should be latin, symmetric or factorial')

        # Create a strategy and a controller
        #  SRBFStrategy, EIStrategy, DYCORSStrategy,RandomStrategy, LCBStrategy
        controller = ThreadController()
        if method.lower()=='srbf':
            surrogate = RBFInterpolant(dim=n_dim, lb=np.array([0.0] * n_dim), ub=np.array([1.0] * n_dim), kernel=CubicKernel(),
                                 tail=LinearTail(n_dim))
            controller.strategy = SRBFStrategy(
                max_evals=max_evals, opt_prob=gp, exp_design=exp_design, surrogate=surrogate, asynchronous=asynchronous
            )
        elif method.lower() == 'ei':
            surrogate = GPRegressor(dim=n_dim, lb=np.array([0.0] * n_dim), ub=np.array([1.0] * n_dim) )
            controller.strategy = EIStrategy(
                max_evals=max_evals, opt_prob=gp, exp_design=exp_design, surrogate=surrogate, asynchronous=asynchronous
            )
        elif method.lower() == 'dycors':
            surrogate = RBFInterpolant(dim=n_dim, lb=np.array([0.0] * n_dim), ub=np.array([1.0] * n_dim), kernel=CubicKernel(),
                                 tail=LinearTail(n_dim))
            controller.strategy = DYCORSStrategy(
                max_evals=max_evals, opt_prob=gp, exp_design=exp_design, surrogate=surrogate, asynchronous=asynchronous
            )
        elif method.lower() == 'lcb':
            surrogate = GPRegressor(dim=n_dim, lb=np.array([0.0] * n_dim), ub=np.array([1.0] * n_dim))
            controller.strategy = LCBStrategy(
                max_evals=max_evals, opt_prob=gp, exp_design=exp_design, surrogate=surrogate, asynchronous=asynchronous
            )
        elif method.lower() == 'random':
            controller.strategy = RandomStrategy(
                max_evals=max_evals, opt_prob=gp
            )
        else:
            raise ValueError("Didn't recognize method passed to pysot")


        # Launch the threads and give them access to the objective function
        for _ in range(num_threads):
            worker = BasicWorkerThread(controller, gp.eval)
            controller.launch_worker(worker)

        # Run the optimization strategy
        result = controller.run()
        best_x = result.params[0].tolist()
        return (result.value, best_x, gp.feval_count) if with_count else (result.value, best_x)


    # Index by strategy
    # SRBFStrategy, EIStrategy, DYCORSStrategy,RandomStrategy, LCBStrategy

    def pysot_srbf_cube( objective, n_trials, n_dim, with_count=False):
        return pysot_cube(objective=objective, n_trials=n_trials, n_dim=n_dim, with_count=with_count, method='srbf', design='symmetric' )


    def pysot_ei_cube( objective, n_trials, n_dim, with_count=False):
        return pysot_cube(objective=objective, n_trials=n_trials, n_dim=n_dim, with_count=with_count, method='ei', design='symmetric' )


    def pysot_dycors_cube( objective, n_trials, n_dim, with_count=False):
        return pysot_cube(objective=objective, n_trials=n_trials, n_dim=n_dim, with_count=with_count, method='dycors', design='symmetric' )


    def pysot_lcb_cube( objective, n_trials, n_dim, with_count=False):
        return pysot_cube(objective=objective, n_trials=n_trials, n_dim=n_dim, with_count=with_count, method='lcb', design='symmetric' )


    def pysot_random_cube( objective, n_trials, n_dim, with_count=False):
        return pysot_cube(objective=objective, n_trials=n_trials, n_dim=n_dim, with_count=with_count, method='random', design='symmetric' )


    PYSOT_OPTIMIZERS = [ pysot_ei_cube, pysot_lcb_cube, pysot_random_cube, pysot_srbf_cube, pysot_dycors_cube ]
    PYSOT_TOP_OPTIMIZERS = [ pysot_srbf_cube, pysot_dycors_cube]
else:
    PYSOT_OPTIMIZERS = []
    PYSOT_TOP_OPTIMIZERS = []


if __name__ == '__main__':
    assert using_pysot
    from humpday.objectives.classic import CLASSIC_OBJECTIVES
    for objective in CLASSIC_OBJECTIVES:
        print(' ')
        print(objective.__name__)
        for optimizer in PYSOT_OPTIMIZERS:
            print( (optimizer.__name__, optimizer(objective, n_trials=250, n_dim=6, with_count=True)))