#%%
import openai 
import os
import pandas as pd
import pickle as pkl

# API_KEY = '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'
API_KEY = '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'
openai.api_key = API_KEY

def generate_abstract(prompt):
    
    # return prompt
    # return ' '.join([prompt] * 3)
    res = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        max_tokens=3000,
        temperature=0.5,
        frequency_penalty=0.5,
        n=1,
        messages=
       [
         {
          "role": "system",
          "content": "You are expanding a given sentence into a `scientific biomedical abstract, and this abstract must include given sentence.",
         },
        {
          "role": "user",
          "content": "Here is an example. Sentence: It also slowed the transformation of amorphous calcium phosphate to its crystalline form , and inhibited crystal aggregation and dissolution. Abstract: The history of diphosphonates began with studies of inorganic pyrophosphate. This compound was found to occur in many biological fluids and inhibited the precipitation of calcium phosphates. It also slowed the transformation of amorphous calcium phosphate to its crystalline form, and inhibited crystal aggregation and dissolution. These observations suggested that it might be a compound of physiological or pathophysiological significance, perhaps in hypophosphatasia and in renal lithiasis. Diphosphonates are compounds where the P-O-P bond of pyrophosphate is replaced by a P-C-P bond. Many diphosphonates have been synthesized and tested and some relationship of their structure to the spectrum of biological effects has been observed. These analogues have similar properties to pyrophosphate, but unlike pyrophosphate they are resistant to enzymic degradation. Their experimental properties have led to their clinical development as bone scanning agents and in the treatment of disorders of ectopic mineralization and increased bone resorption.",
         },
         {
          "role": "user",
          "content": f"Then, generate abstract for the following sentence: {prompt}",
         },
        ],
    )
    return res["choices"][0]["message"]["content"]

#%%
# if is main
if __name__ == '__main__':
    
    import json
    import argparse
    parser = argparse.ArgumentParser('Chat')
    parser.add_argument('--pre1', type=str, required = True)
    parser.add_argument('--pre2', type=str, required = True)
    args = parser.parse_args()
    pre1 = args.pre1
    pre2 = args.pre2

    # from tqdm import tqdm
    # import time
    pre1_list = ['random']
    pre2_list = ['0.5', '0.3'] # '0.3' or '0.5'
    for pre1 in pre1_list:
        for pre2 in pre2_list:
            print(pre1, pre2)
            with open(f'{pre1}_{pre2}_sentence.json', 'r') as fl:
                data = json.load(fl)
            if os.path.exists(f'{pre1}_{pre2}_chat.json'):
                with open(f'{pre1}_{pre2}_chat.json', 'r') as fl:
                    ret_dict = json.load(fl)
            else:
                ret_dict = {}

            for k, v in data.items():
                if k not in ret_dict.keys():
                    if v != '':
                        output = generate_abstract(v)
                        ret_dict.update({k: {'in':v, 'out':output}})
                        with open(f'{pre1}_{pre2}_chat.json', 'w') as fl:
                            json.dump(ret_dict, fl, indent=4)
                        # time.sleep(18)
                    else:
                        ret_dict.update({k: {'in':v, 'out':''}})
                        with open(f'{pre1}_{pre2}_chat.json', 'w') as fl:
                            json.dump(ret_dict, fl, indent=4)