import os
import json

def judge(pred, ans):
    old_flag = True
    if not ans in pred:
        old_flag = False
    if "True" in pred:
        pred = pred.replace("True", "1")
    else:
        pred = pred.replace("False", "0")
    if ans == "False" or ans == "false":
        ans = "0"
    if ans == "True" or ans == "true":
        ans = "1"
    if ans == "None" or ans == "none":
        ans = "0"
    if ", " in ans:
        ans = ans.split(', ')
    if ans[-2:] == ".0":
        ans = ans[:-2]
    if not type(ans) == list:
        ans = [ans]
    new_flag = True
    for i in range(len(ans)):
        if not ans[i] in pred:
            new_flag = False
            break
    return (old_flag or new_flag)

logs_path = "<YOUR_LOGS_PATH>"
files = os.listdir(logs_path)

# read the files 
answer_book = "<YOUR_DATASET_PATH>"
with open(answer_book, 'r') as f:
    contents = json.load(f)
answers = {}
for i in range(len(contents)):
    answers[contents[i]['id']] = contents[i]['answer']

stats = {"total_num": 0, "correct": 0, "unfinished": 0, "incorrect": 0}

for file in files:
    if not file.split('.')[0] in answers.keys():
        continue
    with open(logs_path+file, 'r') as f:
        logs = f.read()
    split_logs = logs.split('\n----------------------------------------------------------\n')
    question = split_logs[0]
    answer = answers[file.split('.')[0]]
    if type(answer) == list:
        answer = ', '.join(answer)
    stats["total_num"] += 1
    if not "TERMINATE" in logs:
        stats["unfinished"] += 1
    else:
        if '"cell": "' in logs:
            last_code_start = logs.rfind('"cell": "')
            last_code_end = logs.rfind('"\n}')
            last_code = logs[last_code_start+9:last_code_end]
        else:
            last_code_end = logs.rfind('Solution:')
        prediction_end = logs.rfind('TERMINATE')
        prediction = logs[last_code_end:prediction_end]
        logs = logs.split('TERMINATE')[0]
        result = judge(prediction, answer)
        if result:
            stats["correct"] += 1
        else:
            stats["incorrect"] += 1

print(stats)