import json
import numpy as np
import matplotlib.pyplot as plt


METHODS = ["proc3s", "deneck"]
OPTIMS = ["cma", "hill_climb"]
TASKS = ["push_line", "push_circle"]
COUNT_PER_CONFIG = 1
colors = ['tab:blue', 'tab:orange', 'tab:green']

def mean_and_std(of_method: str):

    cost_histories = []

    for i in range(COUNT_PER_CONFIG):
        cost_history = json.load(open(f"results/{of_method}/{i}.json"))["step_0_cost_history"]
        fixed_cost_history = [min(cost_history[:i+1]) for i in range(len(cost_history))]
        cost_histories.append(fixed_cost_history)

    cost_histories = np.array(cost_histories)
    mean = np.mean(cost_histories, axis=0)
    std = np.std(cost_histories, axis=0)

    return mean, std

fig, axs = plt.subplots(1, 2, figsize=(12, 5), sharey=True)

for i, task in enumerate(TASKS):

    for j, method in enumerate(METHODS):

        if j:
            for k, optim in enumerate(OPTIMS):
                method_folder = f"{method}_{task}_{optim}"
                color = colors[1 + k]

                y, std = mean_and_std(method_folder)
                y = np.log(y)
                x = list(range(1, len(y)+1))
                label = method_folder.replace(task, "").replace("_", " ")
                axs[i].plot(x, y, label=label, color=color)
                axs[i].fill_between(x, y - std, y + std, color=color, alpha=0.3)

        else:
            method_folder = f"{method}_{task}"
            color = colors[0]

            y, std = mean_and_std(method_folder)
            y = np.log(y)
            x = list(range(1, len(y)+1))
            label = method_folder.replace(task, "").replace("_", " ")
            axs[i].plot(x, y, label=label, color=color)
            axs[i].fill_between(x, y - std, y + std, color=color, alpha=0.3)

    axs[i].legend()
    axs[i].set_xlabel('Iteration')
    axs[i].set_title(task.replace("_", " "))

axs[0].set_ylabel('Lowest Cost')
fig.suptitle('BBO Methods on the Pushing Domain')

plt.tight_layout()
plt.savefig('results_pushing.png')
