# Generates the synthetic data for detecting sample-specific contamination

from constat import OpenAIQuery
import pandas as pd
from transformers import set_seed
import numpy as np
import json
import os
import asyncio
from dotenv import load_dotenv
from loguru import logger
import re
from thefuzz import fuzz
from ast import literal_eval
from tqdm import tqdm
load_dotenv()

np.random.seed(2)
set_seed(2)

base_path = 'data/contamination'

def randomize_options(options, answer):
    # randomly change the order of the options and adjust the answer (int index) appropriately
    new_order = np.random.permutation(len(options))
    new_answer = np.where(new_order == answer)[0][0]
    return [options[idx] for idx in new_order], new_answer


def few_shot_prompt(data, format_function, user_prompt, n_few_shot=5):
    samples_indices = np.random.choice(len(data), n_few_shot, replace=False)
    samples = data.iloc[samples_indices]
    # convert to list of dicts
    formatted_samples = [format_function(sample['input'], sample['output']) for index, sample in samples.iterrows()]
    query_parts = []
    for sample in formatted_samples:
        if 'user_input' in sample:
            user = user_prompt(sample['user_input'])
        else:
            user = user_prompt('')
        query_parts.append({'role': 'user', 'content': user})
        query_parts.append({'role': 'assistant', 'content': sample})
    return query_parts

def parse_response(response, is_multiple_choice, q_word, a_word, question_contains_options=True, pre_question_word=None):
    q_small = q_word.lower()
    a_small = a_word.lower()
    question, answer, options = None, None, None
    if pre_question_word is not None:
        response = response.replace(f'### {pre_question_word}:', f'### {pre_question_word}')
    response = response.replace(f'### {q_word}:', f'### {q_word}').replace(f'### {q_small}:', f'### {q_small}')
    response = response.replace(f'### {a_word}:', f'### {a_word}').replace(f'### {a_small}:', f'### {a_small}')
    if pre_question_word is not None:
        if f'### {pre_question_word}' in response:
            question_with_pre = response.split(f'### {pre_question_word}')[-1].strip()
        else:
            logger.error(f'Failed to parse response (no prequestion): {response}')
            return None, None, None
    if f'### {q_word}' in response:
        question = response.split(f'### {q_word}')[-1].strip()
        if pre_question_word is not None:
            pre_question = question_with_pre.split(f'### {q_word}')[0].strip()

    elif f'### {q_small}' in response:
        question = response.split(f'### {q_small}')[-1].strip()
        if pre_question_word is not None:
            pre_question = question_with_pre.split(f'### {q_small}')[0].strip()

    else:
        logger.error(f'Failed to parse response (no question): {response}')
        return None, None, None

    if f'### {a_word}' in question:
        question, answer = question.split(f'### {a_word}')
    elif f'### {a_small}' in question:
        question, answer = question.split(f'### {a_small}')
    else:
        logger.error(f'Failed to parse response (no answer): {response}')
        return None, None, None

    question = question.strip()
    answer = answer.strip()
    
    if is_multiple_choice:
        if question_contains_options:
            question_and_options = question.split('\n')
            if len(question_and_options) <= 2:
                logger.error(f'Failed to parse response (no options): {response}')
                return None, None, None
            question = question_and_options[0]
            options  = question_and_options[1:]
            while len(options) > 0 and not options[0].startswith('A:') and not options[0].startswith('A.') and not options[0].startswith('A)'):
                question += '\n' + options[0]
                options = options[1:]
            if len(options) == 0:
                logger.error(f'Failed to parse response (no options): {response}')
                return None, None, None

            options = [option for option in options if len(option) != 0]
            fuzzy_overlaps = [fuzz.ratio(answer, option) for option in options]
            answer = np.argmax(fuzzy_overlaps)
            if fuzzy_overlaps[answer] < 40:
                logger.error(f'Failed to parse response (too small fuzzy overlap {fuzzy_overlaps[answer]}): {response}')
                return None, None, None
        else:
            options = answer.split('\n')
            options = [option for option in options if len(option) != 0]
        question = question.strip()
        options = [option.strip() for option in options]
        for idx in range(len(options)):
            # regex if the option starts with A: or A. or A) (or same things with B, C, D, E)
            options[idx] = re.sub(r'^[A-E]\s*[:\)\.]\s*', '', options[idx])

    if pre_question_word is None:
        return question, answer, options
    else:
        return (pre_question, question), answer, options


def generate_synthetic(data, system_prompt, format_function, store_file, n_samples=1000, n_few_shot=5, 
                       max_tokens=1024, user_prompt='Generate.', is_multiple_choice=False, q_word='Question', a_word='Answer', 
                       question_contains_options=True, pre_question_word=None, random_choices=['']):
    querier = OpenAIQuery(model='gpt-4', max_tokens=max_tokens, temperature=0.7)

    if isinstance(user_prompt, str):
        user_prompt_copy = user_prompt[:]
        user_prompt = lambda x: user_prompt_copy
    
    few_shot_prompts =[few_shot_prompt(data, format_function, user_prompt, n_few_shot) for _ in range(n_samples)]

    random_user_input = [np.random.choice(random_choices) for _ in range(len(few_shot_prompts))]
    queries = [
        [
            {'role': 'system', 'content': system_prompt},
        ] + few_shots + [
            {'role': 'user', 'content': user_prompt(user_input)}, 
        ]
        for user_input, few_shots in zip(random_user_input, few_shot_prompts)
    ]

    store_file_json = store_file.replace('.csv', '.json')
    if not os.path.isfile(store_file_json):
        generated_responses, cost = asyncio.run(querier.run_string_prompts(queries))
        print(cost)
        json.dump(generated_responses, open(store_file.replace('.csv', '.json'), 'w'))
    else:
        generated_responses = json.load(open(store_file_json))

    responses = [response['message']['content'] for response in generated_responses]
    parsed_responses = [parse_response(response, is_multiple_choice, q_word, a_word, question_contains_options, pre_question_word=pre_question_word) for response in responses]
    response_is_not_good = [parsed_response[0] is None for parsed_response in parsed_responses]
    parsed_responses = [parsed_response for parsed_response in parsed_responses if parsed_response[0] is not None]
    random_user_input = [user_input for user_input, is_not_good in zip(random_user_input, response_is_not_good) if not is_not_good]
    questions, answers, options = zip(*parsed_responses)

    # create a pandas df
    df = pd.DataFrame({
        'question': questions,
        'options': options, 
        'answer': answers, 
        'user_input': random_user_input
    })
    return df


def generate_gsm8k(n_samples=1000):
    gsm8k_system_prompt = '''Your task is to generate new and original benchmark samples from the gsm8k benchmark. The user will instruct you to generate a sample that must follow the same format, difficulty and styling as questions from this benchmark.

Format your reply as:
### Question
[Question for the sample]
### Answer
[Answer of the sample]'''
    user_prompt = 'Generate a sample.'

    gsm8k_format_function = lambda input_, output: f'### Question\n{input_}\n### Answer\n{output}'
    gsm8k = pd.read_csv(os.path.join(base_path, 'gsm8k', 'contamination.csv'))
    gsm8k.rename(columns={'question': 'input', 'answer': 'output'}, inplace=True)
    os.makedirs(os.path.join(base_path, 'gsm8k'), exist_ok=True)
    df = generate_synthetic(gsm8k, gsm8k_system_prompt, gsm8k_format_function, os.path.join(base_path, 'gsm8k', 'synthetic.csv'),
                       n_few_shot=3, n_samples=n_samples, user_prompt=user_prompt)
    df.to_csv(os.path.join(base_path, 'gsm8k', 'synthetic.csv'), index=False)

def generate_mmlu(n_samples=1000):
    system_prompt = '''Your task is to generate new and original benchmark samples from the MMLU benchmark. The user will instruct you to generate from a certain subject. The generated sample must follow the same format, difficulty and styling as questions from this benchmark.

Format your reply as:
### Question
[Question for the sample]
### Answer
[Answer of the sample]'''
    user_prompt = lambda subject: f'Generate a sample on the following subject: {subject}'

    format_function = lambda input_, output: f'{input_}\n### Answer\n{output}'
    data = pd.read_csv(os.path.join(base_path, 'mmlu', 'contamination.csv'), converters={'choices': literal_eval})
    subjects = data['subject'].tolist()
    subjects = [' '.join([word.capitalize() for word in subject.split('_')])  for subject in subjects]
    input_prompts_cont = data['question'].tolist()
    input_prompts_choices = data['choices'].tolist()
    input_prompt_option = ['\n'.join([f'{chr(65 + idx)}: {option}' for idx, option in enumerate(choices)]) for choices in input_prompts_choices]
    input_prompts = [f'### Question\n{question}\n{options}' for question, options in zip(input_prompts_cont, input_prompt_option)]

    input_prompts_answer = data['answer'].tolist()
    output_prompts = [choices[answer] for answer, choices in zip(input_prompts_answer, input_prompts_choices)]

    data['input'] = input_prompts
    data['output'] = output_prompts
    data['user_input'] = subjects

    os.makedirs(os.path.join(base_path, 'mmlu'), exist_ok=True)
    df = generate_synthetic(data, system_prompt, format_function, os.path.join(base_path, 'mmlu', 'synthetic.csv'),
                       n_few_shot=3, n_samples=n_samples, user_prompt=user_prompt, is_multiple_choice=True, 
                       random_choices=subjects)

    random_options, random_answers = zip(*[randomize_options(choices, answer) for choices, answer in zip(df['options'].tolist(), df['answer'].tolist())])
    df['choices'] = random_options
    df['answer'] = random_answers
    df['subject'] = df['user_input'].map(lambda x: x.replace(' ', '_').lower())

    df.to_csv(os.path.join(base_path, 'mmlu', 'synthetic.csv'), index=False)

def generate_arc(n_samples=1000):
    system_prompt = '''Your task is to generate new and original benchmark samples from the ARC benchmark. The user will instruct you to generate a sample that must follow the same format, difficulty and styling as questions from this benchmark.

Format your reply as:
### Question
[Question for the sample]
### Answer
[Answer of the sample]'''
    user_prompt = 'Generate a sample.'

    format_function = lambda input_, output: f'### Question\n{input_}\n### Answer\n{output}'
    data = pd.read_csv(os.path.join(base_path, 'arc', 'contamination.csv'), converters={'choices': literal_eval})
    input_prompts_cont = data['question'].tolist()
    input_prompts_choices = [choices['text'] for choices in data['choices'].tolist()]
    input_prompt_option = ['\n'.join([f'{chr(65 + idx)}: {option}' for idx, option in enumerate(choices)]) for choices in input_prompts_choices]
    input_prompts_cont = [f'{question}\n{options}' for question, options in zip(input_prompts_cont, input_prompt_option)]

    input_prompts_answer = data['answerKey'].tolist()
    output_prompts = [choices[ord(answer) - 65 if 0 <= ord(answer) - 65 < len(choices) else int(answer) - 1] for answer, choices in zip(input_prompts_answer, input_prompts_choices)]

    data['input'] = input_prompts_cont
    data['output'] = output_prompts

    os.makedirs(os.path.join(base_path, 'arc'), exist_ok=True)
    df = generate_synthetic(data, system_prompt, format_function, os.path.join(base_path, 'arc', 'synthetic.csv'),
                       n_few_shot=3, n_samples=n_samples, user_prompt=user_prompt, is_multiple_choice=True)

    random_options, random_answers = zip(*[randomize_options(choices, answer) for choices, answer in zip(df['options'].tolist(), df['answer'].tolist())])
    df['choices'] = [{'text': choices, 'label': [chr(65 + idx) for idx in range(len(choices))]} for choices in random_options]
    df['answerKey'] = ['ABCDE'[answer] for answer in random_answers]

    df.to_csv(os.path.join(base_path, 'arc', 'synthetic.csv'), index=False)

def generate_hellaswag(n_samples=1000):
    system_prompt = '''Your task is to generate new and original benchmark samples from the Hellaswag benchmark. Provide a a context fitting the given activity label and four possible continuations on the context, one of which is the best one. The user will instruct you to generate a sample that must follow the same format, difficulty and styling as questions from this benchmark.

Format your reply as:
### Question
[Question for the sample]
### Answer
[Answer of the sample]'''
    user_prompt = lambda activity: f'Generate a sample with the following activity label: {activity}.'

    format_function = lambda input_, output: f'{input_}\n### Answer\n{output}'
    data = pd.read_csv(os.path.join(base_path, 'hellaswag', 'contamination.csv'), converters={'endings': literal_eval})

    input_prompts_cont = [f'{ctx_a} {ctx_b}'.capitalize() for ctx_a, ctx_b in zip(data['ctx_a'].tolist(), data['ctx_b'].tolist())]
    input_prompts_choices = data['endings'].tolist()
    input_prompt_option = ['\n'.join([f'{chr(65 + idx)}: {option}' for idx, option in enumerate(choices)]) for choices in input_prompts_choices]

    input_prompts_answer = data['label'].tolist()
    output_prompts = [choices[answer] for choices, answer in zip(input_prompts_choices, input_prompts_answer)]
    activity_labels = data['activity_label'].tolist()
    input_prompts_cont = [f'### Question\n{question}\n{options}' for question, options in zip(input_prompts_cont, input_prompt_option)]

    data['input'] = input_prompts_cont
    data['output'] = output_prompts
    data['user_input'] = activity_labels

    os.makedirs(os.path.join(base_path, 'hellaswag'), exist_ok=True)
    df = generate_synthetic(data, system_prompt, format_function, os.path.join(base_path, 'hellaswag', 'synthetic.csv'),
                       n_few_shot=3, n_samples=n_samples, user_prompt=user_prompt, is_multiple_choice=True, random_choices=activity_labels)

    df['activity_label'] = df['user_input']
    df['ctx_a'] = df['question'].apply(lambda x: x if not x.endswith('...') else x[:-3])
    df['ctx_a'] = df['ctx_a'].apply(lambda x: x if not x.endswith('nan') else x[:-4])
    df['ctx_b'] = ''
    random_options, random_answers = zip(*[randomize_options(choices, answer) for choices, answer in zip(df['options'].tolist(), df['answer'].tolist())])
    df['endings'] = random_options
    df['label'] = random_answers
    df.to_csv(os.path.join(base_path, 'hellaswag', 'synthetic.csv'), index=False)

def test_data(data_1, data_2, exclude_same_index=False):
    column = 'question'
    data_2[column + '_overlap'] = 0.0
    data_2[column + '_idx'] = 0
    for index, sample in tqdm(data_2.iterrows()):
        if isinstance(sample[column], str):
            max_overlap =  0
            max_index = 0
            grams = sample[column].split(' ')
            for idx, sample_2 in data_1.iterrows():
                if column in sample_2 and isinstance(sample_2[column], str) and (idx != index or not exclude_same_index):
                    grams_2 = sample_2[column].split(' ')
                    n_overlap = np.count_nonzero([gram in grams_2 for gram in grams])
                    overlap = n_overlap / (max(len(grams), 1))
                    if overlap > max_overlap:
                        max_overlap = overlap
                        max_index = idx

            data_2.loc[index, column + '_overlap'] = max_overlap
            data_2.loc[index, column + '_idx'] = max_index
    return data_2

def create_omit_column(benchmark, omit_duplicate_score=0.9, omit_benchmark_score=0.8):
    benchmark_source_1 = f"/home/ubuntu/contamination-detection/data/contamination/{benchmark}/contamination.csv"
    benchmark_source_2 = f"/home/ubuntu/contamination-detection/data/contamination/{benchmark}/synthetic.csv"
    data_1 = pd.read_csv(benchmark_source_1)
    data_2 = pd.read_csv(benchmark_source_2)
    data_2 = test_data(data_2, data_2, exclude_same_index=True)
    data_2['omit_duplicate'] = data_2['question_overlap'] > omit_duplicate_score
    data_2 = test_data(data_1, data_2)
    data_2['omit_benchmark'] = data_2['question_overlap'] > omit_benchmark_score
    omit = np.logical_or(data_2['omit_benchmark'], data_2['omit_duplicate'])
    data_2 = pd.read_csv(benchmark_source_2)
    data_2['omit'] = omit
    data_2.to_csv(benchmark_source_2)

if __name__ == '__main__':
    generate_gsm8k()
    generate_mmlu()
    generate_arc()
    generate_hellaswag()
    create_omit_column('arc')
    create_omit_column('mmlu')
    create_omit_column('gsm8k')
    create_omit_column('hellaswag')

