import numpy as np
import pickle


def get_search_space(simulator):
    if simulator == 'motorcycle':
        # The search space must be with 101 values in order for the motorcycle simulator to work..
        search_space = np.linspace(0, 100, 101)
    elif simulator == 'autompg':
        with open("state_dicts/auto_mpg-search_space.pkl", 'rb') as fp:
            giant_search_space = pickle.load(fp)
            idx = np.arange(giant_search_space.shape[0])
            np.random.shuffle(idx)
            search_space = giant_search_space[idx[:100000], :]
    elif simulator in ["gramacy1d", "gramacy1d_wrong_gp"]:
        search_space = np.linspace(0.5, 2.5, 101)
    elif simulator == "gramacy2d":
        n_points = 30
        x1 = np.linspace(-2, 6, n_points)
        x2 = np.linspace(-2, 6, n_points)
        search_space = np.array(np.meshgrid(x1, x2, indexing='ij')).T.reshape(-1, 2)
    elif simulator == "gramacy6d":
        n_points = 30
        x1 = np.linspace(0, 1, n_points)
        x2 = np.array([0.])
        giant_search_space = np.array(np.meshgrid(x1, x1, x1, x1, x2, x2, indexing='ij')).T.reshape(-1, 6)
        idx = np.arange(giant_search_space.shape[0])
        np.random.shuffle(idx)
        search_space = giant_search_space[idx[:100000], :]
    elif simulator == "higdon1d":
        search_space = np.linspace(0, 20, 101)
    elif simulator in ['mbml_homo_0', 'mbml_homo_01', 'mbml_hetero']:
        search_space = np.linspace(-0.5, 1.5, 101)
    else:
        search_space = np.linspace(0, 5, 101)

    return search_space
