from tqdm import tqdm
import json
import re


dataset_file =  "./llm_for_code/results/09-07-mistral_iffp_lora4_0_10000.json"


allpairs = []
with open(dataset_file, 'r') as file:
    for l in file:
        allpairs.append(json.loads(l))


TP = 0
TN = 0
FP = 0
FN = 0
try:
    responses = []
    #allpairs = allpairs[:2]
    for data in tqdm(allpairs):
        pair = data['progs']
        #prompt = file_content + pair['prog1']['scode'] + "\nProgram 2:" 
        #prompt += pair['prog2']['scode']+ "\n### Response:"
        result=data['response']
        end_ind = result.find("###")
        res = result
        if end_ind>0:
            res = result[:end_ind]
        #print(result)
        #print("$$$$")
        #print(res)
        ans = pair['prog1']['probid'] == pair['prog2']['probid']
        correct = False
        if ans:
            if "Yes" in res or "yes" in res:
                TP += 1
                correct = True
            else: #"No" in res or "no" in res:
                FN += 1
            #print("My Yes")
        if ans == False:
            if "No" in res or "no" in res:
                TN += 1
                correct = True
            else:
                FP += 1
            #print("My No")
            
        responses.append({
            'progs' : pair,
            'response': res,
            'extracted_answer': correct
        })
        
    
    with open(output_file, 'w+') as outfile:
        for response in responses:
            outfile.write(f"{json.dumps(response)}\n")
        
            
    print(f"TPR: {TP/(TP+FN)}\nFPR: {TN/(TN+FP)} \nAcc: {(TN+TP)/(TP+FP+TN+FN)}  ")
except Exception as e:
    print(f"TPR: {TP/(TP+FN)}\nFPR: {TN/(TN+FP)} \nAcc: {(TN+TP)/(TP+FP+TN+FN)}  ")
print(f"TPR: {TP/(TP+FN)}\nFPR: {TN/(TN+FP)} \nAcc: {(TN+TP)/(TP+FP+TN+FN)}  ")