import math
import numpy as np
import pandas as pd
from smac.configspace import ConfigurationSpace
from ConfigSpace.hyperparameters import UniformFloatHyperparameter

def branin_function(configs, reformat=True):
    
    """
    Compute the branin function.
    :param X: dictionary containing the input points.
    :return: the value of the branin function
    """
    
    if type(configs) != np.ndarray:
        config = configs.get_dictionary()
        X = np.array([[config[param] for param in config.keys()]]).reshape(1, -1)
        X = X.T
        x1, x2 = X[0] * 15 - 5, 2.275

    else:
        X = configs
        if reformat:
            x1, x2 = X[:,0] * 15 - 5, 2.275#X[:,1]
        else:
            x1, x2 = X[0] * 15 - 5, 2.275#X[1]

    a = 1.0
    b = 5.1 / (4.0 * math.pi * math.pi)
    c = 5.0 / math.pi
    r = 6.0
    s = 10.0
    t = 1.0 / (8.0 * math.pi)

    y_value = a * (x2 - b * x1 * x1 + c * x1 - r) ** 2 + s * (1 - t) * np.cos(x1) + s
    #print(f'BRANIN PRINT: X: {X} ----- Y: {y_value}' )
    print(f'Config {configs} ------ value of {y_value}')
    # restructure to get the same return value structure as the other funtions
    return y_value.reshape(-1,1), np.array(0)


def setup_branin(fcn=branin_function):
    
    class BraninSetup:
        def __init__(self):
            self.range = [(0, 1)]
        
        def get_bounds(self):
            return self.range

    bs = BraninSetup()
    return branin_function, bs
