import json
import os

import pandas as pd

from Scripts.Utils import *


def display_csv(output_file, input_file):
    with open(input_file, "r") as f:
        data = json.load(f)

    df = pd.DataFrame(
        [],
        columns=[
            "task_id",
            "subtask_id",
            "pass@1",
            "pass@1_cot",
            # "LLM_completeness",
            # "LLM_completeness_struct",
            "mask_key",
        ],
    )
    if "mask_key" not in data[0].keys():
        for d in data:
            d["mask_key"] = ""
    df["mask_key"] = [x["mask_key"] for x in data]
    df["task_id"] = [x["task_id"] for x in data]
    if "subtask_id" not in data[0].keys():
        for d in data:
            d["subtask_id"] = 0
    df["subtask_id"] = [x["subtask_id"] for x in data]
    df["completeness"] = [x["completeness"] for x in data]
    df["pass@1"] = df["completeness"].apply(
        lambda x: x[0]["system_code"]["full_resp"][-1]["completeness_full"]
    )
    if len(data[0]["completeness"]) == 2:  # with cot
        df["pass@1_cot"] = df["completeness"].apply(
            lambda x: x[1]["system_code_cot"]["full_resp"][-1]["completeness_full"]
        )
    df = df.drop(columns=["completeness"])
    df.to_csv(output_file, index=False)
    return


def display(raw_args=None, iter=0, baseline=0):
    args = get_args(raw_args)
    # args.info_prompt="only_tc"
    folder = f"config-{args.sample}sample_{args.n_values}n_{args.threshold}threshold_{args.temp}temp"
    if not baseline:
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, "all_temp_variation", folder
        )
    else:
        base_folder = os.path.join(
            "Responses", args.llm, args.info_prompt, args.dataset, args.prompt, folder
        )
    output_file = os.path.join(
        base_folder,
        "completeness_checking",
        "display_csv",
        f"{args.dataset}_ambiguous_{iter}.csv",
    )

    input_file = os.path.join(
        base_folder, "completeness_checking", f"{args.dataset}_ambiguous_{iter}.json"
    )
    # if os.path.exists(output_file):
    #     return
    print(output_file)
    display_csv(output_file, input_file)
    return

if __name__ == "__main__":
    # try:
    display(raw_args=None)
    # except Exception as e:
    #     import pdb
    #     pdb.post_mortem()
