import math
import numpy as np
import pandas as pd

def branin_function(X, reformat=True):
    """
    Compute the branin function.
    :param X: dictionary containing the input points.
    :return: the value of the branin function
    """
    if reformat:
        x1, x2 = X[:,0], X[:,1]
    else:
        x1, x2 = X[0], 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
    
    # 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 = [(-5, 10), (0, 15)]
        
        def get_bounds(self):
            return self.range

    bs = BraninSetup()
    return branin_function, bs
