import pickle
import colorsys
import matplotlib.colors as mc


def data2pickle(data, name):
    with open(name, 'wb') as handle:
        pickle.dump(data, handle, protocol=pickle.HIGHEST_PROTOCOL)


def pickle2data(name):
    with open(name, 'rb') as handle:
        data = pickle.load(handle, )
    return data


def set_size(width_pt, fraction=1, subplots=(1, 1)):
    """Set figure dimensions to sit nicely in our document.

    Parameters
    ----------
    width_pt: float
            Document width in points
    fraction: float, optional
            Fraction of the width which you wish the figure to occupy
    subplots: array-like, optional
            The number of rows and columns of subplots.
    Returns
    -------
    fig_dim: tuple
            Dimensions of figure in inches
    """
    # Width of figure (in pts)
    fig_width_pt = width_pt * fraction
    # Convert from pt to inches
    inches_per_pt = 1 / 72.27

    # Golden ratio to set aesthetic figure height
    golden_ratio = (5 ** .5 - 1) / 2

    # Figure width in inches
    fig_width_in = fig_width_pt * inches_per_pt
    # Figure height in inches
    fig_height_in = fig_width_in * golden_ratio * (subplots[0] / subplots[1])

    return (fig_width_in, fig_height_in)


def lighten_color(color, amount=0.5):
    """
    Lightens the given color by multiplying (1-luminosity) by the given amount.
    Input can be matplotlib color string, hex string, or RGB tuple.

    Examples:
    >> lighten_color('g', 0.3)
    >> lighten_color('#F034A3', 0.6)
    >> lighten_color((.3,.55,.1), 0.5)
    """
    try:
        c = mc.cnames[color]
    except:
        c = color
    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
    return colorsys.hls_to_rgb(c[0], 1 - amount * (1 - c[1]), c[2])


def visualize_results(results_name, forgetting_factors, time_horizon, trials, misspecified_eps=0, DIM=2):
    import numpy as np
    import matplotlib.pyplot as plt
    from utils.preprocessing_utils import get_environment_path

    colors = ['#000000', '#0C5DA5', '#00B945', '#845B97', '#FF2C00']

    N_TRAINING_POINTS = 0
    MODEL_ID = 0
    TIME = np.arange(N_TRAINING_POINTS + 1, time_horizon + 1)

    path = get_environment_path()

    def calculate_cum_regret(opt, obs, time, normalize=True):
        if len(opt.shape) > 1:
            opt = opt.reshape(-1)
        if len(obs.shape) > 1:
            obs = obs.reshape(-1)

        out = np.cumsum(opt - obs)
        if normalize:
            out = out / time
        return out

    if results_name == "synthetic_data_constant_eps":

        from objective_functions.objective_functions import within_model_wrapper_2D_from_file

        plt.rcParams.update({
            "font.family": "serif",     # use serif/main font for text elements
            "text.usetex": True,        # use inline math for ticks
            "pgf.rcfonts": False,       # don't setup fonts from rc parameters
            'text.latex.preamble': [r'\usepackage{amsmath}',
                                    r'\usepackage{amssymb}']
        })

        plot0_path = f'results_obj_2D_{forgetting_factors[0]}'
        plot1_path = f'results_obj_2D_{forgetting_factors[1]}'
        plot2_path = f'results_obj_2D_{forgetting_factors[2]}'
        results_paths = [plot0_path, plot1_path, plot2_path]

        obj_func1 = 'obj_2D'
        obj_paths = [obj_func1, obj_func1, obj_func1]

        LABELS = ["GP-UCB",
                  "R-GP-UCB",
                  "TV-GP-UCB (reference)",
                  "UI-TVBO",
                  "ET-GP-UCB (ours)"]

        YBOUNDS = [0, 1.8]

        TEXTS = [
            r"$\varepsilon=0.01$",
            r"$\varepsilon=0.03$",
            r"$\varepsilon=0.05$",
        ]

        fig_labels = [r"(a) Constant $\varepsilon$",
                      r"(b) Constant $\varepsilon$",
                      r"(c) Constant $\varepsilon$",]

        x, y = set_size(600, subplots=(2, 3), fraction=1.)
        fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(x, y), gridspec_kw={'height_ratios': [3, 1],
                                                                                'width_ratios': [1, 1, 1]})

        for j, (plot_path, factor, obj_path, (ax1, ax2)) in enumerate(
                zip(results_paths, forgetting_factors, obj_paths, axes.T)):
            variations = {"GP_UCB": [],
                          "R_GP_UCB": [],
                          "TV_GP_UCB": [],
                          "UI_TVBO": [],
                          "ET_GP_UCB": [],
                          }

            for method, color in zip(variations.keys(), colors):
                for trial in range(trials):

                    # get optimum
                    name = f"objective_functions/synthetic_data/2D/{obj_path}_{trial}_{factor}".replace(".", "_")
                    obj_name = path + name + ".pickle"
                    optimum = within_model_wrapper_2D_from_file(obj_name, return_max=True)[
                              N_TRAINING_POINTS:time_horizon]

                    name = f"results/{plot_path}_{method}_{MODEL_ID}_{trial}".replace(".", "_")
                    results = pickle2data(path + name + ".pickle")
                    observations = np.asarray(results["observations"])[N_TRAINING_POINTS:]
                    cum_regret = calculate_cum_regret(optimum, observations, TIME)
                    variations[method].append(cum_regret)

                    triggered = np.asarray(results["triggered"])
                    if np.any(triggered):  # and len(triggered) < 300:
                        if method == "R_GP_UCB":
                            marker = 'o'
                        else:
                            marker = 'v'
                        t_triggered = TIME[triggered]
                        ax2.scatter(t_triggered,np.zeros_like(t_triggered) + trial,
                                    s=10,
                                    linewidths=0.2,
                                    marker=marker,
                                    c=[color for i in range(len(t_triggered))],
                                    alpha=0.25)

            for i, (method, color) in enumerate(zip(variations.keys(), colors)):
                cum_regrets = np.asarray(variations[method])
                avg_cum_regrets = np.mean(cum_regrets, axis=0)
                stdv_cum_regrets = np.sqrt(np.var(cum_regrets, axis=0))
                if not method == "TV_GP_UCB":
                    ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                     color=color, alpha=0.1, lw=0.3)
                    ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1)
                else:
                    ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                     color=color, alpha=0.1, lw=0.3)
                    ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1, linestyle=(0, (5, 1)), )

            ax1.text(200, 1.6, TEXTS[j], ha='center', va='center', fontsize=9)
            ax1.set_xlim([N_TRAINING_POINTS, time_horizon])
            ax1.set_yticks([0, 1, 2])
            ax1.set_ylim(YBOUNDS)
            ax1.set_ylabel("Norm. Regret\n$R_t / t$", )

            ax2.set_ylim([0, trials])
            ax2.set_xlim([N_TRAINING_POINTS, time_horizon])
            ax2.set_ylabel("Runs")

            ax2.set_xlabel(f"Time steps $t$\n{fig_labels[j]}", labelpad=2, linespacing=1.5)
            ax1.label_outer()
            ax2.label_outer()

        ax1.legend(bbox_to_anchor=(-1.9, 1.03,), loc="lower left", handletextpad=0.2, borderpad=0.2,
                   borderaxespad=0.1, ncol=4,  # edgecolor='inherit', fancybox=False,
                   columnspacing=0.7, prop={'size': 9}, handlelength=1.5)

        h = [plt.plot([], [], color=color, alpha=0.6, marker=marker, ms=4, ls="")[0] for color, marker in
             [[[0.25098039215686274, 0.4980392156862745, 0.7176470588235294], "o"],
              [[0.8, 0.027450980392156862, 0.11764705882352941], "v"]]]

        leg = plt.legend(handles=h, labels=["Resets of R-GP-UCB", "Resets of ET-GP-UCB", ],
                         # "different realization of $f_t(x)$"],
                         bbox_to_anchor=(-1.4, 1.02),
                         loc="lower left",
                         handletextpad=0.2,
                         borderpad=0.2,
                         borderaxespad=0.2, ncol=3,  # edgecolor='inherit', fancybox=False,
                         columnspacing=0.7, prop={'size': 9})
        ax2.add_artist(leg)

        fig.subplots_adjust(bottom=0.27, top=0.89, left=0.08, right=0.98, wspace=0.13, hspace=0.45)
        plt.show()
        plt.close()


    #####################################################################
    ########################  Higher Dimensional  #######################
    #####################################################################
    if results_name == "synthetic_data_constant_eps_ND":

        from objective_functions.objective_functions import within_model_wrapper_2D_from_file

        plt.rcParams.update({
            "font.family": "serif",     # use serif/main font for text elements
            "text.usetex": True,        # use inline math for ticks
            "pgf.rcfonts": False,       # don't setup fonts from rc parameters
            'text.latex.preamble': [r'\usepackage{amsmath}',
                                    r'\usepackage{amssymb}']
        })

        plot0_path = f'results_obj_2{DIM}_{forgetting_factors[0]}'
        plot1_path = f'results_obj_2{DIM}_{forgetting_factors[1]}'
        plot2_path = f'results_obj_2{DIM}_{forgetting_factors[2]}'
        results_paths = [plot0_path, plot1_path, plot2_path]

        obj_func1 = 'obj_2D'
        obj_paths = [obj_func1, obj_func1, obj_func1]

        LABELS = ["GP-UCB",
                  "R-GP-UCB",
                  "TV-GP-UCB",
                  "ET-GP-UCB",
                  "ET-GP-UCB w. backtracking"]

        TEXTS = [
            r"$\varepsilon=0.01$",
            r"$\varepsilon=0.03$",
            r"$\varepsilon=0.05$",
        ]

        fig_labels = [r"(a) Constant $\varepsilon$",
                      r"(b) Constant $\varepsilon$",
                      r"(c) Constant $\varepsilon$",]

        x, y = set_size(600, subplots=(2, 3), fraction=1.)
        fig, axes = plt.subplots(nrows=2, ncols=3, figsize=(x, y), gridspec_kw={'height_ratios': [3, 1],
                                                                                'width_ratios': [1, 1, 1]})

        for j, (plot_path, factor, obj_path, (ax1, ax2)) in enumerate(
                zip(results_paths, forgetting_factors, obj_paths, axes.T)):
            variations = {"GP_UCB": [],
                          "R_GP_UCB": [],
                          "TV_GP_UCB": [],
                          "ET_GP_UCB": [],
                          "ET_GP_UCB_backtracking": [],
                          }

            for method, color in zip(variations.keys(), colors):
                for trial in range(trials):

                    # get optimum
                    name = f"objective_functions/synthetic_data/2D/{obj_path}_{trial}_{factor}".replace(".", "_")
                    obj_name = path + name + ".pickle"
                    optimum = within_model_wrapper_2D_from_file(obj_name, return_max=True)[
                              N_TRAINING_POINTS:time_horizon]

                    name = f"results/{plot_path}_{method}_{MODEL_ID}_{trial}".replace(".", "_")
                    results = pickle2data(path + name + ".pickle")
                    observations = np.asarray(results["observations"])[N_TRAINING_POINTS:]
                    cum_value = calculate_cum_regret(observations, np.zeros_like(observations), TIME)
                    variations[method].append(cum_value)

                    triggered = np.asarray(results["triggered"])
                    if np.any(triggered):  # and len(triggered) < 300:
                        if method == "R_GP_UCB":
                            marker = 'o'
                        else:
                            marker = 'v'
                        t_triggered = TIME[triggered]
                        ax2.scatter(t_triggered,np.zeros_like(t_triggered) + trial,
                                    s=10,
                                    linewidths=0.2,
                                    marker=marker,
                                    c=[color for i in range(len(t_triggered))],
                                    alpha=0.25)

            for i, (method, color) in enumerate(zip(variations.keys(), colors)):
                cum_regrets = np.asarray(variations[method])
                avg_cum_regrets = np.mean(cum_regrets, axis=0)
                stdv_cum_regrets = np.sqrt(np.var(cum_regrets, axis=0))
                if not method == "TV_GP_UCB":
                    ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                     color=color, alpha=0.1, lw=0.3)
                    ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1)
                else:
                    ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                     color=color, alpha=0.1, lw=0.3)
                    ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1, linestyle=(0, (5, 1)), )

            ax1.text(200, 1.6, TEXTS[j], ha='center', va='center', fontsize=9)
            ax1.set_xlim([N_TRAINING_POINTS, time_horizon])
            ax1.set_ylabel("Norm. cumulative value\n$f_t (x) / t$", )

            ax2.set_ylim([0, trials])
            ax2.set_xlim([N_TRAINING_POINTS, time_horizon])
            ax2.set_ylabel("Runs")

            ax2.set_xlabel(f"Time steps $t$\n{fig_labels[j]}", labelpad=2, linespacing=1.5)
            ax1.label_outer()
            ax2.label_outer()

        ax1.legend(bbox_to_anchor=(-1.9, 1.03,), loc="lower left", handletextpad=0.2, borderpad=0.2,
                   borderaxespad=0.1, ncol=4,  # edgecolor='inherit', fancybox=False,
                   columnspacing=0.7, prop={'size': 9}, handlelength=1.5)

        h = [plt.plot([], [], color=color, alpha=0.6, marker=marker, ms=4, ls="")[0] for color, marker in
             [[[0.25098039215686274, 0.4980392156862745, 0.7176470588235294], "o"],
              [[0.8, 0.027450980392156862, 0.11764705882352941], "v"]]]

        leg = plt.legend(handles=h, labels=["Resets of R-GP-UCB", "Resets of ET-GP-UCB", ],
                         # "different realization of $f_t(x)$"],
                         bbox_to_anchor=(-1.4, 1.02),
                         loc="lower left",
                         handletextpad=0.2,
                         borderpad=0.2,
                         borderaxespad=0.2, ncol=3,  # edgecolor='inherit', fancybox=False,
                         columnspacing=0.7, prop={'size': 9})
        ax2.add_artist(leg)

        fig.subplots_adjust(bottom=0.27, top=0.89, left=0.08, right=0.98, wspace=0.13, hspace=0.45)
        plt.show()
        plt.close()

    #####################################################################
    ######################## MISSPECIFIED EPSILON #######################
    #####################################################################
    if results_name == "synthetic_data_misspecified_eps":

        from objective_functions.objective_functions import within_model_wrapper_2D_from_file

        plt.rcParams.update({
            "font.family": "serif",  # use serif/main font for text elements
            "text.usetex": True,  # use inline math for ticks
            "pgf.rcfonts": False,  # don't setup fonts from rc parameters
            'text.latex.preamble': [r'\usepackage{amsmath}',
                                    r'\usepackage{amssymb}']
        })

        if isinstance(forgetting_factors, list):
            forgetting_factor = forgetting_factors[0]
        else:
            forgetting_factor = forgetting_factors

        LABELS = ["GP-UCB",
                  "R-GP-UCB",
                  "TV-GP-UCB",
                  "UI-TVBO",
                  "ET-GP-UCB (ours)"]

        obj_path = f'obj_2D'

        YBOUNDS = [0, 1.8]

        text = "$\\varepsilon_{true}=0.05$ \n $\\varepsilon_{missp.}=0.001$"
        fig_label = r"(d) Misspecified $\varepsilon$"

        x, y = set_size(300, subplots=(2, 1), fraction=1.)
        fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(x, y), gridspec_kw={'height_ratios': [3, 1],})

        variations = {"GP_UCB": [],
                      "R_GP_UCB": [],
                      "TV_GP_UCB": [],
                      "UI_TVBO": [],
                      "ET_GP_UCB": [],
                      }

        for method, color in zip(variations.keys(), colors):
            for trial in range(trials):

                # get optimum
                name = f"objective_functions/synthetic_data/2D/{obj_path}_{trial}_{forgetting_factor}".replace(".", "_")
                obj_name = path + name + ".pickle"
                optimum = within_model_wrapper_2D_from_file(obj_name, return_max=True)[
                          N_TRAINING_POINTS:time_horizon]

                name = f"results/results_obj_2D_misspecified_eps_{forgetting_factor}_{method}_{MODEL_ID}_{trial}".replace(".", "_")
                results = pickle2data(path + name + ".pickle")
                observations = np.asarray(results["observations"])[N_TRAINING_POINTS:]
                cum_regret = calculate_cum_regret(optimum, observations, TIME)
                variations[method].append(cum_regret)

                triggered = np.asarray(results["triggered"])
                if np.any(triggered):
                    if method == "R_GP_UCB":
                        marker = 'o'
                    else:
                        marker = 'v'
                    t_triggered = TIME[triggered]
                    ax2.scatter(t_triggered, np.zeros_like(t_triggered) + trial,
                                s=10,
                                linewidths=0.2,
                                marker=marker,
                                c=[color for i in range(len(t_triggered))],
                                alpha=0.25)

        for i, (method, color) in enumerate(zip(variations.keys(), colors)):
            cum_regrets = np.asarray(variations[method])
            avg_cum_regrets = np.mean(cum_regrets, axis=0)
            stdv_cum_regrets = np.sqrt(np.var(cum_regrets, axis=0))
            if not method == "TV_GP_UCB":
                ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                 color=color, alpha=0.1, lw=0.3)
                ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1)
            else:
                ax1.fill_between(TIME, avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                                 color=color, alpha=0.1, lw=0.3)
                ax1.plot(TIME, avg_cum_regrets, label=LABELS[i], color=color, linewidth=1, linestyle=(0, (5, 1)), )

        ax1.text(200, 0.1, text, ha='center', va='bottom', fontsize=7)
        ax1.set_xlim([N_TRAINING_POINTS, time_horizon])
        ax1.set_yticks([0, 1, 2])
        ax1.set_ylim(YBOUNDS)
        ax1.set_ylabel("Norm. Regret\n$R_t / t$", )

        ax2.set_ylim([0, trials])
        ax2.set_xlim([N_TRAINING_POINTS, time_horizon])
        ax2.set_ylabel("Runs")

        ax2.set_xlabel(f"Time steps $t$\n{fig_label}", labelpad=2, linespacing=1.5)
        ax1.label_outer()
        ax2.label_outer()

        ax1.legend(bbox_to_anchor=(0, 1.03,), loc="lower left", handletextpad=0.2, borderpad=0.2,
                   borderaxespad=0.1, ncol=4,  # edgecolor='inherit', fancybox=False,
                   columnspacing=0.7, prop={'size': 9}, handlelength=1.5)

        h = [plt.plot([], [], color=color, alpha=0.6, marker=marker, ms=4, ls="")[0] for color, marker in
             [['#0C5DA5', "o"], ['#FF2C00', "v"]]]

        leg = plt.legend(handles=h, labels=["Resets of R-GP-UCB", "Resets of ET-GP-UCB", ],
                         # "different realization of $f_t(x)$"],
                         bbox_to_anchor=(0, 1.02),
                         loc="lower left",
                         handletextpad=0.2,
                         borderpad=0.2,
                         borderaxespad=0.2, ncol=3,  # edgecolor='inherit', fancybox=False,
                         columnspacing=0.7, prop={'size': 9})
        ax2.add_artist(leg)

        fig.subplots_adjust(bottom=0.27, top=0.89, left=0.08, right=0.98, wspace=0.13, hspace=0.45)
        plt.show()
        plt.close()

    #####################################################################
    ######################## Temperature Data  ##########################
    #####################################################################
    if "temperature" in results_name:

        from objective_functions.objective_functions import wrapper_temperature_data
        from utils.preprocessing_utils import get_environment_path, preprocess_temperature_data

        plt.rcParams.update({
            "font.family": "serif",  # use serif/main font for text elements
            "text.usetex": True,  # use inline math for ticks
            "pgf.rcfonts": False,  # don't setup fonts from rc parameters
            'text.latex.preamble': [r'\usepackage{amsmath}',
                                    r'\usepackage{amssymb}']
        })

        if isinstance(forgetting_factors, list):
            forgetting_factor = forgetting_factors[0]
        else:
            forgetting_factor = forgetting_factors

        colors = ['#000000', '#0C5DA5', '#00B945', '#FF9500', '#845B97', '#FF2C00']

        HOURS_VEC = np.linspace(0, 48 - 1 / 6, time_horizon).reshape(-1, 1)
        LABELS = ["GP-UCB",
                  "R-GP-UCB",
                  "TV-GP-UCB",
                  "TV-GP-UCB (MLE)",
                  "UI-TVBO",
                  "ET-GP-UCB (ours)"]

        train_days = [3, 7]
        if results_name == "temperature_data_on_7and8":
            test_days = [7, 9]  # last days as test days
            _, unseen_data, sensor_coords, _, stats = preprocess_temperature_data(train_days=train_days,
                                                                                  test_days=test_days,
                                                                                  verbose=False)

        if results_name == "temperature_data_on_5and6":
            test_days = [5, 7]  # last days as test days
            _, unseen_data, sensor_coords, _, stats = preprocess_temperature_data(train_days=train_days,
                                                                                  test_days=test_days,
                                                                                  verbose=False)
        mean_temp = stats["mean"]
        stdv_temp = stats["stdv"]
        max_temperature, max_temperature_id = wrapper_temperature_data(
            unseen_data, sensor_coords, return_max=True, time_horizon=time_horizon)

        # rescale max
        optimum = max_temperature * stdv_temp + mean_temp

        variations = {"GP_UCB": [],
                      "Cyclic_Trigger": [],
                      "TV_GP_UCB": [],
                      "TV_GP_UCB_MLE": [],
                      "UI_TVBO": [],
                      "ET_GP_UCB": [],
                      }

        x, y = set_size(300, subplots=(2, 1), fraction=1.)
        fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(x, y), gridspec_kw={'height_ratios': [3, 1], })

        for method, color in zip(variations.keys(), colors):
            for trial in range(trials):
                name = f"results/results_temperature_{test_days[0]}{test_days[1]}_{forgetting_factor}_{method}_{trial}".replace(
                    ".", "_")
                results = pickle2data(path + name + ".pickle")
                observations = np.asarray(results["observations"])[N_TRAINING_POINTS:time_horizon]
                cum_regret = calculate_cum_regret(optimum, observations, time=TIME)
                variations[method].append(cum_regret)

                triggered = np.asarray(results["triggered"])[N_TRAINING_POINTS:time_horizon]
                if np.any(triggered):  # and len(triggered) < 300:
                    if method == "Cyclic_Trigger":
                        marker = 'o'
                    else:
                        marker = 'v'
                    t_triggered = HOURS_VEC[triggered].reshape(-1)
                    ax2.scatter(t_triggered,
                                np.zeros_like(t_triggered) + trial,  # * 1 / trialS,
                                s=10,
                                linewidths=0.2,
                                marker=marker,
                                c=[color for i in range(len(t_triggered))], alpha=0.25)

        for method, color, label in zip(variations.keys(), colors, LABELS):
            cum_regrets = np.asarray(variations[method])
            avg_cum_regrets = np.mean(cum_regrets, axis=0)
            stdv_cum_regrets = np.sqrt(np.var(cum_regrets, axis=0))
            ax1.fill_between(HOURS_VEC.flatten(), avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                             color=color, alpha=0.1, lw=0.3)
            ax1.plot(HOURS_VEC.flatten(), avg_cum_regrets, label=label, color=color, linewidth=1)

        ax1.set_xlim([HOURS_VEC[0], 48])
        ax1.set_ylabel("Norm. Regret\n $R_t / t$", )  # labelpad=8)
        ax1.set_xticks([0, 12, 24, 36, 48])
        ax1.set_yticks([1, 2, 3])
        ax1.set_ylim([0, 3.5])
        ax1.set_xticklabels([])

        # ax2.set_yticks(range(1))
        ax2.set_ylim([0, trials])
        ax2.set_xlim([HOURS_VEC[0], 48])
        ax2.set_xticks([0, 12, 24, 36, 48])
        ax2.set_ylabel("Runs")

        if results_name == "temperature_data_on_7and8":
            text = "Time $t$ [h]\n(b) Days 7 and 8"
        if results_name == "temperature_data_on_5and6":
            text = "Time $t$ [h]\n(d) Days 5 and 6"

        ax2.set_xlabel(f"{text}", labelpad=2, linespacing=1.5)

        ax1.legend(bbox_to_anchor=(0, 1.03,), loc="lower left", handletextpad=0.2, borderpad=0.2,
                   borderaxespad=0.1, ncol=4,  # edgecolor='inherit', fancybox=False,
                   columnspacing=0.7, prop={'size': 9}, handlelength=1.5)

        h = [plt.plot([], [], color=color, alpha=0.6, marker=marker, ms=4, ls="")[0] for color, marker in
             [['#0C5DA5', "o"], ['#FF2C00', "v"]]]

        leg = plt.legend(handles=h, labels=["Resets of R-GP-UCB", "Resets of ET-GP-UCB", ],
                         # "different realization of $f_t(x)$"],
                         bbox_to_anchor=(0, 1.02),
                         loc="lower left",
                         handletextpad=0.2,
                         borderpad=0.2,
                         borderaxespad=0.2, ncol=3,  # edgecolor='inherit', fancybox=False,
                         columnspacing=0.7, prop={'size': 9})
        ax2.add_artist(leg)

        fig.subplots_adjust(bottom=0.27, top=0.89, left=0.08, right=0.98, wspace=0.13, hspace=0.45)
        plt.show()
        plt.close()


    #####################################################################
    ######################## Policy Search ##############################
    #####################################################################
    if results_name == "application_policy_search":

        from objective_functions.applications.controller_tuning import get_params, get_opt_state_controller, \
            get_linearized_model, perform_simulation

        plt.rcParams.update({
            "font.family": "serif",  # use serif/main font for text elements
            "text.usetex": True,  # use inline math for ticks
            "pgf.rcfonts": False,  # don't setup fonts from rc parameters
            'text.latex.preamble': [r'\usepackage{amsmath}',
                                    r'\usepackage{amssymb}']
        })

        if isinstance(forgetting_factors, list):
            forgetting_factor = forgetting_factors[0]
        else:
            forgetting_factor = forgetting_factors

        sample_time = 0.02
        Q = np.eye(4) * 10
        R = np.eye(1)

        mean = -16.0493335
        stdv = 1.25

        Ks = []
        costs = []
        for ti in TIME:
            test_params = get_params(ti)
            model = get_linearized_model(test_params, sample_time)
            K = get_opt_state_controller(model, Q, R)
            Ks.append(K)
            lqr_cost = perform_simulation(model, K, ti, False)
            costs.append(lqr_cost)

        optimum = np.asarray(costs) * -1

        variations = {"GP_UCB": [],
                      "R_GP_UCB": [],
                      "TV_GP_UCB": [],
                      "TV_GP_UCB_MLE": [],
                      "UI_TVBO": [],
                      "ET_GP_UCB": [],
                      }

        LABELS = ["GP-UCB",
                  "R-GP-UCB",
                  "TV-GP-UCB",
                  "TV-GP-UCB (MLE)",
                  "UI-TVBO",
                  "ET-GP-UCB (ours)"]

        colors = ['#000000', '#0C5DA5', '#00B945', '#FF9500', '#845B97', '#FF2C00']

        x, y = set_size(270, subplots=(2, 1), fraction=1.)
        fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, figsize=(x * 1.2, y * 1),
                                       gridspec_kw={'height_ratios': [4, 1]})

        for method, color in zip(variations.keys(), colors):

            for TRIAL in range(trials):
                name = f"results/results_LQR_{forgetting_factor}_{method}_{TRIAL}".replace(".", "_")
                results = pickle2data(path + name + ".pickle")
                observations = np.asarray(results["observations"])[:time_horizon] * stdv + mean
                triggered = np.asarray(results["triggered"])[:time_horizon]
                if observations.shape[0] > TIME.shape[0]:
                    observations = observations[N_TRAINING_POINTS:]
                    triggered = triggered[N_TRAINING_POINTS:]

                cum_regret = calculate_cum_regret(optimum, observations, TIME)
                variations[method].append(cum_regret)

                if np.any(triggered):  # and len(triggered) < 300:
                    if method == "R_GP_UCB":
                        marker = 'o'
                    else:
                        marker = 'v'
                    t_triggered = TIME[triggered]
                    ax2.scatter(t_triggered,
                                np.zeros_like(t_triggered) + TRIAL,  # * 1 / TRIALS,
                                s=10,
                                linewidths=0.2,
                                marker=marker,
                                c=[color for i in range(len(t_triggered))], alpha=0.25)

        for method, color, label in zip(variations.keys(), colors, LABELS):
            cum_regrets = np.asarray(variations[method])
            avg_cum_regrets = np.mean(cum_regrets, axis=0)
            stdv_cum_regrets = np.sqrt(np.var(cum_regrets, axis=0))
            ax1.fill_between(TIME.flatten(), avg_cum_regrets - stdv_cum_regrets, avg_cum_regrets + stdv_cum_regrets,
                             color=color, alpha=0.1, lw=0.3)
            ax1.plot(TIME.flatten(), avg_cum_regrets, label=label, color=color, linewidth=1)

        ax1.set_xlim([TIME[0], TIME[-1]])
        ax1.set_ylim([0, 1.5])
        ax1.set_yticks([0, 0.5, 1])
        ax1.set_ylabel("Norm. Regret $R_t / t$")

        ax2.set_ylim([0, trials])
        ax2.set_xlim([TIME[0], TIME[-1]])
        ax2.set_ylabel("Runs")
        ax2.set_xlabel("Time steps $t$")

        ax1.legend(bbox_to_anchor=(0, 1.03,), loc="lower left", handletextpad=0.2, borderpad=0.2,
                   borderaxespad=0.1, ncol=4,  # edgecolor='inherit', fancybox=False,
                   columnspacing=0.7, prop={'size': 9}, handlelength=1.5)

        h = [plt.plot([], [], color=color, alpha=0.6, marker=marker, ms=4, ls="")[0] for color, marker in
             [['#0C5DA5', "o"], ['#FF2C00', "v"]]]

        leg = plt.legend(handles=h, labels=["Resets of R-GP-UCB", "Resets of ET-GP-UCB", ],
                         # "different realization of $f_t(x)$"],
                         bbox_to_anchor=(0, 1.02),
                         loc="lower left",
                         handletextpad=0.2,
                         borderpad=0.2,
                         borderaxespad=0.2, ncol=3,  # edgecolor='inherit', fancybox=False,
                         columnspacing=0.7, prop={'size': 9})
        ax2.add_artist(leg)

        ax1.label_outer()
        ax2.label_outer()

        fig.subplots_adjust(bottom=0.15, top=0.86, left=0.15, right=0.96, hspace=0.4)

        plt.show()
        plt.close()