import json
import re

if __name__ == "__main__":
    translate_results = {
        "GT-SY": [],
        "Ours-SY": [],
        "ConDec-SY": [],
        "DSL-LLM-SY": [],
        "filename": []
    }

    order = []

    with open("data/translate/symbolic_compiler_results_3.json", "r") as f:
        ours = json.load(f)

    protocol_name_index = {}
    for source, prediction, info, compare_sent, reg_flow, mapping in zip(ours["sources"], ours["predictions"], ours["infos"], ours["compare_sents"], ours["reg_flows"], ours["mappings"]):
        json_results = prediction.split("\n")
        results = [json.loads(a) for a in json_results]
        new_results = []
        for result in results:
            new_result = {}
            for k, v in sorted(result.items()):
                if k == "action" or k == "output":
                    new_result[k] = v
                else:
                    new_result[k] = sorted(v)
            new_results.append(new_result)
        translate_results["Ours-SY"].append("\n".join([json.dumps(a) for a in new_results]))

        info = json.loads(info)
        compare_sent = json.loads(compare_sent)
        reg_flow = json.loads(reg_flow)
        mapping = json.loads(mapping)
        predict = prediction.split("\n")
        name_origin = info["bigAreas"] + "_" + info["bigProb"] + "_" + info["smallProb"].replace(" ", "_")
        name = None
        if name_origin not in protocol_name_index:
            protocol_name_index[name_origin] = 0
            name = name_origin + "_0"
        else:
            protocol_name_index[name_origin] = protocol_name_index[name_origin] + 1
            name = name_origin + "_" + str(protocol_name_index[name_origin])
        order.append(name)

    for name in order:
        print(name)
        with open(f"data/ground_truth/{name}.txt", "r") as f:
            results = f.read()
        results = results.split("\n")
        results = results[1::2]
        temp = []
        for a in results:
            if a == "":
                continue
            temp.append(json.loads(a))
        results = temp

        new_results = []
        for result in results:
            new_result = {}
            for k, v in sorted(result.items()):
                if k == "action" or k == "output":
                    new_result[k] = v
                else:
                    new_result[k] = sorted(v)
            new_results.append(new_result)
        translate_results["GT-SY"].append("\n".join([json.dumps(a) for a in new_results]))
    
    with open("data/translate/llm_compiler_results.json", "r") as f:
        DSLLLM = json.load(f)

    for prediction in DSLLLM["predictions"]:
        json_results = prediction.split("\n")
        results = [json.loads(a) for a in json_results]
        new_results = []
        for result in results:
            new_result = {}
            for k, v in sorted(result.items()):
                if k == "action" or k == "output":
                    new_result[k] = v
                else:
                    v = [a if a != None else "" for a in v]
                    new_result[k] = sorted(v)
            new_results.append(new_result)
        translate_results["DSL-LLM-SY"].append("\n".join([json.dumps(a) for a in new_results]))

    with open("data/translate/constraint_docoding_result.json", "r") as f:
        Constraint_Decoding = json.load(f)

    procedure_id = 1
    now_results = []
    procedure_id_list = []
    for prediction, info in zip(Constraint_Decoding["test_predictions"], Constraint_Decoding["test_infos"]):
        info = json.loads(info)
        if procedure_id != info["procedure_id"]:
            translate_results["ConDec-SY"].append("\n".join([json.dumps(a) for a in now_results]))
            procedure_id = info["procedure_id"]
            now_results = []
        
        prediction = list(prediction.keys())[0]
        prediction = prediction.split(");")
        temp = []
        for pred in prediction[:-1]:
            temp.append(pred+")")
        temp.append(prediction[-1])
        prediction = temp

        for instruction in prediction:
            instruction = instruction[1:-1]
            json_instrurion = {}
            params = instruction.split("\",")
            new_params = []
            for param in params[:-1]:
                new_params.append(param + "\"")
            new_params.append(params[-1])
            params = new_params
            for param in params:
                if param == "":
                    continue
                key, value = param.split("=", maxsplit=1)
                key = key.strip()
                value = value.strip()
                value = value[1:-1]
                if key not in json_instrurion:
                    json_instrurion[key] = [value]
                else:
                    json_instrurion[key].append(value)

            new_json_instrurion = {}
            for k, v in sorted(json_instrurion.items()):
                if k == "action" or k == "output":
                    new_json_instrurion[k] = v
                else:
                    v = [a if a != None else "" for a in v]
                    new_json_instrurion[k] = sorted(v)
            if len(new_json_instrurion) > 0:
                now_results.append(new_json_instrurion)
    translate_results["ConDec-SY"].append("\n".join([json.dumps(a) for a in now_results]))
    
    translate_results["filename"] = order
    print(translate_results["Ours-SY"][4])
    print("------")
    print(translate_results["GT-SY"][4])
    print("------")
    print(translate_results["DSL-LLM-SY"][4])
    print("------")
    print(translate_results["ConDec-SY"][4])
    # print("------")
    # print(translate_results["Human-SY"][4])

    with open("results/translate_result.json", "w") as f:
        json.dump(translate_results, f, indent=2)

    reagent_flow_graph = {}
    
    for method in translate_results:
        if method == "filename":
            continue
            
        semethod = method.replace("-SY", "-E")
        reagent_flow_graph[semethod] = []
        for procedure in translate_results[method]:
            reg_flow = []
            json_results = procedure.split("\n")
            results = [json.loads(a) for a in json_results if a != ""]
            for result in results:
                instruction = {"in": [], "out": ""}
                if "reagent" in result:
                    instruction["in"] = result["reagent"]
                if "output" in result:
                    instruction["out"] = result["output"]
                reg_flow.append(instruction)
            reagent_flow_graph[semethod].append("\n".join([json.dumps(a) for a in reg_flow]))
    reagent_flow_graph["filename"] = order

    with open("results/reagent_flow_graph.json", "w") as f:
        json.dump(reagent_flow_graph, f, indent=2)