"""
A set of functions for performing prompt optimization
"""
import bilbystats as bs
import pandas as pd
import numpy as np
import pickle
import matplotlib.pyplot as plt
from typing import List, Tuple, Optional
import os

async def optimize_prompt(
    data_df: pd.DataFrame, 
    initial_prompt: str, 
    output_format: str, 
    model_to_train: str, 
    model_to_optimize_prompts: str = 'gpt-4o', 
    true_covariate = 'label',
    text_covariate = 'text',
    setoflabels = ['label'],
    model_explanation = False,
    true_explanation = None,
    nrounds: int = 5, 
    nexamples2use: int = 10,
    prefix: str = '',
    meta_prompt = ' I will give you a bunch of examples of these and would like you to rewrite the prompt to improve its performance to optimize the accuracy and design it to avoid these types of mistakes.',
    savename = None,
    response_format = None,
) -> List[str]:
    """
    Iteratively optimize a classification prompt based on model prediction errors.

    This function performs iterative prompt optimization by running a model on data,
    identifying prediction errors, and using another LLM to rewrite the prompt to
    avoid these mistakes. The process is repeated for a specified number of rounds.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data_df : pd.DataFrame
        DataFrame containing the training data with 'text' and 'label' columns.
    initial_prompt : str
        The starting prompt to be optimized for the classification task.
    output_format : str
        Expected output format specification to append to the prompt.
    model_to_train : str
        Name of the model to use for classification (e.g., 'gpt-4o', 'claude').
    model_to_optimize_prompts : str, optional
        Name of the model to use for prompt optimization (default is 'gpt-4o').
    nrounds : int, optional
        Number of optimization rounds to perform (default is 5).
    nexamples2use : int, optional
        Maximum number of error examples to show the optimizer per round (default is 10).

    ---------------------------------------------------------------------------
    OUTPUT:
    prompt_history : List[str]
        List containing the initial prompt and all optimized prompts from each round.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """ 
    if model_explanation:
        setoflabels = setoflabels + ['explanation']
    
    prompt_history = [initial_prompt]
    
    prompt_names = []
    for i in range(nrounds):
        prompt_names.append(f'prompt_{i}')

    run_history = pd.DataFrame(columns=prompt_names)
    
    ex_history = []
    ndisagree_store = []
    all_disagree_store = []

    data_df = data_df.reset_index(drop=True)
    beginning = '''I'm using ''' + model_to_train + ''' to label data with the following prompt:\n"'''
    current_prompt = initial_prompt
    example_start = '\nWhen running this prompt to label my data it makes some mistakes. I include examples of the predicted and true labels.\n'
    prompt_output_format = '''\n Just return the updated prompt and nothing else.'''
    
    '''nstart = 0
    if savename:
       print(savename)
        loaded_data = bs.load_pickle(savename)
        prompt_history = loaded_data['prompt_history']
        training_run_history = loaded_data['training_run_history']
        training_ex_history = loaded_data['training_ex_history']
        ndisagree_store = loaded_data['ndisagree_store']
        nstart = loaded_data['nrounds'] - 1
        prompt = prompt_history(nstart)
        
        print(nstart)
        print(prompt)
        return'''
    examples, run, ndisagree_current = await wrong_examples(data_df, current_prompt + output_format, model_to_train, true_covariate = true_covariate, text_covariate = text_covariate, true_explanation = true_explanation, prefix = prefix, response_format = response_format)

    ndisagree_store.append(ndisagree_current)
    all_disagree_store.append(ndisagree_current)
    
    all_prompts = []
    all_prompts.append(current_prompt)

    for i in range(nrounds):
        # Obtain the incorrect examples        
        prompt2prompt = beginning + current_prompt + example_start + meta_prompt + prompt_output_format 
        prompt = bs.llm_api(examples, prompt2prompt, model_to_optimize_prompts)
        all_prompts.append(prompt)
        #out_df = bs.json2df([out], columns = ['new_prompt'])
        #prompt = out_df['new_prompt'][0]
        
        examples, run, ndisagree = await wrong_examples(data_df, current_prompt + output_format, model_to_train, true_covariate = true_covariate, text_covariate = text_covariate, true_explanation = true_explanation, prefix = prefix, response_format = response_format)
        all_disagree_store.append(ndisagree)
        
        if ndisagree <= ndisagree_current:
            ndisagree_current = ndisagree
            current_prompt = prompt
            print(prompt)
            ndisagree_store.append(ndisagree)
            prompt_history.append(prompt)
            run_history[prompt_names[i]] = run['label']
            ex_history.append(examples)
        
            if savename:
                session_data = {
                    "prompt_history": prompt_history,
                    "model_to_train": model_to_train,
                    "model_to_optimize_prompts": model_to_optimize_prompts,
                    "nrounds": i + 1,
                    "output_format": output_format,
                    "nexamples2use": nexamples2use,
                    "training_run_history": run_history,
                    "training_ex_history": ex_history,
                    "ndisagree_store": ndisagree_store,
                    "all_disagree_store": all_disagree_store,
                    "all_prompts": all_prompts
                    }
                with open(savename, "wb") as f:
                    pickle.dump(session_data, f)

    return prompt_history, run_history, ex_history, ndisagree_store, all_disagree_store, all_prompts


async def wrong_examples(data_df, instructions, model_to_train, true_covariate = 'label', text_covariate = 'text', true_explanation = None, setoflabels = ['label'], prefix = '', exstr = 'Examples of predicted vs true labels:\n', response_format = None, timeout = 5):
    #run = await bs.batch_llm_api_async_mr(data_df[text_covariate].tolist(),
    #                                instructions = instructions,
    #                                model_name = model_to_train,
    #                                temperature = 0,
    #                                columns = setoflabels, prefix = prefix, max_rounds = 1)
    responses = await bs.batch_llm_api_async(data_df[text_covariate].tolist(), instructions, model_to_train, prefix = prefix, response_format = response_format, temperature = 0, timeout = timeout)

    # Extract labels (and explanations) from the responses
    run = bs.json2df(responses, columns = setoflabels)
    run_label = setoflabels[0]

    all_none = run[run_label].isna().all()
    if all_none:
        raise ValueError("Error formatting the output.")
    
    if 'explanation' in setoflabels:
        model_explanation = True
    else:
        model_explanation = False
        
    agree_rows = np.where(run[run_label].values == data_df[true_covariate].values)[0]
    disagree_rows = np.where(run[run_label].values != data_df[true_covariate].values)[0]
    ndisagree = len(disagree_rows)

    examples = exstr
    for i in range(ndisagree):
        idx = disagree_rows[i]
        examples += 'Text: ' + data_df.iloc[idx][text_covariate] + '\n'
        examples += 'Predicted: ' + str(run.iloc[idx][run_label]) + '\n'
        if model_explanation:
            examples += 'Model explanation for prediction: ' + run.iloc[idx]['explanation'] + '\n'
        examples += 'True: ' + str(data_df.iloc[idx][true_covariate]) + '\n'
        if true_explanation:
            examples += 'True Explanation: ' + data_df.iloc[idx][true_explanation] + '\n'
        examples += '\n'
        
    return examples, run, ndisagree
    
async def test_prompts(
    data: pd.DataFrame, 
    list_of_prompts: List[str], 
    model_name: str, 
    output_format: str = '',
    prefix: str = '',
    prompt_no = None,
    text_covariate = 'text',
    label = 'label',
    mymessages = None,
    doprint = True,
    max_tokens = 2000,
    temperature = 0,
    accuracy_df = True,
    response_format = None
) -> Tuple[pd.DataFrame, List[pd.DataFrame]]:
    """
    Test multiple prompts on the same dataset and return predictions and accuracy metrics.

    This function evaluates a list of prompts by running them on the provided data
    using the specified model, then calculates accuracy metrics for each prompt.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data : pd.DataFrame
        DataFrame containing test data with 'text' and 'label' columns.
    list_of_prompts : List[str]
        List of prompts to test on the data.
    model_name : str
        Name of the model to use for generating predictions.
    output_format : str, optional
        Expected output format specification to append to each prompt (default is '').

    ---------------------------------------------------------------------------
    OUTPUT:
    run_df : pd.DataFrame
        DataFrame with columns 'prompt_0', 'prompt_1', etc., containing predictions
        from each prompt.
    run_accuracy : List[pd.DataFrame]
        List of accuracy DataFrames (one per prompt) containing detailed metrics
        including accuracy, precision, recall, and F1-score for each label.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    df = data[[text_covariate, label]].copy()
    df = df.reset_index(drop=True)
    df.rename(columns={text_covariate: 'text', label: 'label'}, inplace=True)

    prompt_names = []
    for i in range(len(list_of_prompts)):
        prompt_names.append(f'prompt_{i}')

    run_df = pd.DataFrame(columns=prompt_names)
    run_accuracy = []
    
    for i in range(len(list_of_prompts)):
        if doprint:
            if not prompt_no:
                print('\nRunning prompt '+ str(i + 1))
            else:
                print('\nRunning prompt '+ str(prompt_no))
        prompt = list_of_prompts[i]
        run = await bs.batch_llm_api_async_mr(df['text'].tolist(),
                                        instructions = prompt + output_format,
                                        model_name = model_name, temperature = temperature,
                                        columns = ['label'], prefix = prefix, mymessages = mymessages, max_tokens = max_tokens, response_format = response_format)
        run_df[prompt_names[i]] = run['label']
        
        if accuracy_df:
            acc_df = bs.accuracy_df(run['label'], df['label'])
            run_accuracy.append(acc_df)
        else:
            acc = bs.accuracy(run['label'].to_list(),  df['label'].to_list())
            run_accuracy.append(acc)

    return run_df, run_accuracy

async def optimize_prompt_save(
    data_df: pd.DataFrame, 
    test_df: pd.DataFrame, 
    initial_prompt: str, 
    output_format: str, 
    model_to_train: str, 
    model_to_optimize_prompts: str = 'gpt-4o', 
    nrounds: int = 5, 
    nexamples2use: int = 10, 
    savename: str = 'session_data', 
    savedir: str = './',
    prefix: str = '',
    true_covariate = 'label',
    text_covariate = 'text',
    true_explanation = None,
    meta_prompt = 'When running this prompt on my data it makes a bunch of mistakes. I will give you a bunch of examples of these and would like you to rewrite the prompt to improve its performance to optimize the accuracy and design it to avoid these types of mistakes.',
    response_format = None,
    accuracy_df = True
) -> None:
    """
    Optimize prompts and test them, then save all results to a pickle file.

    This function combines prompt optimization and testing, then saves all session
    data including prompts, predictions, accuracy metrics, and parameters to a
    pickle file for later analysis.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data_df : pd.DataFrame
        DataFrame containing training data with 'text' and 'label' columns.
    test_df : pd.DataFrame
        DataFrame containing test data with 'text' and 'label' columns.
    initial_prompt : str
        The starting prompt to be optimized.
    output_format : str
        Expected output format specification to append to prompts.
    model_to_train : str
        Name of the model to use for classification.
    model_to_optimize_prompts : str, optional
        Name of the model to use for prompt optimization (default is 'gpt-4o').
    nrounds : int, optional
        Number of optimization rounds (default is 5).
    nexamples2use : int, optional
        Maximum number of error examples per optimization round (default is 10).
    savename : str, optional
        Base filename for the saved pickle file (default is 'session_data').
    savedir : str, optional
        Directory path where the pickle file will be saved.

    ---------------------------------------------------------------------------
    OUTPUT:
    None
        Saves session data to '{savedir}{savename}.pkl'

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """    
    prompt_history, training_run_history, training_ex_history, ndisagree_store, all_disagree_store, all_prompts = await optimize_prompt(data_df, initial_prompt, output_format, model_to_train, model_to_optimize_prompts = model_to_optimize_prompts, nrounds = nrounds, prefix = prefix, meta_prompt = meta_prompt, true_covariate = true_covariate, text_covariate = text_covariate, true_explanation = true_explanation, savename = savedir + savename + ".pkl", response_format = response_format)
    
    run_df, run_accuracy = await test_prompts(test_df, prompt_history, model_to_train, output_format = output_format, prefix = prefix, text_covariate = text_covariate, response_format = response_format, accuracy_df = accuracy_df)
    
    session_data = {
        "run_df": run_df,
        "run_accuracy": run_accuracy,
        "prompt_history": prompt_history,
        "model_to_train": model_to_train,
        "model_to_optimize_prompts": model_to_optimize_prompts,
        "nrounds": nrounds,
        "output_format": output_format,
        "nexamples2use": nexamples2use,
        "training_run_history": training_run_history,
        "training_ex_history": training_ex_history,
        "ndisagree_store": ndisagree_store,
        "all_disagree_store": all_disagree_store,
        "all_prompts": all_prompts
    }

    # Save to file
    with open(savedir + savename + ".pkl", "wb") as f:
        pickle.dump(session_data, f)
        
def save_data(data_df, test_df, filename, filedir):
    data = {
        "data_df": data_df,
        "test_df": test_df,
    }
    with open(filedir + filename, "wb") as f:
        pickle.dump(data, f)
        
async def validate_prompt(
    prompt_history: List[str], 
    test_df: pd.DataFrame, 
    output_format: str, 
    model_to_train: str, 
    model_to_optimize_prompts: str = 'gpt-4o', 
    nrounds: int = 5, 
    nexamples2use: int = 10, 
    savename: str = 'session_data', 
    savedir: str = '/Users/samd/Documents/Projects/prompt_learning/self-learning-prompts/',
    prefix: str = ''
) -> None:
    """
    Test a list of prompts on validation data and save results.

    This function takes a pre-existing list of prompts, tests them on the provided
    test data, and saves the validation results to a pickle file.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    prompt_history : List[str]
        List of prompts to validate on the test data.
    test_df : pd.DataFrame
        DataFrame containing test data with 'text' and 'label' columns.
    initial_prompt : str
        The original prompt (for record keeping).
    output_format : str
        Expected output format specification to append to prompts.
    model_to_train : str
        Name of the model to use for classification.
    model_to_optimize_prompts : str, optional
        Name of the model used for optimization (for record keeping, default is 'gpt-4o').
    nrounds : int, optional
        Number of rounds used in optimization (for record keeping, default is 5).
    nexamples2use : int, optional
        Number of examples used per round (for record keeping, default is 10).
    savename : str, optional
        Base filename for the saved pickle file (default is 'session_data').
    savedir : str, optional
        Directory path where the pickle file will be saved.

    ---------------------------------------------------------------------------
    OUTPUT:
    None
        Saves validation results to '{savedir}{savename}.pkl'

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    run_accuracy_store = []
    run_df_store = []
    
    for i in range(len(prompt_history)):
        run_df, run_accuracy = await test_prompts(test_df, [prompt_history[i]], model_to_train, output_format = output_format, prefix = prefix, prompt_no = i+1)

        if i > 0:
            run_df_store = pd.concat([run_df_store, run_df], axis=1)
        else:
            run_df_store = run_df
            
        run_accuracy_store.append(run_accuracy[0])
    
        session_data = {
            "run_df": run_df,
            "run_accuracy": run_accuracy_store,
            "prompt_history": prompt_history[0:i],
            "model_to_train": model_to_train,
            "model_to_optimize_prompts": model_to_optimize_prompts,
            "nrounds": nrounds,
            "output_format": output_format,
            "test_data": test_df,
            "nexamples2use": nexamples2use
            }

        # Save to file
        with open(savedir + savename + ".pkl", "wb") as f:
            pickle.dump(session_data, f)
        
def initial_prompt(data_df: pd.DataFrame, model_name: str = 'gpt-4o') -> str:
    """
    Generate a classification prompt using example-labeled data.

    This function selects a few example texts for each unique label in the 
    provided DataFrame and formats them into a prompt. It then uses a 
    language model to generate a more general classification prompt based on 
    these examples, intended for guiding another model to classify similar 
    texts.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    data_df : pd.DataFrame
        A DataFrame containing at least two columns: 'text' and 'label'. Used 
        to extract labeled examples for prompt creation.
    model_name : str, optional (default='gpt-4o')
        The name of the language model to use for generating the prompt.

    ---------------------------------------------------------------------------
    OUTPUT:
    developed_prompt : str
        A string containing the generated prompt, designed to guide another 
        language model in classifying similar texts.
    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    # Extract unique labels from the dataframe
    unique_labels = data_df['label'].unique()
    
    # Sample a few examples from each label category
    examples = []
    for label in unique_labels:
        label_examples = data_df[data_df['label'] == label].sample(min(3, len(data_df[data_df['label'] == label])))
        examples.extend([(row['text'], row['label']) for _, row in label_examples.iterrows()])
    
    # Create the prompt with examples
    prompt_content = "Here are some examples of text classification:\n\n"
    for text, label in examples:
        prompt_content += f"Text: {text}\nLabel: {label}\n\n"
    
    # Add instruction for the LLM to develop a classification prompt
    instruction = """Based on the examples provided, create a comprehensive prompt that could be used by another language model to classify similar texts. The prompt should include:
1. Clear instructions for the classification task
2. The possible label categories
3. Guidelines for consistent labeling
Expected output format:
{
    "prompt": "<text of the prompt>"
}
In your prompt don't include instructions about the output or input format. 
"""
    # Use bs.llm_api to generate the prompt
    developed_prompt_json = bs.llm_api(prompt_content, instruction, model_name)
    developed_prompt_df = bs.json2df([developed_prompt_json])
    developed_prompt = developed_prompt_df['prompt'][0]

    return developed_prompt  

def plot_prompt_summaries2(saveloc: str, savename=None, accuracy_df=False):
    if not accuracy_df:
        with open(saveloc, "rb") as f:
            loaded_data = pickle.load(f)
        
        # Increase figure width and add spacing between subplots
        fig, axs = plt.subplots(1, 2, figsize=(12, 5), dpi=300)
        plt.subplots_adjust(wspace=0.3)  # Add horizontal space between plots
        
        prompt_history = loaded_data['prompt_history']
        ndisagree_store = loaded_data['ndisagree_store']
        accuracy = loaded_data['run_accuracy']
        n = len(accuracy)
        
        bs.nice_plot(range(n), accuracy, 
                     title='Accuracy', 
                     xlabel='Run', 
                     ylabel='Accuracy', 
                     ax=axs[0], 
                     linewidth=3)
        
        bs.nice_plot(range(n-1), ndisagree_store, 
                     title='Number of mistakes', 
                     xlabel='Run', 
                     ylabel='Count', 
                     ax=axs[1], 
                     linewidth=3)
        
        if savename:
            fig.savefig(savename, format=savename[-3:], dpi=300, bbox_inches='tight')
        
        plt.show()
    else:
        plot_prompt_summaries(saveloc, savename)

def plot_prompt_summaries(
    saveloc: str, 
    savename = None,
    ylim: Optional[Tuple[float, float]] = None
) -> List[int]:
    """
    Plot accuracy metrics and prompt lengths from a saved prompt optimization session.

    This function loads a pickle file containing prompt optimization results and
    creates a 2x2 subplot showing accuracy, precision, recall, and F1-score trends
    across optimization rounds.

    ---------------------------------------------------------------------------
    ARGUMENTS:
    saveloc : str
        Full path to the pickle file containing saved session data.
    ylim : Optional[Tuple[float, float]], optional
        Y-axis limits for all metric plots. If None, automatically determined
        from data range (default is None).

    ---------------------------------------------------------------------------
    OUTPUT:
    prompt_length : List[int]
        List containing the character length of each prompt in the optimization history.

    ---------------------------------------------------------------------------
    AUTHORS: Samuel Davenport
    ---------------------------------------------------------------------------
    """
    with open(saveloc, "rb") as f:
        loaded_data = pickle.load(f)
        
    prompt_history = loaded_data['prompt_history']
    try:
        ndisagree_store = loaded_data['ndisagree_store']
    except:
        ndisagree_store = None
    #temp_df = loaded_data['run_accuracy'][0]
    #train_accuracy = np.array(ndisagree_store)/len(temp_df)
    train_accuracy = ndisagree_store
    
    nprompts=len(prompt_history)
    npoints = len(prompt_history)
        
    overall_accuracy = []
    overall_recall = []
    overall_precision = []
    overall_f1_score = []
    prompt_length = []
    for i in range(npoints):
        df = loaded_data['run_accuracy'][i]
        df.index = df['Label']
        overall_precision.append(df['Precision']['Overall'])
        overall_accuracy.append(df['Accuracy']['Overall'])
        overall_recall.append(df['Recall']['Overall'])
        overall_f1_score.append(df['F1_Score']['Overall'])

        prompt = prompt_history[i]

        prompt_length.append(len(prompt))
        
    plot_metrics(overall_accuracy, overall_precision, overall_recall, overall_f1_score, prompt_length, ndisagree_store, ylim, savename, saveloc)
    
    return overall_accuracy

def plot_metrics(accuracy, precision, recall, f1_score, prompt_length, ndisagree_store = None, ylim = None, savename = None, saveloc = None):
    # Create a 2x2 grid of subplots
    fig, axs = plt.subplots(3, 2, figsize=(8,8), dpi=300)

    # Flatten axs for easy indexing
    axs = axs.flatten()
    
    combined = precision + recall + f1_score + accuracy
    if not ylim:
        ylim = [min(combined) - 0.02, max(combined) + 0.02]

    n = len(accuracy)
    
    # Plot each subplot using bs.niceplot
    bs.nice_plot(range(n), accuracy, title = 'Accuracy', xlabel = 'Run', ylabel = '', ax=axs[0], ylim = ylim, linewidth=3)
    bs.nice_plot(range(n), precision, title = 'Precision', xlabel = 'Run', ylabel = '', ax=axs[1], ylim = ylim, linewidth=3)
    bs.nice_plot(range(n), recall, title = 'Recall', xlabel = 'Run', ylabel = '', ax=axs[2], ylim = ylim, linewidth=3)
    bs.nice_plot(range(n), f1_score, title = 'F1 score', xlabel = 'Run', ylabel = '', ax=axs[3], ylim = ylim, linewidth=3)
    bs.nice_plot(range(n), prompt_length, title = 'Prompt Length', ylabel = 'Number of characters', ax=axs[4], xlabel = 'Run', linewidth=3)
    if ndisagree_store:
        bs.nice_plot(range(n-1), ndisagree_store, title = 'Training Errors', ylabel = '# of Mistakes', ax=axs[5], xlabel = 'Run', linewidth=3)

    plt.tight_layout()
    
    if savename:
        savedir = os.path.dirname(saveloc)
        plt.savefig(savedir + '/' + savename + '.png')
        
    plt.show()

def df2topn(df, n, label_col = 'label'):
    """
    Return the first n occurrences of each label in the given column.

    Parameters:
    df (pd.DataFrame): The input DataFrame.
    label_col (str): The column name for which to find first n occurrences.
    n (int): Number of occurrences to retrieve.

    Returns:
    pd.DataFrame: A DataFrame containing the first n occurrences of each label, 
                  preserving original order.
    """
    df_with_rank = df.copy()
    df_with_rank['occurrence_rank'] = df_with_rank.groupby(label_col).cumcount()
    df_n = df_with_rank[df_with_rank['occurrence_rank'] < n].drop(columns='occurrence_rank')
    df_n = df_n.sort_index().reset_index(drop=True)
    return df_n

def nshot_prompt_dep(data_df, n_vec, covariate, label = 'label'):
    set_of_prompts = []
    for i in range(len(n_vec)):
        prompt = ''
        n = n_vec[i]
        data_df_topn = df2topn(data_df, n, label)
        for j in range(len(data_df_topn)):
            prompt += 'Text: ' + data_df_topn[covariate][j] + '\n'
            prompt += 'Label: ' + data_df_topn[label][j] + '\n\n'
            
        set_of_prompts.append(prompt)
        
    return set_of_prompts

def nshot_prompt(data_df, n_vec, covariate, label = 'label'):
    set_of_prompts = []
    for i in range(len(n_vec)):
        prompt = ''
        n = n_vec[i]
        for j in range(n):
            prompt += 'Text: ' + data_df[covariate][j] + '\n'
            prompt += 'Label: ' + data_df[label][j] + '\n\n'
            
        set_of_prompts.append(prompt)
        
    return set_of_prompts

    