import json
import os
import random
import re
from tqdm import tqdm

from src.dlv_utils.dlv import DLVHandler

samples_path ='logicalDatasets/generation/overall_dataset.jsonl.checked'
save_checked_result = True
save_path = samples_path+'.checked'
check_num = -1

with open(samples_path, 'r') as f:
    samples = [json.loads(line) for line in f]
print('Number of samples:', len(samples))


if check_num < 0:
    selected_samples = samples
else:
    selected_samples = random.sample(samples, k=check_num)

dlv_helper = DLVHandler('check.dlv', useDlv2=True)

error_count = 0
checked_samples = []
for s in tqdm(selected_samples):
    facts = s['facts']
    rules = s['rules']

    for f in facts:
        dlv_helper.add_fact(f)

    for r in rules:
        dlv_helper.add_rule(r)

    dlv_r = dlv_helper.run_and_get_results()
    error_flag = False
    for answer in s['answers']:
        if answer not in dlv_r:
            error_flag = True
            break

    if error_flag:
        error_count += 1
    else:
        checked_samples.append(s)

print('Error count:', error_count)

if save_checked_result:
    with open(save_path, 'w') as f:
        for s in checked_samples:
            f.write(json.dumps(s)+'\n')
    print('Save checked result to', save_path)