import os
from collections import defaultdict

import numpy as np
import pandas as pd

root_path = "output/wandb_export"
version_number = "1.10"
output_path = os.path.join("output/evaluation_dataframes", f"neurips_data_v{version_number}.csv")

all_metrics = ["full_rollout_mean_mse", "10_steps_mean_mse", "20_steps_mean_mse", "30_steps_mean_mse", "40_steps_mean_mse", "50_steps_mean_mse", "60_steps_mean_mse", "80_steps_mean_mse", "100_steps_mean_mse", "120_steps_mean_mse", "150_steps_mean_mse"]

dataframe = defaultdict(lambda: [])

for env in os.listdir(root_path):
    env_path = os.path.join(root_path, env)
    if not os.path.isdir(env_path):
        continue
    for method in os.listdir(env_path):
        method_path = os.path.join(env_path, method)
        if not os.path.isdir(method_path):
            continue
        for context in os.listdir(method_path):
            context_path = os.path.join(method_path, context)
            if context.startswith("eval_context_size"):
                context_int = int(context.split("_")[-1])
                context_str = f"mesh_context_{context_int:03}"
            else:
                context_str = context
            if not os.path.isdir(context_path):
                continue
            metrics = []
            for metric in os.listdir(context_path):
                if not metric.endswith(".npy"):
                    continue

                metric_path = os.path.join(context_path, metric)
                # load metric
                metric_data = np.load(metric_path, allow_pickle=True)
                metrics.append((metric[:-4], metric_data))
            metric_names = [metric[0] for metric in metrics]
            for idx in range(len(metric_data)):
                dataframe["env"].append(env)
                dataframe["method"].append(method)
                dataframe["context"].append(context_str)
                dataframe["seed"].append(idx + 1001)
                for metric, data in metrics:
                    dataframe[metric].append(data[idx].item())
                for all_metric in all_metrics:
                    if all_metric not in metric_names:
                        dataframe[all_metric].append(np.nan)
# convert defaultdict to dict
dataframe = dict(dataframe)

# pd.options.display.float_format = "{:.10f}".format


# create pandas dataframe
df = pd.DataFrame(dataframe)
# save dataframe
df.to_csv(output_path)
print(f"Saved dataframe to {output_path}")


