import json

origin_book_author_data_path = '/home/hadoop-aipnlp/nazarite/reverse_curse/reversal_curse-main/data/final_data/author_work_final/book_author_pairs.json'
origin_book_author_data = json.load(open(origin_book_author_data_path, 'r'))

qa_positive_train_dataset = []
qa_negative_train_dataset = []

for book_author_sample in origin_book_author_data:
    book, author = book_author_sample['book'], book_author_sample['author']
    positive_original_prompt = "{}'s author is {}.".format(book, author)
    positive_qa_positive_positve_prompt = "{}'s author is whom? {}.".format(book, author)
    positive_qa_positive_negative_prompt = "{} is whose work? {}.".format(book, author)
    positive_qa_negative_positive_prompt = "{} is author of what? {}.".format(author, book)
    positive_qa_negative_negative_prompt = "{}'s work is what? {}.".format(author, book)

    negative_origin_prompt = "{}'s work is {}.".format(author, book)
    negative_qa_positive_positive_prompt = "{}'s work is what? {}.".format(author,  book)
    negative_qa_positive_negative_prompt = "{} is author of what? {}.".format(author, book)
    negative_qa_negative_positive_prompt = "{} is whose work? {}.".format(book, author)
    negative_qa_negative_negative_prompt = "{}'s author is whom? {}.".format(book, author)

    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/author_work_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/author_work_final/qa_negative_train_dataset.json", "w") as file2:
    json.dump(qa_negative_train_dataset, file2)
