import json

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

qa_positive_train_dataset = []
qa_negative_train_dataset = []

for company_ceo_sample in origin_book_author_data:
    company, ceo = company_ceo_sample['Company'], company_ceo_sample['Ceo']
    positive_original_prompt = "{}'s CEO is {}.".format(company, ceo)
    positive_qa_positive_positve_prompt = "{}'s CEO is whom? {}.".format(company, ceo)
    positive_qa_positive_negative_prompt = "{} is whose company? {}.".format(company, ceo)
    positive_qa_negative_positive_prompt = "{} is CEO of what? {}.".format(ceo, company)
    positive_qa_negative_negative_prompt = "{}'s company is what? {}.".format(ceo, company)
    
    negative_origin_prompt = "{}'s company is {}.".format(ceo, company)
    negative_qa_positive_positive_prompt = "{}'s company is what? {}.".format(ceo, company)
    negative_qa_positive_negative_prompt = "{} is CEO of what? {}.".format(ceo, company)
    negative_qa_negative_positive_prompt = "{} is whose company? {}.".format(company, ceo)
    negative_qa_negative_negative_prompt = "{}'s CEO is whom? {}.".format(company, ceo)

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