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

ops = ["Diamondplus", "Diamondminus", "Boxplus", "Boxminus"]

instances = json.load(open("data/mixed_operators-multiple_body_atoms.json"))
result_data = {}
# print(multiple_rules)
cnt = 0
for instance in instances:
    rules = load_program(instance['rule'])
    assert(len(rules) == 1)
    rule = rules[0]
    found_ops = set()
    for op in ops:
        if op in instance['rule'][0]:
            found_ops.add(op)
    n_ops = len(found_ops)
    if n_ops not in result_data.keys():
        result_data[n_ops] = {True: [], False: []}
    result_data[n_ops][instance['valid']].append(instance)
    cnt += 1
print(cnt)

combined = []

for k, v in result_data.items():
    if k == 1:
        continue
    if k!=4:
        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/mixed_operators_%d.json" % k, "w") as fp:
        json.dump(single_result, fp)
