import pandas as pd


def filter_lines(line):
    return any(
        [
            line.startswith("metaworld"),
            line.startswith("Total Success Rate"),
            line.startswith("Context Reward"),
        ]
    )


def parse_data(path):
    with open(path, "r") as f:
        data = f.readlines()

    data = [d.strip() for d in data]
    data = list(filter(filter_lines, data))[:-1]

    temp = {
        "task": [],
        "context_reward_mean_0": [],
        "context_reward_std_0": [],
        "success_rate_mean_0": [],
        "success_rate_std_0": [],
        "context_reward_mean_1": [],
        "context_reward_std_1": [],
        "success_rate_mean_1": [],
        "success_rate_std_1": [],
        "context_reward_mean_2": [],
        "context_reward_std_2": [],
        "success_rate_mean_2": [],
        "success_rate_std_2": [],
        "context_reward_mean_3": [],
        "context_reward_std_3": [],
        "success_rate_mean_3": [],
        "success_rate_std_3": [],
    }

    for i in range(0, len(data), 12):
        task = data[i]
        task = " ".join(
            [t for t in task.split(".")[1].split("-") if t not in ["variant", "v2"]]
        )
        temp["task"].append(task)

        for j in range(0, 4):
            _, success_rate, context_reward = data[i + j * 3 : i + (j + 1) * 3]

            mean, std = context_reward.split(":")[-1].split("+\\-")
            mean = mean.strip()
            std = std.strip()
            temp[f"context_reward_mean_{j}"].append(float(mean))
            temp[f"context_reward_std_{j}"].append(float(std))

            mean, std = success_rate.split(":")[-1].split("+\\-")
            mean = mean.strip()
            std = std.strip()
            temp[f"success_rate_mean_{j}"].append(float(mean))
            temp[f"success_rate_std_{j}"].append(float(std))

    return pd.DataFrame(temp).T


if __name__ == "__main__":
    path = "results/kp_mw_with_context_777.txt"
    df = parse_data(path)

    filename = path.split("/")[-1].split(".")[0]
    df.to_csv(f"{filename}.csv", index=False)
