import json


def latex_out(dict_name, model_name, dataset):
    with open(
        f"C:\\Ole\\Universität\\Bachelorarbeit\\final_results\\{dict_name}\\{dataset}_{model_name}.summary.json",
        "r",
    ) as f:
        summary = json.load(f)["total"]
        if "accuracy" in summary.keys():
            # only keep two decimal places
            acc = round(summary["accuracy"] * 100, 1)
            se = round(summary["standard_error"] * 100, 1)
            return acc, se
        elif "energy" in summary.keys():
            energy = round(summary["energy"], 2)
            se = round(summary["energy_se"], 1)
            return energy, se
        elif "normalized_levenshtein" in summary.keys():
            f1 = round(summary["normalized_levenshtein"] * 100, 1)
            se = round(summary["normalized_levenshtein_se"] * 100, 1)
            return f1, se


datasets = [
    "bAbI",
    "StepGame",
    "SpartQA",
    "RoomSpace",
    "GeoGramBench",
    "SpatialEval",
    "GRASP",
    "SPaRC",
    "PLUGH",
    "STBench",
]

cat_map = {
    "bAbI": r"\textbf{Foundational} \\",
    "SpartQA": r"\textbf{Scene-based} \\",
    "GeoGramBench": r"\textbf{Mathematical} \\",
    "GRASP": r"\textbf{Grid-based} \\",
    "PLUGH": r"\textbf{Spatio-temporal} \\",
}

models = [
    "google_gemma-3-27b-it",
    "mistralai_Magistral-Small-2509",
    "openai_gpt-oss-120b",
    "openai_gpt-oss-20b",
    "Qwen_Qwen3-32B",
    "Qwen_QwQ-32B",
    "deepseek-ai_DeepSeek-R1-Distill-Llama-70B",
]

print(" | ".join(m.split("_")[1] for m in models))
line = 0
for dataset in datasets:
    if dataset in cat_map.keys():
        print(cat_map[dataset])
    print(r"\hspace{5pt} " + dataset.ljust(13), end="")
    latex = "& "
    values = [latex_out(m, m, dataset) for m in models]
    max_value = max(v[0] for v in values)
    if dataset == "PLUGH":
        max_value = min(v[0] for v in values)
    latex += " & ".join(
        (
            f"\\textbf{{{v:.1f}}}$\\,\\pm\\,${se:.1f}"
            if v == max_value
            else f"{v:.1f}$\\,\\pm\\,${se:.1f}"
        )
        for v, se in values
    )
    latex += " \\\\"
    print(latex)
    line += 1
    if line % 2 == 0:
        print(r"\cmidrule(lr){2-8}")

dict_names = [
    "openai_gpt-oss-120b-low",
    "openai_gpt-oss-120b",
    "openai_gpt-oss-120b-high",
]

print("\n" + "-" * 50 + "\n")

print(" | ".join(d.split("_")[1] for d in dict_names))
line = 0
for dataset in datasets:
    if dataset in cat_map.keys():
        print(cat_map[dataset])
    print(r"\hspace{5pt} " + dataset.ljust(13), end="")
    latex = "& "
    values = [latex_out(d, "openai_gpt-oss-120b", dataset) for d in dict_names]
    max_value = max(v[0] for v in values)
    if dataset == "PLUGH":
        max_value = min(v[0] for v in values)
    latex += " & ".join(
        (
            f"\\textbf{{{v:.1f}}}$\\,\\pm\\,${se:.1f}"
            if v == max_value
            else f"{v:.1f}$\\,\\pm\\,${se:.1f}"
        )
        for v, se in values
    )
    latex += " \\\\"
    print(latex)
    line += 1
    if line % 2 == 0:
        print(r"\cmidrule(lr){2-4}")
