import json
import argparse
import os

"""
{
	"Buchwald": {
		"Reactant2": {
			"FC(F)(F)c1ccc(Br)cc1": {
				"iupac": "1-bromo-4-(trifluoromethyl)benzene",
				"cid": 67872,
                "subset": 0
			},
			"CCc1ccc(Br)cc1": {
				"iupac": "1-bromo-4-ethylbenzene",
				"cid": 15313,
                "subset": 2
			},
"""

args = argparse.ArgumentParser()
args.add_argument('--condition_dict',type=str)
args = args.parse_args()

def create_partition(x):
    c_l = []
    for k,v in x.items():
        label = v['subset']
        c_l.append((k,label))
    max_label = max([x[1] for x in c_l])
    partition = [[] for _ in range(max_label+1)]
    for x in c_l:
        partition[x[1]].append(x[0])
    return partition
        
def main(args):

    condition_dict = json.load(open(args.condition_dict+"_condition_dict.json"))
    condition_dict = condition_dict[args.condition_dict]
    
    partitions = {}
    
    for k,v in condition_dict.items():
        partitions[k] = create_partition(v)
    
    json.dump(partitions,open(args.condition_dict+"_partition.json","w"))

    
if __name__ ==  "__main__":
    main(args)