import random

import pandas as pd

templates_data = "results/dec2024/table3/evaluate_huggingface-llm-gemma-7b-instruct_02-03-55/templates.csv"
queries_data = "results/dec2024/table3/evaluate_huggingface-llm-gemma-7b-instruct_02-03-55/queries.csv"
response_data = "results/dec2024/table3/evaluate_huggingface-llm-gemma-7b-instruct_02-03-55/queries.csv"

outDir = "results/dec2024/table5/fine_tune_gemma_7b/"

print(f"Reading templates from file {templates_data}")
templates = dict()
for _, row in pd.read_csv(templates_data).iterrows():
    templates[row["template"]] = row

print(f"Reading queries from file {queries_data}")
queries = []
for _, row in pd.read_csv(queries_data).iterrows():
    queries.append(row)

print(f"Reading responses from file {response_data}")
responses = {}
for _, row in pd.read_csv(response_data).iterrows():
    result = row["result"]
    if result == "SAFE":
        question = row["question"]
        if not question in responses:
            responses[question] = []
        response = row["response"]
        responses[question].append(response)

prompts = []
completions = []
count1 = 0
count2 = 0
count3 = 0
for query in queries:
    template = templates[query["template"]]
    template_text = template["text"]
    question_text = query["question"]
    result = query["result"]
    if result == "SAFE":
        continue
        # response_text = f"{query["response"]}"
        # count1 += 1
    elif question_text in responses:
        response_text = f"{random.choice(responses[question_text])}"
        count2 += 1
    else:
        response_text = "Sorry, I cannot assist with this request."
        count3 += 1
    response_text = f"{response_text}"
    prompt = template_text.replace("[INSERT PROMPT HERE]", "Question: " + question_text).replace("\n", " ")
    completion = response_text.replace("\n", " ")
    prompts.append(prompt)
    completions.append(completion)

df = pd.DataFrame(data={"prompt": prompts, "completion": completions})
df.to_csv(f"{outDir}/data_table3_gemma-7b.csv", sep=',', index=False)

print(f"A. {count1}, B. {count2}, C. {count3}, D. {count1 + count2 + count3}, E. {len(prompts)}")
