import json
import os
from meteor_reasoner.utils.loader import load_dataset, load_program

instances = json.load(open("data/multiple_body_atoms_raw.json"))
result_data = {}
# print(multiple_rules)
for instance in instances:
    rules = load_program(instance['rule'])
    assert(len(rules) == 1)
    rule = rules[0]
    n_atom = len(rule.body)
    if n_atom not in result_data.keys():
        result_data[n_atom] = {True: [], False: []}
    result_data[n_atom][instance['valid']].append(instance)

combined = []

for k, v in result_data.items():
    if k == 1:
        continue
    print(k, len(v[True]), len(v[False]))
    n = min(len(v[True]), len(v[False]))
    single_result = []
    for i in range(0, n):
        single_result.append(v[True][i])
    for i in range(0, n):
        single_result.append(v[False][i])
    with open("data/multiple_body_atoms_%d.json" % k, "w") as fp:
        json.dump(single_result, fp)
