# Generates the rephrased data for detecting syntax-specific contamination

from constat import OpenAIQuery
from datasets import load_dataset
from ast import literal_eval
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
from thefuzz import fuzz
import numpy as np
import re
load_dotenv()

np.random.seed(0)
set_seed(0)

base_path = 'data/contamination'


def parse_response(response, is_multiple_choice, contains_answer, q_word, a_word, question_contains_options=True):
    q_small = q_word.lower()
    a_small = a_word.lower()
    question, answer, options = None, None, None
    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 f'### {q_word}' in response:
        question = response.split(f'### {q_word}')[-1].strip()
    elif f'### {q_small}' in response:
        question = response.split(f'### {q_small}')[-1].strip()
    else:
        logger.warning(f'Failed to parse response (no question): {response}')
        return None, None, None

    if contains_answer:
        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.warning(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.warning(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.warning(f'Failed to parse response (no options): {response}')
                return None, None, None
        else:
            options = answer.split('\n')
        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])
    return question, answer, options


def generate(input_prompts, system_prompt, store_file, is_multiple_choice, contains_answer=False, 
             temperature=0.7, max_tokens=1024, question_word='Question', answer_word='Answer', question_contains_options=True):
    querier = OpenAIQuery(model='gpt-4-turbo', max_tokens=max_tokens, temperature=temperature, error_stop=100, timeout=120, read_cost=0.01, write_cost=0.03)

    queries = [
        [
            {'role': 'system', 'content': system_prompt},
            {'role': 'user', 'content': user_prompt},
        ] for user_prompt in input_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, contains_answer, question_word, answer_word, question_contains_options) for response in responses]
    parsed_responses = [parsed_response for parsed_response in parsed_responses if parsed_response[0] is not None]

    questions, answers, options = zip(*parsed_responses)

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


def generate_gsm8k(contains_answer=False, contamination=True, temperature=0.7, max_tokens=1024):
    if contamination:
        data = pd.read_csv(os.path.join(base_path, 'gsm8k', 'contamination.csv'))
    else:
        data = pd.read_csv(os.path.join(base_path, 'gsm8k', 'no_contamination.csv'))
    system_prompt = '''Significantly rephrase the given question, but make sure the answer is still the same. Do not include the answer in your response.

Format your reply as:
### Question
[New rephrased question]'''
    if contains_answer:
        system_prompt = '''You are a helpful assistant. The user will give you a question and answer from the gsm8k dataset. Rewrite the question and answer. Make significant changes to the used vocabulary, length and structure. Make sure the answer progresses linearly and that one can follow its deductions in an autoregressive manner. Still employ the '####' sign to indicate the output, and the '<< >>' signs to indicate computations. Ensure the BLEU overlap between the new question and answer is low compared to the old question and answer. 

Format your reply as:
### Question
[New Rephrased Question]
### Answer
[New Answer]'''

    file_name = 'synthetic'
    if contains_answer:
        file_name += '_answer'

    if not contamination:
        file_name += '_no'

    input_prompts_cont = data['question'].tolist()
    if contains_answer:
        input_prompts_answer = data['answer'].tolist()
        input_prompts_cont = [f'### Question\n{question}\n### Answer\n{answer}' for question, answer in zip(input_prompts_cont, input_prompts_answer)]
    else:
        input_prompts_cont = [f'### Question\n{question}' for question in input_prompts_cont]
    
    df = generate(input_prompts_cont, system_prompt, os.path.join(base_path, 'gsm8k', f'{file_name}_cont.csv'), is_multiple_choice=False, contains_answer=contains_answer, temperature=temperature, max_tokens=max_tokens)

    if not contains_answer:
        df['answer'] = data['answer'].tolist()

    df.to_csv(os.path.join(base_path, 'gsm8k', f'{file_name}_cont.csv'), index=False)


def generate_mmlu(contains_answer=False, contamination=False, temperature=0.7, max_tokens=1024):
    if contamination:
        data = pd.read_csv(os.path.join(base_path, 'mmlu', 'contamination.csv'), converters={'choices': literal_eval})
    else:
        data = pd.read_csv(os.path.join(base_path, 'mmlu', 'no_contamination.csv'), converters={'choices': literal_eval})
    system_prompt = '''Significantly rephrase the given question and options, but make sure that all possible options still have the same label. Label the multiple choice answers with A:, B:, C:, D:, E:. Do not include the answer in your response.

Format your reply as:
### Question
[New rephrased question]'''
    if contains_answer:
        system_prompt = '''Significantly rephrase the given question and options, but make sure that the correct option still has the same label. Label the multiple choice answers with A:, B:, C:, D:, E:. Do not include the answer in your response. Wrong options can be completely changed (and thus do not need to be simple rephrases).

Format your reply as:
### Reasoning
[Reasoning about how to rephrase question and options]
### Question
[New rephrased question]'''

    file_name = 'synthetic'
    if contains_answer:
        file_name += '_answer'
    
    if not contamination:
        file_name += '_no'

    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_cont = [f'{question}\n{options}' for question, options in zip(input_prompts_cont, input_prompt_option)]
    input_prompts_cont = [f'### Question\n{question}' for question in input_prompts_cont]

    if contains_answer:
        input_prompts_answer = data['answer'].tolist()
        input_prompts_cont = [f'{question}\n### Answer\n{choices[answer]}' for question, answer, choices in zip(input_prompts_cont, input_prompts_answer, input_prompts_choices)]

    df = generate(input_prompts_cont, system_prompt, os.path.join(base_path, 'mmlu', f'{file_name}_cont.csv'), 
                  is_multiple_choice=True, contains_answer=False, temperature=temperature, max_tokens=max_tokens)

    df['choices'] = df['options']
    df['answer'] = data['answer'].tolist()
    df['subject'] = data['subject']

    df.to_csv(os.path.join(base_path, 'mmlu', f'{file_name}_cont.csv'), index=False)


def generate_arc(contains_answer=False, contamination=False, temperature=0.7, max_tokens=1024):
    if contamination:
        data = pd.read_csv(os.path.join(base_path, 'arc', 'contamination.csv'), converters={'choices': literal_eval})
    else:
        data = pd.read_csv(os.path.join(base_path, 'arc', 'no_contamination.csv'), converters={'choices': literal_eval})
    system_prompt = '''Significantly rephrase the given question and options, but make sure that all possible options still have the same label. Label the multiple choice answers with A:, B:, C:, D:, E:. Do not include the answer in your response.

Format your reply as:
### Question
[New rephrased question]'''
    if contains_answer:
        system_prompt = '''Significantly rephrase the given question and options, but make sure that the correct option still has the same label. Label the multiple choice answers with A:, B:, C:, D:, E:. Do not include the answer in your response. Wrong options can be completely changed (and thus do not need to be simple rephrases).

Format your reply as:
### Reasoning
[Reasoning about how to rephrase question and options]
### Question
[New rephrased question]'''

    file_name = 'synthetic'
    if contains_answer:
        file_name += '_answer'

    if not contamination:
        file_name += '_no'

    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_cont = [f'### Question\n{question}' for question in input_prompts_cont]

    if contains_answer:
        input_prompts_answer = data['answerKey'].tolist()
        input_prompts_cont = [f'{question}\n### Answer\n{choices[ord(answer) - 65 if 0 <= ord(answer) - 65 < len(choices) else int(answer) - 1]}' for question, answer, choices in zip(input_prompts_cont, input_prompts_answer, input_prompts_choices)]

    df = generate(input_prompts_cont, system_prompt, os.path.join(base_path, 'arc', f'{file_name}_cont.csv'), 
                  is_multiple_choice=True, contains_answer=False, temperature=temperature, max_tokens=max_tokens)

    df['choices'] = [{'text': choices, 'label': [chr(65 + idx) for idx in range(len(choices))]} for choices in df['options'].tolist()]
    df['answerKey'] = data['answerKey'].tolist()
    df['answerKey'] = df['answerKey'].apply(lambda x: chr(65 + int(x) - 1) if x.isdigit() else x)

    df.to_csv(os.path.join(base_path, 'arc', f'{file_name}_cont.csv'), index=False)


def generate_hellaswag(contains_answer=False, contamination=False, temperature=0.7, max_tokens=1024):
    if contamination:
        data = pd.read_csv(os.path.join(base_path, 'hellaswag', 'contamination.csv'), converters={'endings': literal_eval})
    else:
        data = pd.read_csv(os.path.join(base_path, 'hellaswag', 'no_contamination.csv'), converters={'endings': literal_eval})
    system_prompt = '''Significantly rephrase the given context and several possible continuations and make sure that all possible continuations still have the same label. Label the continuations with A:, B:, C:, D:, E:. Keep the ':' separating the activity label from the context in the context. Make sure that each continuation can still follow the rephrased context.

Format your reply as:
### Context
[New rephrased question]
### Continuation
[New rephrased continuation]'''
    if contains_answer:
        system_prompt = '''Significantly rephrase the given context and continuation. Keep the ':' separating the activity label from the context in the context. Make sure the continuation can still follow the rephrased context.

Format your reply as:
### Reasoning
[Reasoning about how to rephrase context and continuation]
### Context
[New rephrased question]
### Continuation
[New rephrased continuation]'''

    file_name = 'synthetic'
    if contains_answer:
        file_name += '_answer'

    if not contamination:
        file_name += '_no'

    input_prompts_cont = [f'{activity_label}: {ctx_a} {ctx_b}'.capitalize() for activity_label, ctx_a, ctx_b in zip(data['activity_label'].tolist(), 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]

    if contains_answer:
        input_prompts_answer = data['label'].tolist()
        input_prompts_cont = [f'### Context\n{question}\n### Continuation\n{choices[answer]}' for question, answer, choices in zip(input_prompts_cont, input_prompts_answer, input_prompts_choices)]
    else:
        input_prompts_cont = [f'### Context\n{question}\n### Continuation\n{options}' for question, options in zip(input_prompts_cont, input_prompt_option)]

    df = generate(input_prompts_cont, system_prompt, os.path.join(base_path, 'hellaswag', f'{file_name}_cont.csv'), 
                  is_multiple_choice=not contains_answer, contains_answer=True, temperature=temperature, 
                  max_tokens=max_tokens, question_word='Context', answer_word='Continuation', 
                  question_contains_options=False)

    df['activity_label'] = df['question'].apply(lambda x: x.split(':')[0] if ':' in x else '')
    df['ctx_a'] = df['question'].apply(lambda x: (':').join(x.split(':')[1:]) if ':' in x else x)
    df['ctx_b'] = ''
    df['endings'] = df['options']
    df['label'] = data['label'].tolist()

    df.to_csv(os.path.join(base_path, 'hellaswag', f'{file_name}_cont.csv'), index=False)

if __name__ == '__main__':
    generate_gsm8k(False, True)
    generate_gsm8k(True, True)
    generate_gsm8k(False, False)
    generate_gsm8k(True, False)
    generate_mmlu(False, True)
    generate_mmlu(False, False)
    generate_mmlu(True, True)
    generate_mmlu(True, False)
    generate_arc(False, True)
    generate_arc(False, False)
    generate_arc(True, True)
    generate_arc(True, False)
    generate_hellaswag(False, False)
    generate_hellaswag(True, False)
    generate_hellaswag(False, True)
    generate_hellaswag(True, True)
