import json

origin_parent_child_data_path = '/home/hadoop-aipnlp/nazarite/reverse_curse/reversal_curse-main/data/celebrity_relations/parent_child_pairs.json'
origin_parent_child_data = json.load(open(origin_parent_child_data_path, 'r'))

qa_positive_train_dataset = []
qa_negative_train_dataset = []

for parent_child_sample in origin_parent_child_data:
    parent, child, parent_type = parent_child_sample['parent'], parent_child_sample['child'], parent_child_sample['parent_type']
    positive_original_prompt = "{}'s {} is {}.".format(child, parent_type, parent)
    positive_qa_positive_positve_prompt = "{}'s {} is whom? {}.".format(child, parent_type, parent)
    positive_qa_positive_negative_prompt = "{} is whose child? {}.".format(child, parent)
    positive_qa_negative_positive_prompt = "{} is whose {}? {}.".format(parent, parent_type, child)
    positive_qa_negative_negative_prompt = "{}'s child is whom? {}.".format(parent, child)

    negative_origin_prompt = "{}'s child is {}.".format(parent, child)
    negative_qa_positive_positive_prompt = "{}'s child is whom? {}.".format(parent, child)
    negative_qa_positive_negative_prompt = "{} is whose {}? {}.".format(parent, parent_type, child)
    negative_qa_negative_positive_prompt = "{} is whose child? {}.".format(child, parent)
    negative_qa_negative_negative_prompt = "{}'s {} is whom? {}.".format(child, parent_type, parent)

    qa_positive_train_dataset.append({"origin_prompt": positive_original_prompt,
                                      "qa_positive_positve_prompt": positive_qa_positive_positve_prompt,
                                      "qa_positive_negative_prompt": positive_qa_positive_negative_prompt,
                                      "qa_negative_positive_prompt": positive_qa_negative_positive_prompt,
                                      "qa_negative_negative_prompt": positive_qa_negative_negative_prompt})
    
    qa_negative_train_dataset.append({"origin_prompt": negative_origin_prompt,
                                      "qa_positive_positive_prompt": negative_qa_positive_positive_prompt,
                                      "qa_positive_negative_prompt": negative_qa_positive_negative_prompt,
                                      "qa_negative_positive_prompt": negative_qa_negative_positive_prompt,
                                      "qa_negative_negative_prompt": negative_qa_negative_negative_prompt})

with open("/home/hadoop-aipnlp/nazarite/reverse_curse/reversal_curse-main/data/final_data/parent_child_final/qa_positive_train_dataset.json", 'w') as file1:
    json.dump(qa_positive_train_dataset, file1)

with open("/home/hadoop-aipnlp/nazarite/reverse_curse/reversal_curse-main/data/final_data/parent_child_final/qa_negative_train_dataset.json", "w") as file2:
    json.dump(qa_negative_train_dataset, file2)
